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}