datum/
instrument.rs

1//! Opt-in stream instrumentation for operational tooling.
2//!
3//! Instrumentation is explicit: ordinary stream blueprints do not carry a registry,
4//! branch, timestamp call, or atomic counter. Calling [`Source::instrumented`] inserts
5//! a small counting boundary into that source only. Each materialization registers a
6//! fresh run in a [`StreamInstrumentationRegistry`], and future control-plane tools can
7//! read snapshots from the registry without putting actors on the element path.
8//!
9//! [`Source::instrumented`]: crate::Source::instrumented
10
11use crate::stream::{BoxStream, StreamError, StreamResult};
12use std::{
13    collections::BTreeMap,
14    sync::{
15        Arc, Mutex,
16        atomic::{AtomicU8, AtomicU64, Ordering},
17    },
18    time::{Duration, SystemTime, UNIX_EPOCH},
19};
20
21const STATE_RUNNING: u8 = 0;
22const STATE_DRAINING: u8 = 1;
23const STATE_COMPLETED: u8 = 2;
24const STATE_FAILED: u8 = 3;
25
26/// Stable identifier for one instrumented stream materialization.
27#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
28pub struct StreamInstrumentationId(u64);
29
30impl StreamInstrumentationId {
31    /// Return the numeric id assigned by the registry.
32    #[must_use]
33    pub const fn get(self) -> u64 {
34        self.0
35    }
36}
37
38/// The externally visible state of one instrumented stream materialization.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum StreamInstrumentationState {
41    /// The stream has been materialized and has not reached a terminal state.
42    Running,
43    /// The stream is draining or was dropped before the boundary observed normal completion.
44    Draining,
45    /// The boundary observed upstream completion.
46    Completed,
47    /// The boundary observed a non-cancellation stream failure.
48    Failed,
49}
50
51impl StreamInstrumentationState {
52    #[must_use]
53    const fn from_code(code: u8) -> Self {
54        match code {
55            STATE_DRAINING => Self::Draining,
56            STATE_COMPLETED => Self::Completed,
57            STATE_FAILED => Self::Failed,
58            _ => Self::Running,
59        }
60    }
61
62    #[must_use]
63    const fn code(self) -> u8 {
64        match self {
65            Self::Running => STATE_RUNNING,
66            Self::Draining => STATE_DRAINING,
67            Self::Completed => STATE_COMPLETED,
68            Self::Failed => STATE_FAILED,
69        }
70    }
71
72    #[must_use]
73    const fn is_terminal(self) -> bool {
74        matches!(self, Self::Completed | Self::Failed)
75    }
76}
77
78/// Point-in-time values for one instrumented stream materialization.
79#[derive(Clone, Debug, PartialEq, Eq)]
80pub struct StreamInstrumentationSnapshot {
81    /// Stable id assigned when the stream materialized.
82    pub id: StreamInstrumentationId,
83    /// User-provided instrumentation name.
84    pub name: String,
85    /// Count of successful elements observed by the instrumentation boundary.
86    pub elements_through: u64,
87    /// Restart count recorded by the owner of this materialization.
88    pub restarts: u64,
89    /// Current stream state.
90    pub state: StreamInstrumentationState,
91    /// Wall-clock time when this materialization registered.
92    pub started_at: SystemTime,
93    /// Wall-clock time of the most recent state transition.
94    pub state_changed_at: SystemTime,
95    /// Wall-clock terminal timestamp when the boundary observed completion or failure.
96    pub finished_at: Option<SystemTime>,
97    /// Elapsed time since `started_at`, or until `finished_at` for terminal runs.
98    pub uptime: Duration,
99}
100
101/// Cloneable registry handle for instrumented stream materializations.
102///
103/// The registry is not consulted by ordinary streams. It is only captured by streams that
104/// explicitly call [`Source::instrumented`], and registration happens once per materialization.
105///
106/// [`Source::instrumented`]: crate::Source::instrumented
107#[derive(Clone, Debug, Default)]
108pub struct StreamInstrumentationRegistry {
109    inner: Arc<RegistryInner>,
110}
111
112#[derive(Debug, Default)]
113struct RegistryInner {
114    next_id: AtomicU64,
115    runs: Mutex<BTreeMap<StreamInstrumentationId, Arc<StreamInstrumentationCounters>>>,
116}
117
118impl StreamInstrumentationRegistry {
119    /// Create an empty registry.
120    #[must_use]
121    pub fn new() -> Self {
122        Self::default()
123    }
124
125    /// Register one materialized stream run.
126    ///
127    /// This is public so a future control plane can create a run handle for work it
128    /// supervises directly. Most stream users should prefer [`Source::instrumented`].
129    ///
130    /// [`Source::instrumented`]: crate::Source::instrumented
131    #[must_use]
132    pub fn register(&self, name: impl Into<String>) -> StreamInstrumentationRun {
133        let id = StreamInstrumentationId(self.inner.next_id.fetch_add(1, Ordering::Relaxed) + 1);
134        let counters = Arc::new(StreamInstrumentationCounters::new(id, name.into()));
135        self.inner
136            .runs
137            .lock()
138            .expect("stream instrumentation registry poisoned")
139            .insert(id, Arc::clone(&counters));
140        StreamInstrumentationRun { counters }
141    }
142
143    /// Return a snapshot for one materialized run, if it is still retained.
144    #[must_use]
145    pub fn snapshot(&self, id: StreamInstrumentationId) -> Option<StreamInstrumentationSnapshot> {
146        self.inner
147            .runs
148            .lock()
149            .expect("stream instrumentation registry poisoned")
150            .get(&id)
151            .map(|counters| counters.snapshot())
152    }
153
154    /// Return snapshots for every retained materialized run, ordered by id.
155    #[must_use]
156    pub fn snapshots(&self) -> Vec<StreamInstrumentationSnapshot> {
157        self.inner
158            .runs
159            .lock()
160            .expect("stream instrumentation registry poisoned")
161            .values()
162            .map(|counters| counters.snapshot())
163            .collect()
164    }
165
166    /// Drop a retained run from the registry.
167    ///
168    /// Removing a run only affects future snapshots; existing [`StreamInstrumentationRun`]
169    /// handles continue to own their counters.
170    pub fn remove(&self, id: StreamInstrumentationId) -> bool {
171        self.inner
172            .runs
173            .lock()
174            .expect("stream instrumentation registry poisoned")
175            .remove(&id)
176            .is_some()
177    }
178}
179
180/// Handle for updating one instrumented materialization.
181#[derive(Clone, Debug)]
182pub struct StreamInstrumentationRun {
183    counters: Arc<StreamInstrumentationCounters>,
184}
185
186impl StreamInstrumentationRun {
187    /// The id assigned by the registry.
188    #[must_use]
189    pub fn id(&self) -> StreamInstrumentationId {
190        self.counters.id
191    }
192
193    /// Record one successfully observed element.
194    pub fn record_element(&self) {
195        self.record_elements(1);
196    }
197
198    /// Record multiple successfully observed elements.
199    pub fn record_elements(&self, elements: u64) {
200        self.counters
201            .elements_through
202            .fetch_add(elements, Ordering::Relaxed);
203    }
204
205    /// Record one restart for this materialization/job owner.
206    pub fn record_restart(&self) {
207        self.record_restarts(1);
208    }
209
210    /// Record multiple restarts for this materialization/job owner.
211    pub fn record_restarts(&self, restarts: u64) {
212        self.counters
213            .restarts
214            .fetch_add(restarts, Ordering::Relaxed);
215    }
216
217    /// Mark the run as running.
218    pub fn mark_running(&self) {
219        self.counters
220            .mark_state(StreamInstrumentationState::Running);
221    }
222
223    /// Mark the run as draining.
224    pub fn mark_draining(&self) {
225        self.counters
226            .mark_state(StreamInstrumentationState::Draining);
227    }
228
229    /// Mark the run as completed.
230    pub fn mark_completed(&self) {
231        self.counters
232            .mark_state(StreamInstrumentationState::Completed);
233    }
234
235    /// Mark the run as failed.
236    pub fn mark_failed(&self) {
237        self.counters.mark_state(StreamInstrumentationState::Failed);
238    }
239
240    /// Return a point-in-time snapshot of this run.
241    #[must_use]
242    pub fn snapshot(&self) -> StreamInstrumentationSnapshot {
243        self.counters.snapshot()
244    }
245}
246
247#[derive(Debug)]
248struct StreamInstrumentationCounters {
249    id: StreamInstrumentationId,
250    name: Arc<str>,
251    elements_through: AtomicU64,
252    restarts: AtomicU64,
253    state: AtomicU8,
254    started_at_millis: u64,
255    state_changed_at_millis: AtomicU64,
256    finished_at_millis: AtomicU64,
257}
258
259impl StreamInstrumentationCounters {
260    fn new(id: StreamInstrumentationId, name: String) -> Self {
261        let now = unix_time_millis(SystemTime::now());
262        Self {
263            id,
264            name: Arc::from(name),
265            elements_through: AtomicU64::new(0),
266            restarts: AtomicU64::new(0),
267            state: AtomicU8::new(STATE_RUNNING),
268            started_at_millis: now,
269            state_changed_at_millis: AtomicU64::new(now),
270            finished_at_millis: AtomicU64::new(0),
271        }
272    }
273
274    fn mark_state(&self, state: StreamInstrumentationState) {
275        let now = unix_time_millis(SystemTime::now());
276        self.state.store(state.code(), Ordering::Relaxed);
277        self.state_changed_at_millis.store(now, Ordering::Relaxed);
278        if state.is_terminal() {
279            self.finished_at_millis.store(now, Ordering::Relaxed);
280        }
281    }
282
283    fn snapshot(&self) -> StreamInstrumentationSnapshot {
284        let state = StreamInstrumentationState::from_code(self.state.load(Ordering::Relaxed));
285        let started_at = system_time_from_millis(self.started_at_millis);
286        let state_changed_at =
287            system_time_from_millis(self.state_changed_at_millis.load(Ordering::Relaxed));
288        let finished_at_millis = self.finished_at_millis.load(Ordering::Relaxed);
289        let finished_at =
290            (finished_at_millis != 0).then(|| system_time_from_millis(finished_at_millis));
291        let uptime_end = finished_at.unwrap_or_else(SystemTime::now);
292        let uptime = uptime_end
293            .duration_since(started_at)
294            .unwrap_or(Duration::ZERO);
295
296        StreamInstrumentationSnapshot {
297            id: self.id,
298            name: self.name.to_string(),
299            elements_through: self.elements_through.load(Ordering::Relaxed),
300            restarts: self.restarts.load(Ordering::Relaxed),
301            state,
302            started_at,
303            state_changed_at,
304            finished_at,
305            uptime,
306        }
307    }
308}
309
310pub(crate) struct InstrumentedStream<T> {
311    input: BoxStream<T>,
312    run: StreamInstrumentationRun,
313    terminal_observed: bool,
314}
315
316impl<T> InstrumentedStream<T> {
317    pub(crate) fn new(input: BoxStream<T>, run: StreamInstrumentationRun) -> Self {
318        Self {
319            input,
320            run,
321            terminal_observed: false,
322        }
323    }
324}
325
326impl<T> Iterator for InstrumentedStream<T> {
327    type Item = StreamResult<T>;
328
329    fn next(&mut self) -> Option<Self::Item> {
330        if self.terminal_observed {
331            return None;
332        }
333
334        match self.input.next() {
335            Some(Ok(item)) => {
336                self.run.record_element();
337                Some(Ok(item))
338            }
339            Some(Err(error)) => {
340                self.terminal_observed = true;
341                if matches!(error, StreamError::Cancelled) {
342                    self.run.mark_draining();
343                } else {
344                    self.run.mark_failed();
345                }
346                Some(Err(error))
347            }
348            None => {
349                self.terminal_observed = true;
350                self.run.mark_completed();
351                None
352            }
353        }
354    }
355}
356
357impl<T> Drop for InstrumentedStream<T> {
358    fn drop(&mut self) {
359        if !self.terminal_observed {
360            self.run.mark_draining();
361        }
362    }
363}
364
365fn unix_time_millis(time: SystemTime) -> u64 {
366    time.duration_since(UNIX_EPOCH)
367        .map(|duration| duration.as_millis().min(u128::from(u64::MAX)) as u64)
368        .unwrap_or(0)
369}
370
371fn system_time_from_millis(millis: u64) -> SystemTime {
372    UNIX_EPOCH + Duration::from_millis(millis)
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378    use crate::{Keep, NotUsed, Sink, Source};
379    use std::{
380        sync::{
381            Arc,
382            atomic::{AtomicBool, Ordering},
383        },
384        thread,
385    };
386
387    #[test]
388    fn enabled_counters_track_successful_stream() {
389        let registry = StreamInstrumentationRegistry::new();
390
391        let values = Source::from_iter(0_u64..4)
392            .instrumented("success", &registry)
393            .run_collect()
394            .expect("instrumented stream succeeds");
395
396        assert_eq!(values, vec![0, 1, 2, 3]);
397        let snapshots = registry.snapshots();
398        assert_eq!(snapshots.len(), 1);
399        let snapshot = &snapshots[0];
400        assert_eq!(snapshot.name, "success");
401        assert_eq!(snapshot.elements_through, 4);
402        assert_eq!(snapshot.restarts, 0);
403        assert_eq!(snapshot.state, StreamInstrumentationState::Completed);
404        assert!(snapshot.finished_at.is_some());
405        assert!(snapshot.uptime >= Duration::ZERO);
406    }
407
408    #[test]
409    fn enabled_counters_track_failure_after_successful_elements() {
410        let registry = StreamInstrumentationRegistry::new();
411        let error = StreamError::Failed("boom".into());
412
413        let result = Source::from_iter([1_u64, 2])
414            .concat(Source::failed(error.clone()))
415            .instrumented("failure", &registry)
416            .run_collect();
417
418        assert_eq!(result, Err(error));
419        let snapshot = registry
420            .snapshots()
421            .into_iter()
422            .next()
423            .expect("snapshot registered");
424        assert_eq!(snapshot.elements_through, 2);
425        assert_eq!(snapshot.state, StreamInstrumentationState::Failed);
426        assert!(snapshot.finished_at.is_some());
427    }
428
429    #[test]
430    fn enabled_counters_mark_cancelled_stream_as_draining() {
431        let registry = StreamInstrumentationRegistry::new();
432        let emitted = Arc::new(AtomicBool::new(false));
433        let source = {
434            let emitted = Arc::clone(&emitted);
435            Source::from_materialized_factory(move |_materializer| {
436                let emitted = Arc::clone(&emitted);
437                Ok((
438                    Box::new(std::iter::from_fn(move || {
439                        if !emitted.swap(true, Ordering::SeqCst) {
440                            return Some(Ok(1_u64));
441                        }
442                        loop {
443                            if crate::stream::current_stream_cancelled()
444                                .as_ref()
445                                .is_some_and(|cancelled| cancelled.load(Ordering::SeqCst))
446                            {
447                                return Some(Err(StreamError::Cancelled));
448                            }
449                            thread::park_timeout(Duration::from_millis(1));
450                        }
451                    })) as BoxStream<u64>,
452                    NotUsed,
453                ))
454            })
455        };
456
457        let completion = source
458            .instrumented("cancel", &registry)
459            .to_mat(Sink::ignore(), Keep::right)
460            .run()
461            .expect("stream materializes");
462
463        wait_for_snapshot(&registry, |snapshot| snapshot.elements_through == 1);
464        drop(completion);
465        let snapshot = wait_for_snapshot(&registry, |snapshot| {
466            snapshot.state == StreamInstrumentationState::Draining
467        });
468
469        assert_eq!(snapshot.elements_through, 1);
470        assert_eq!(snapshot.state, StreamInstrumentationState::Draining);
471    }
472
473    #[test]
474    fn disabled_behavior_matches_plain_source() {
475        let plain = Source::from_iter(0_u64..8)
476            .map(|item| item * 2)
477            .run_collect()
478            .expect("plain stream succeeds");
479        let registry = StreamInstrumentationRegistry::new();
480        let instrumented = Source::from_iter(0_u64..8)
481            .map(|item| item * 2)
482            .instrumented("enabled", &registry)
483            .run_collect()
484            .expect("instrumented stream succeeds");
485
486        assert_eq!(plain, instrumented);
487        assert!(StreamInstrumentationRegistry::new().snapshots().is_empty());
488    }
489
490    #[test]
491    fn run_handle_records_restarts() {
492        let registry = StreamInstrumentationRegistry::new();
493        let run = registry.register("job");
494
495        run.record_restart();
496        run.record_restarts(2);
497
498        let snapshot = registry.snapshot(run.id()).expect("snapshot retained");
499        assert_eq!(snapshot.restarts, 3);
500        assert_eq!(snapshot.state, StreamInstrumentationState::Running);
501    }
502
503    fn wait_for_snapshot(
504        registry: &StreamInstrumentationRegistry,
505        predicate: impl Fn(&StreamInstrumentationSnapshot) -> bool,
506    ) -> StreamInstrumentationSnapshot {
507        for _ in 0..500 {
508            if let Some(snapshot) = registry.snapshots().into_iter().next()
509                && predicate(&snapshot)
510            {
511                return snapshot;
512            }
513            thread::sleep(Duration::from_millis(2));
514        }
515        panic!("timed out waiting for instrumentation snapshot");
516    }
517}
datum/instrument.rs

datum/
instrument.rs
