metrale_telemetry/
sampler.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The device sampler thread, the I/O half of L0 and energy. It
4//! reads a [`DeviceSource`] at the configured period and hands each reading to
5//! [`Telemetry::absorb_device`]. It runs on its own thread, so a slow driver
6//! call delays the next sample, not a scheduler step.
7//!
8//! Owner: telemetry.
9//! Invariants: stopping or dropping a [`SamplerHandle`] stops its thread and
10//! joins it.
11
12use std::sync::Arc;
13use std::sync::atomic::{AtomicBool, Ordering};
14use std::thread::JoinHandle;
15use std::time::{Duration, Instant};
16
17use crate::device::{DeviceSource, NvmlSource};
18use crate::energy::CounterTracker;
19use crate::hub::Telemetry;
20use crate::level::{Level, TelemetryConfig};
21use metrale_gpu_sys::nvml::NvmlUnavailable;
22
23/// 2026-09-26: A running sampler. Dropping it stops the thread and joins it.
24pub struct SamplerHandle {
25    stop: Arc<AtomicBool>,
26    thread: Option<JoinHandle<()>>,
27}
28
29impl SamplerHandle {
30    /// 2026-09-26: Stop and join.
31    pub fn stop(mut self) {
32        self.shutdown();
33    }
34
35    fn shutdown(&mut self) {
36        self.stop.store(true, Ordering::Release);
37        if let Some(t) = self.thread.take() {
38            t.thread().unpark();
39            // 2026-09-26: A panicked sampler has already stopped; its join error
40            // is ignored.
41            let _ = t.join();
42        }
43    }
44}
45
46impl Drop for SamplerHandle {
47    fn drop(&mut self) {
48        self.shutdown();
49    }
50}
51
52/// 2026-09-26: Sample `source` into `t` every `period` until stopped.
53pub fn spawn(
54    t: &'static Telemetry,
55    mut source: Box<dyn DeviceSource>,
56    period: Duration,
57) -> std::io::Result<SamplerHandle> {
58    let stop = Arc::new(AtomicBool::new(false));
59    let flag = stop.clone();
60    let thread = std::thread::Builder::new()
61        .name("metrale-telemetry-device".into())
62        .spawn(move || {
63            let mut tracker = CounterTracker::default();
64            let mut next = Instant::now();
65            while !flag.load(Ordering::Acquire) {
66                let reading = source.read();
67                t.absorb_device(&mut tracker, &reading, source.read_errors());
68                next += period;
69                let now = Instant::now();
70                if next <= now {
71                    // 2026-09-26: Fell behind: the next sample is one period
72                    // from now.
73                    next = now + period;
74                }
75                std::thread::park_timeout(next - now);
76            }
77        })?;
78    Ok(SamplerHandle {
79        stop,
80        thread: Some(thread),
81    })
82}
83
84/// 2026-09-26: Configure `t` and, for a level that samples, start the NVML
85/// sampler.
86///
87/// `Ok(None)` when the level is `Off`: nothing is started or allocated.
88/// `Err(reason)` when the device cannot be read (NVML missing, no such
89/// device, no thread): `t` is marked unavailable with the reason, which is
90/// also returned.
91pub fn start(
92    t: &'static Telemetry,
93    cfg: &TelemetryConfig,
94) -> Result<Option<SamplerHandle>, String> {
95    start_with(t, cfg, |index| {
96        NvmlSource::open(index).map(|s| Box::new(s) as Box<dyn DeviceSource>)
97    })
98}
99
100/// 2026-09-26: [`start`] over an explicit device opener.
101pub fn start_with(
102    t: &'static Telemetry,
103    cfg: &TelemetryConfig,
104    open: impl FnOnce(u32) -> Result<Box<dyn DeviceSource>, NvmlUnavailable>,
105) -> Result<Option<SamplerHandle>, String> {
106    t.configure(cfg);
107    if cfg.level() == Level::Off {
108        return Ok(None);
109    }
110    let failed = |why: String| {
111        t.mark_device_unavailable(why.clone());
112        Err(why)
113    };
114    match open(cfg.device_index()) {
115        Ok(src) => match spawn(t, src, cfg.device_period()) {
116            Ok(h) => Ok(Some(h)),
117            Err(e) => failed(format!("device sampler thread did not start: {e}")),
118        },
119        Err(e) => failed(e.to_string()),
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::clock::MonotonicClock;
127    use crate::device::DeviceReading;
128    use crate::hub::DeviceState;
129
130    struct Counting {
131        counter_mj: u64,
132    }
133
134    impl DeviceSource for Counting {
135        fn read(&mut self) -> DeviceReading {
136            self.counter_mj += 1_000;
137            DeviceReading {
138                energy_counter_mj: Some(self.counter_mj),
139                ..Default::default()
140            }
141        }
142        fn read_errors(&self) -> u64 {
143            0
144        }
145    }
146
147    #[test]
148    fn the_thread_samples_until_stopped_then_stops() {
149        static CLOCK: MonotonicClock = MonotonicClock;
150        let t: &'static Telemetry = Box::leak(Box::new(Telemetry::new(&CLOCK)));
151        t.configure(&TelemetryConfig::new(Level::Basic, Duration::from_millis(10), 1, 0).unwrap());
152        let h = spawn(
153            t,
154            Box::new(Counting { counter_mj: 0 }),
155            Duration::from_millis(10),
156        )
157        .unwrap();
158        let deadline = Instant::now() + Duration::from_secs(10);
159        while t.device_reading().map_or(0, |(_, _, n)| n) < 3 {
160            assert!(
161                Instant::now() < deadline,
162                "the sampler never produced 3 samples"
163            );
164            std::thread::sleep(Duration::from_millis(5));
165        }
166        h.stop();
167        let after_stop = t.device_reading().unwrap().2;
168        std::thread::sleep(Duration::from_millis(50));
169        assert_eq!(
170            t.device_reading().unwrap().2,
171            after_stop,
172            "no sample after stop"
173        );
174        assert_eq!(t.device_state(), DeviceState::Available);
175        // 2026-09-26: Each sample after the first adds 1000 mJ.
176        assert_eq!(t.energy_mj.get(), (after_stop - 1) * 1_000);
177    }
178
179    #[test]
180    fn level_off_starts_nothing_and_touches_no_driver() {
181        static CLOCK: MonotonicClock = MonotonicClock;
182        let t: &'static Telemetry = Box::leak(Box::new(Telemetry::new(&CLOCK)));
183        let h = start(t, &TelemetryConfig::serving(Level::Off, 0)).unwrap();
184        assert!(h.is_none());
185        assert_eq!(t.device_state(), DeviceState::NotStarted);
186    }
187}
metrale_telemetry/sampler.rs

metrale_telemetry/
sampler.rs
