kmp_observability/
embedded_telemetry_guard.rs

1use std::io;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::mpsc::{Receiver, RecvTimeoutError};
5use std::thread::JoinHandle;
6use std::time::Duration;
7
8use crate::QualityTelemetryObservation;
9
10/// Owns a runtime-neutral telemetry worker and closes it without requiring
11/// senders to disconnect first.
12pub struct EmbeddedTelemetryGuard {
13    shutdown: Arc<AtomicBool>,
14    handle: Option<JoinHandle<()>>,
15}
16
17impl EmbeddedTelemetryGuard {
18    pub fn try_spawn<W, F>(
19        receiver: Receiver<QualityTelemetryObservation>,
20        batch_size: usize,
21        flush_interval: Duration,
22        mut write_batch: W,
23        finalize: F,
24    ) -> Result<Self, io::Error>
25    where
26        W: FnMut(Vec<QualityTelemetryObservation>) + Send + 'static,
27        F: FnOnce() + Send + 'static,
28    {
29        let shutdown = Arc::new(AtomicBool::new(false));
30        let worker_shutdown = Arc::clone(&shutdown);
31        let batch_size = batch_size.max(1);
32        let flush_interval = if flush_interval.is_zero() {
33            Duration::from_millis(1)
34        } else {
35            flush_interval
36        };
37        let handle = std::thread::Builder::new()
38            .name("quality-telemetry".to_string())
39            .spawn(move || {
40                let mut pending = Vec::with_capacity(batch_size);
41                loop {
42                    if worker_shutdown.load(Ordering::Acquire) {
43                        pending.extend(receiver.try_iter());
44                        flush_pending(&mut pending, batch_size, &mut write_batch);
45                        break;
46                    }
47                    match receiver.recv_timeout(flush_interval) {
48                        Ok(observation) => {
49                            pending.push(observation);
50                            if pending.len() >= batch_size {
51                                flush_pending(&mut pending, batch_size, &mut write_batch);
52                            }
53                        }
54                        Err(RecvTimeoutError::Timeout) => {
55                            flush_pending(&mut pending, batch_size, &mut write_batch);
56                        }
57                        Err(RecvTimeoutError::Disconnected) => {
58                            flush_pending(&mut pending, batch_size, &mut write_batch);
59                            break;
60                        }
61                    }
62                }
63                finalize();
64            })?;
65        Ok(Self {
66            shutdown,
67            handle: Some(handle),
68        })
69    }
70
71    pub fn close(mut self) {
72        self.stop();
73    }
74
75    fn stop(&mut self) {
76        self.shutdown.store(true, Ordering::Release);
77        if let Some(handle) = self.handle.take() {
78            let _ = handle.join();
79        }
80    }
81}
82
83impl Drop for EmbeddedTelemetryGuard {
84    fn drop(&mut self) {
85        self.stop();
86    }
87}
88
89fn flush_pending<W>(
90    pending: &mut Vec<QualityTelemetryObservation>,
91    batch_size: usize,
92    write_batch: &mut W,
93) where
94    W: FnMut(Vec<QualityTelemetryObservation>),
95{
96    if pending.is_empty() {
97        return;
98    }
99    let batch = std::mem::replace(pending, Vec::with_capacity(batch_size));
100    write_batch(batch);
101}
102
103#[cfg(test)]
104mod tests {
105    use std::sync::Arc;
106    use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
107    use std::time::Duration;
108
109    use kmp_domain::{BundleQualityMetrics, QualityMetricsObserver, QualityObservationContext};
110
111    use crate::BufferedQualityMetricsObserver;
112
113    use super::EmbeddedTelemetryGuard;
114
115    #[test]
116    fn drop_drains_and_joins_even_while_the_observer_is_alive() {
117        let (observer, receiver) = BufferedQualityMetricsObserver::with_capacity(4);
118        let written = Arc::new(AtomicU64::new(0));
119        let finalized = Arc::new(AtomicBool::new(false));
120        let worker_written = Arc::clone(&written);
121        let worker_finalized = Arc::clone(&finalized);
122        let guard = EmbeddedTelemetryGuard::try_spawn(
123            receiver,
124            4,
125            Duration::from_millis(5),
126            move |batch| {
127                worker_written.fetch_add(batch.len() as u64, Ordering::Relaxed);
128            },
129            move || worker_finalized.store(true, Ordering::Release),
130        )
131        .expect("worker starts");
132        let metrics = BundleQualityMetrics::new(1, 1.0, 0.0, 0.0, 0.0).expect("valid metrics");
133        observer.observe(
134            &metrics,
135            &QualityObservationContext {
136                rpc: "kernel_wake".to_string(),
137                root_node_id: "question:guard".to_string(),
138                role: "resumer".to_string(),
139            },
140        );
141
142        drop(guard);
143
144        assert_eq!(written.load(Ordering::Relaxed), 1);
145        assert!(finalized.load(Ordering::Acquire));
146        assert_eq!(observer.dropped_observations(), 0);
147    }
148}
kmp_observability/embedded_telemetry_guard.rs

kmp_observability/
embedded_telemetry_guard.rs
