Skip to main content

dial9_core/
source.rs

1//! Source trait for abstracting flush-thread data sources.
2
3use std::any::Any;
4
5use crate::collector::CentralCollector;
6use crate::encoder::{self, Encodable, ThreadLocalEncoder};
7use crate::primitives::sync::Arc;
8use crate::primitives::sync::atomic::AtomicU64;
9
10/// Context passed to [`Source::flush`] for recording events into the trace.
11///
12/// Use [`record_event`] to emit an encodable event, or [`with_encoder`] when
13/// you need direct encoder access (e.g. to intern stack frames).
14///
15/// [`record_event`]: FlushContext::record_event
16/// [`with_encoder`]: FlushContext::with_encoder
17pub struct FlushContext<'a> {
18    collector: &'a Arc<CentralCollector>,
19    drain_epoch: &'a AtomicU64,
20}
21
22impl std::fmt::Debug for FlushContext<'_> {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        f.debug_struct("FlushContext")
25            .field(
26                "drain_epoch",
27                &self.drain_epoch.load(std::sync::atomic::Ordering::Relaxed),
28            )
29            .finish_non_exhaustive()
30    }
31}
32
33impl<'a> FlushContext<'a> {
34    pub(crate) fn new(collector: &'a Arc<CentralCollector>, drain_epoch: &'a AtomicU64) -> Self {
35        Self {
36            collector,
37            drain_epoch,
38        }
39    }
40
41    /// Record an event into the trace from this flush cycle.
42    pub fn record_event(&self, event: &dyn Encodable) {
43        let _ = encoder::record_encodable_event(event, self.collector, self.drain_epoch);
44    }
45
46    /// Record via the thread-local encoder directly.
47    ///
48    /// Use this when you need encoder-level access, e.g. to call
49    /// `enc.intern_stack_frames(..)` before encoding an event.
50    pub fn with_encoder(&self, f: impl FnOnce(&mut ThreadLocalEncoder<'_>)) {
51        let _ = encoder::with_encoder(f, self.collector, self.drain_epoch);
52    }
53}
54
55/// A data source drained by the flush thread each cycle.
56///
57/// Implement this trait to feed custom events into the dial9 trace. Register
58/// the source with [`RecorderBuilder::source`] before starting the flush
59/// thread; the flush thread calls [`flush`] once per cycle.
60///
61/// [`flush`]: Source::flush
62/// [`RecorderBuilder::source`]: crate::recorder::RecorderBuilder::source
63pub trait Source: Any + Send {
64    /// Drain pending data into the trace. Called once per flush cycle.
65    fn flush(&mut self, ctx: &FlushContext<'_>);
66
67    /// Diagnostic name for this source (e.g. `"cpu_profile"`, `"sched"`).
68    fn name(&self) -> &'static str;
69
70    /// Called when a thread joins the recorder: a Tokio worker's first poll, or
71    /// an explicit [`Dial9Handle::track_current_thread`].
72    ///
73    /// Per-thread sources (e.g. `SchedProfiler`) use this to begin tracking
74    /// the current thread. Returns an error if setup fails.
75    ///
76    /// [`Dial9Handle::track_current_thread`]: crate::handle::Dial9Handle::track_current_thread
77    fn on_thread_start(&mut self) -> std::io::Result<()> {
78        Ok(())
79    }
80
81    /// Called when a thread stops. Per-thread sources use this to stop
82    /// tracking the current thread.
83    fn on_thread_stop(&mut self) {}
84
85    /// Append this source's segment-metadata entries to `out` **iff** they have
86    /// changed since the last call.
87    ///
88    /// The flush loop merges only when `out` is non-empty, so appending nothing
89    /// on an unchanged cycle keeps steady-state cycles allocation-free. The
90    /// default reports a source with no metadata.
91    fn segment_metadata(&mut self, out: &mut Vec<(String, String)>) {
92        let _ = out;
93    }
94
95    /// A segment-processing stage this source's data needs, folded into the
96    /// recorder's default pipeline at build. Called once, at build.
97    ///
98    /// A custom pipeline replaces the default outright and does not pick these
99    /// up; you must chain the stage yourself in those cases.
100    #[cfg(feature = "pipeline")]
101    fn segment_processor(&mut self) -> Option<Box<dyn crate::pipeline::SegmentProcessor>> {
102        None
103    }
104}
105
106/// Collect current segment metadata from every source by calling the
107/// change-aware [`Source::segment_metadata`] once each. A freshly-built source
108/// reports its metadata on the first call, so this yields the full set.
109#[cfg(feature = "test-util")]
110pub fn collect_segment_metadata(sources: &mut [Box<dyn Source>]) -> Vec<(String, String)> {
111    let mut out = Vec::new();
112    for source in sources.iter_mut() {
113        source.segment_metadata(&mut out);
114    }
115    out
116}