use std::sync::atomic::{AtomicBool, Ordering};
use crate::buffer::{BufferMode, Disk, SegmentWriter};
use crate::clock;
use crate::handle::Dial9Handle;
use crate::primitives::sync::Arc;
use crate::recording::{Recorder, RecordingStartHook};
use crate::shared_state::SharedState;
use crate::source::Source;
type RecordingThreadHook = Arc<dyn Fn() -> Box<dyn FnOnce() + Send> + Send + Sync>;
fn noop_thread_hook() -> RecordingThreadHook {
Arc::new(|| Box::new(|| {}) as Box<dyn FnOnce() + Send>)
}
pub(crate) fn merge_segment_metadata(
existing: &mut Vec<(String, String)>,
entries: impl IntoIterator<Item = (String, String)>,
) {
let incoming: Vec<(String, String)> = entries.into_iter().collect();
existing.retain(|(k, _)| !incoming.iter().any(|(ik, _)| ik == k));
existing.extend(incoming);
}
pub fn recorder<M: BufferMode>(writer: SegmentWriter<M>) -> RecorderBuilder<M> {
builder_with(Some(writer))
}
pub fn recorder_or_disabled<M: BufferMode>(
writer: std::io::Result<SegmentWriter<M>>,
) -> RecorderBuilder<M> {
match writer {
Ok(writer) => builder_with(Some(writer)),
Err(e) => {
tracing::error!(
target: "dial9_telemetry",
"dial9: trace writer setup failed; running without telemetry: {e}"
);
builder_with(None)
}
}
}
fn builder_with<M: BufferMode>(writer: Option<SegmentWriter<M>>) -> RecorderBuilder<M> {
RecorderBuilder {
writer,
sources: Vec::new(),
recording_start_hooks: Vec::new(),
segment_metadata: Vec::new(),
metrics_sink: None,
thread_init: noop_thread_hook(),
#[cfg(feature = "pipeline")]
pipeline: None,
#[cfg(feature = "pipeline")]
terminal_processor: None,
#[cfg(feature = "pipeline")]
worker_poll_interval: None,
#[cfg(feature = "pipeline")]
trigger: None,
#[cfg(feature = "pipeline")]
pending_dump_trigger: None,
enabled: true,
}
}
pub fn recorder_disabled() -> crate::recording::Recorder {
crate::recording::Recorder::new(crate::handle::Dial9Handle::disabled(), None)
}
#[derive(Debug)]
pub(crate) struct SoleRecorderGuard;
impl SoleRecorderGuard {
pub(crate) fn claim() -> Option<Self> {
match cfg!(feature = "test-util") || !RECORDER_TAKEN.swap(true, Ordering::SeqCst) {
true => Some(Self),
false => None,
}
}
}
impl Drop for SoleRecorderGuard {
fn drop(&mut self) {
RECORDER_TAKEN.store(false, Ordering::SeqCst);
}
}
static RECORDER_TAKEN: AtomicBool = AtomicBool::new(false);
#[cfg(feature = "pipeline")]
fn default_pipeline(
source_stages: Vec<Box<dyn crate::pipeline::SegmentProcessor>>,
terminal: Option<Box<dyn crate::pipeline::SegmentProcessor>>,
is_disk: bool,
) -> Vec<Box<dyn crate::pipeline::SegmentProcessor>> {
if source_stages.is_empty() && terminal.is_none() {
return Vec::new();
}
let mut processors = source_stages;
processors.push(Box::new(crate::worker::processors::GzipCompressor));
match terminal {
Some(terminal) => processors.push(terminal),
None if is_disk => processors.push(Box::new(
crate::worker::processors::WriteBackProcessor::default(),
)),
None => {}
}
processors
}
#[must_use = "call `.build()` to start recording"]
pub struct RecorderBuilder<M: BufferMode = Disk> {
writer: Option<SegmentWriter<M>>,
sources: Vec<Box<dyn Source>>,
recording_start_hooks: Vec<RecordingStartHook>,
segment_metadata: Vec<(String, String)>,
metrics_sink: Option<metrique::writer::BoxEntrySink>,
thread_init: RecordingThreadHook,
#[cfg(feature = "pipeline")]
pipeline: Option<Vec<Box<dyn crate::pipeline::SegmentProcessor>>>,
#[cfg(feature = "pipeline")]
terminal_processor: Option<Box<dyn crate::pipeline::SegmentProcessor>>,
#[cfg(feature = "pipeline")]
worker_poll_interval: Option<std::time::Duration>,
#[cfg(feature = "pipeline")]
trigger: Option<crate::dump::DumpRx>,
#[cfg(feature = "pipeline")]
pending_dump_trigger: Option<crate::dump::DumpTrigger>,
enabled: bool,
}
impl<M: BufferMode> std::fmt::Debug for RecorderBuilder<M> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RecorderBuilder")
.field("sources", &self.sources.len())
.finish_non_exhaustive()
}
}
impl<M: BufferMode> RecorderBuilder<M> {
pub fn source(mut self, source: impl Source + 'static) -> Self {
self.sources.push(Box::new(source));
self
}
pub fn source_names(&self) -> impl Iterator<Item = &str> + '_ {
self.sources.iter().map(|s| s.name())
}
pub fn writer_boot_id(&self) -> Option<&str> {
self.writer.as_ref()?.boot_id()
}
pub fn segment_metadata(mut self, entries: impl IntoIterator<Item = (String, String)>) -> Self {
merge_segment_metadata(&mut self.segment_metadata, entries);
self
}
pub fn metrics_sink(mut self, sink: metrique::writer::BoxEntrySink) -> Self {
self.metrics_sink = Some(sink);
self
}
pub fn build(self) -> Recorder {
#[allow(unused_mut)]
let Some(mut writer) = self.writer else {
return recorder_disabled();
};
let Some(sole_recorder) = SoleRecorderGuard::claim() else {
tracing::error!(
target: "dial9",
"dial9: this process already has a recorder, this one will run without telemetry."
);
return recorder_disabled();
};
let shared = Arc::new(SharedState::new(clock::clock_monotonic_ns()));
#[cfg(feature = "pipeline")]
if let Some(trigger) = self.pending_dump_trigger {
shared.set_dump_trigger(trigger);
}
let mut segment_metadata = self.segment_metadata;
if let Some(boot_id) = writer.boot_id().map(str::to_owned) {
for (key, value) in &mut segment_metadata {
if key == "boot_id" {
*value = boot_id.clone();
}
}
}
if !segment_metadata.is_empty() {
writer.update_segment_metadata(segment_metadata);
}
#[allow(unused_mut)]
let mut sources = self.sources;
#[cfg(feature = "pipeline")]
let processors = match self.pipeline {
Some(processors) => processors,
None => {
let source_stages = sources
.iter_mut()
.filter_map(|source| source.segment_processor())
.collect();
default_pipeline(source_stages, self.terminal_processor, M::IS_DISK)
}
};
for source in sources {
shared.push_source(source);
}
#[cfg(feature = "pipeline")]
let worker = if processors.is_empty() {
None
} else {
let poll = self
.worker_poll_interval
.unwrap_or(crate::worker::DEFAULT_POLL_INTERVAL);
let metrics = self
.metrics_sink
.clone()
.unwrap_or_else(metrique::writer::sink::DevNullSink::boxed);
let config = crate::worker::BackgroundTaskConfig::builder()
.maybe_trace_dir(M::IS_DISK.then(|| writer.trace_dir().to_path_buf()))
.maybe_trace_stem(M::IS_DISK.then(|| writer.trace_stem().to_string()))
.poll_interval(poll)
.processors(processors)
.metrics_sink(metrics)
.maybe_trigger(self.trigger)
.build();
let (tx, rx) = tokio::sync::oneshot::channel();
let hook = self.thread_init.clone();
crate::worker::spawn(&writer, config, rx, move || hook())
.map(|wt| crate::recording::WorkerHandle::new(tx, wt))
};
let hook = self.thread_init.clone();
#[allow(unused_mut)]
let mut recorder = Recorder::start(shared, writer, self.metrics_sink, move || hook());
recorder.hold_process(sole_recorder);
#[cfg(feature = "pipeline")]
if let Some(worker) = worker {
recorder.attach_worker(worker);
}
recorder.set_recording_start_hooks(self.recording_start_hooks);
if self.enabled {
recorder.enable();
}
recorder
}
pub fn paused(mut self) -> Self {
self.enabled = false;
self
}
}
pub trait RecorderSourceExt: Sized {
fn source(self, source: impl Source + 'static) -> Self;
fn on_recording_start(self, hook: impl FnOnce(&Dial9Handle) + Send + 'static) -> Self;
fn on_recording_thread_start<F, T>(self, hook: F) -> Self
where
F: Fn() -> T + Send + Sync + 'static,
T: FnOnce() + Send + 'static;
fn with_custom_events<F>(
self,
config: crate::custom_events::CustomEventsConfig,
callback: F,
) -> Self
where
F: for<'a> FnMut(&mut crate::custom_events::CustomEventsContext<'a>) + Send + 'static,
{
self.source(crate::custom_events::CustomEventsSource::new(
config, callback,
))
}
}
impl<M: BufferMode> RecorderSourceExt for RecorderBuilder<M> {
fn source(mut self, source: impl Source + 'static) -> Self {
self.sources.push(Box::new(source));
self
}
fn on_recording_start(mut self, hook: impl FnOnce(&Dial9Handle) + Send + 'static) -> Self {
self.recording_start_hooks.push(Box::new(hook));
self
}
fn on_recording_thread_start<F, T>(mut self, hook: F) -> Self
where
F: Fn() -> T + Send + Sync + 'static,
T: FnOnce() + Send + 'static,
{
self.thread_init = Arc::new(move || Box::new(hook()) as Box<dyn FnOnce() + Send>);
self
}
}
#[cfg(feature = "pipeline")]
impl<M: BufferMode> RecorderBuilder<M> {
pub fn pipe(mut self, processor: impl crate::pipeline::SegmentProcessor + 'static) -> Self {
self.pipeline
.get_or_insert_default()
.push(Box::new(processor));
self
}
pub fn processors(
mut self,
processors: Vec<Box<dyn crate::pipeline::SegmentProcessor>>,
) -> Self {
self.pipeline = Some(processors);
self
}
pub fn terminal_processor(
mut self,
processor: impl crate::pipeline::SegmentProcessor + 'static,
) -> Self {
self.terminal_processor = Some(Box::new(processor));
self
}
pub fn worker_poll_interval(mut self, interval: std::time::Duration) -> Self {
self.worker_poll_interval = Some(interval);
self
}
pub fn trigger(mut self, trigger: crate::dump::DumpRx) -> Self {
self.trigger = Some(trigger);
self
}
pub fn with_dump_trigger<F>(mut self, configure: F) -> Self
where
F: FnOnce(&mut crate::dump::DumpTriggerConfig),
{
let mut config = crate::dump::DumpTriggerConfig::new();
configure(&mut config);
let (mut trigger, rx) = crate::dump::channel();
if let Some(window) = config.debounce_window() {
trigger = trigger.with_debounce(window);
}
self.trigger = Some(rx);
self.pending_dump_trigger = Some(trigger);
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::buffer::{DiskBuffer, MemoryBuffer};
use crate::source::FlushContext;
use dial9_trace_format::TraceEvent;
use dial9_trace_format::decoder::Decoder;
use std::path::{Path, PathBuf};
use std::time::Duration;
#[derive(Debug, serde::Deserialize, TraceEvent)]
struct TestEvent {
#[traceevent(timestamp)]
timestamp_ns: u64,
value: u64,
}
struct OnceSource {
emitted: bool,
value: u64,
}
impl Source for OnceSource {
fn flush(&mut self, ctx: &FlushContext<'_>) {
if !self.emitted {
self.emitted = true;
ctx.record_event(&TestEvent {
timestamp_ns: clock::clock_monotonic_ns(),
value: self.value,
});
}
}
fn name(&self) -> &'static str {
"once"
}
}
fn sealed_segment(dir: &Path) -> PathBuf {
std::fs::read_dir(dir)
.expect("trace dir readable")
.filter_map(|e| e.ok().map(|e| e.path()))
.find(|p| {
let name = p.file_name().unwrap().to_string_lossy();
name.ends_with(".bin") && !name.ends_with(".active")
})
.expect("a sealed .bin segment")
}
fn decoded_test_values(bytes: &[u8]) -> Vec<u64> {
let mut decoder = Decoder::new(bytes).expect("valid trace header");
let mut values = Vec::new();
decoder
.for_each_event(|raw| {
if raw.name == "TestEvent" {
let event: TestEvent = raw.deserialize().expect("TestEvent decodes");
values.push(event.value);
}
})
.expect("decode events");
values
}
#[test]
fn records_source_events_to_disk() {
let dir = tempfile::tempdir().expect("tempdir");
let writer = DiskBuffer::single_file(dir.path().join("trace.bin")).expect("writer");
let recorder = recorder(writer)
.segment_metadata([("service".to_string(), "recorder-test".to_string())])
.source(OnceSource {
emitted: false,
value: 7,
})
.build();
recorder.graceful_shutdown(Duration::ZERO);
let bytes = std::fs::read(sealed_segment(dir.path())).expect("read segment");
assert!(
decoded_test_values(&bytes).contains(&7),
"the source's event should round-trip through the trace file"
);
}
#[test]
fn build_starts_recording() {
let writer = MemoryBuffer::new(1 << 20).expect("writer");
let recorder = recorder(writer).build();
assert!(
recorder.shared().expect("live recorder").is_enabled(),
"build() must start recording"
);
}
#[test]
fn paused_build_waits_for_enable() {
let writer = MemoryBuffer::new(1 << 20).expect("writer");
let recorder = recorder(writer).paused().build();
assert!(
!recorder.shared().expect("live recorder").is_enabled(),
"paused() must leave recording off"
);
assert!(
!recorder.handle().is_enabled(),
"a paused handle must report disabled"
);
assert!(
recorder.handle().shared().is_some(),
"a paused handle is still connected"
);
recorder.enable();
assert!(
recorder.shared().expect("live recorder").is_enabled(),
"recording on after enable()"
);
assert!(
recorder.handle().is_enabled(),
"handle enabled after enable()"
);
}
#[test]
fn on_recording_start_runs_once_when_recording_starts() {
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
let runs = StdArc::new(AtomicUsize::new(0));
let runs_hook = StdArc::clone(&runs);
let live = recorder(MemoryBuffer::new(1 << 20).expect("writer"))
.on_recording_start(move |_handle| {
runs_hook.fetch_add(1, Ordering::SeqCst);
})
.build();
assert_eq!(runs.load(Ordering::SeqCst), 1, "hook runs at build");
live.enable();
assert_eq!(runs.load(Ordering::SeqCst), 1, "hook runs at most once");
drop(live);
let paused_runs = StdArc::new(AtomicUsize::new(0));
let paused_hook = StdArc::clone(&paused_runs);
let paused = recorder(MemoryBuffer::new(1 << 20).expect("writer"))
.on_recording_start(move |_handle| {
paused_hook.fetch_add(1, Ordering::SeqCst);
})
.paused()
.build();
assert_eq!(
paused_runs.load(Ordering::SeqCst),
0,
"paused build must not run the hook yet"
);
paused.enable();
assert_eq!(paused_runs.load(Ordering::SeqCst), 1, "hook runs on enable");
}
#[cfg(feature = "pipeline")]
#[test]
fn pipe_runs_the_background_worker() {
use crate::pipeline::{ProcessError, SegmentData, SegmentProcessor};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Debug)]
struct CountingProcessor(StdArc<AtomicUsize>);
impl SegmentProcessor for CountingProcessor {
fn name(&self) -> &'static str {
"Counting"
}
fn process(
&mut self,
data: SegmentData,
) -> Pin<Box<dyn Future<Output = Result<SegmentData, ProcessError>> + Send + '_>>
{
self.0.fetch_add(1, Ordering::SeqCst);
Box::pin(async move { Ok(data) })
}
}
let dir = tempfile::tempdir().expect("tempdir");
let writer = DiskBuffer::single_file(dir.path().join("trace.bin")).expect("writer");
let processed = StdArc::new(AtomicUsize::new(0));
let recorder = recorder(writer)
.source(OnceSource {
emitted: false,
value: 11,
})
.pipe(CountingProcessor(StdArc::clone(&processed)))
.build();
recorder.graceful_shutdown(Duration::from_secs(5));
assert!(
processed.load(Ordering::SeqCst) >= 1,
"the background worker should process the sealed segment"
);
}
#[test]
fn segment_metadata_merges_last_key_wins() {
let mut md = vec![
("service".to_string(), "checkout".to_string()),
("region".to_string(), "us-east-1".to_string()),
];
super::merge_segment_metadata(
&mut md,
[
("region".to_string(), "eu-west-1".to_string()),
("bucket".to_string(), "traces".to_string()),
],
);
assert!(md.contains(&("service".to_string(), "checkout".to_string())));
assert!(md.contains(&("region".to_string(), "eu-west-1".to_string())));
assert!(md.contains(&("bucket".to_string(), "traces".to_string())));
assert!(
!md.iter().any(|(k, v)| k == "region" && v == "us-east-1"),
"the colliding key's old value must be gone"
);
}
#[cfg(feature = "pipeline")]
#[test]
fn source_stage_joins_the_default_pipeline() {
use crate::pipeline::{ProcessError, SegmentData, SegmentProcessor};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
struct MarkerStage(StdArc<AtomicUsize>);
impl SegmentProcessor for MarkerStage {
fn name(&self) -> &'static str {
"Marker"
}
fn process(
&mut self,
data: SegmentData,
) -> Pin<Box<dyn Future<Output = Result<SegmentData, ProcessError>> + Send + '_>>
{
self.0.fetch_add(1, Ordering::SeqCst);
Box::pin(async move { Ok(data) })
}
}
struct StagedSource(StdArc<AtomicUsize>);
impl Source for StagedSource {
fn flush(&mut self, _ctx: &FlushContext<'_>) {}
fn name(&self) -> &'static str {
"staged"
}
fn segment_processor(&mut self) -> Option<Box<dyn SegmentProcessor>> {
Some(Box::new(MarkerStage(StdArc::clone(&self.0))))
}
}
let dir = tempfile::tempdir().expect("tempdir");
let writer = DiskBuffer::single_file(dir.path().join("trace.bin")).expect("writer");
let ran = StdArc::new(AtomicUsize::new(0));
let recorder = recorder(writer)
.source(StagedSource(StdArc::clone(&ran)))
.source(OnceSource {
emitted: false,
value: 3,
})
.build();
recorder.graceful_shutdown(Duration::from_secs(5));
assert!(
ran.load(Ordering::SeqCst) >= 1,
"the source's stage should run in the default pipeline"
);
let gzipped = std::fs::read_dir(dir.path())
.expect("trace dir readable")
.filter_map(|e| e.ok().map(|e| e.path()))
.any(|p| p.to_string_lossy().ends_with(".bin.gz"));
assert!(gzipped, "the default pipeline compresses and writes back");
}
#[cfg(feature = "pipeline")]
#[test]
fn default_pipeline_is_empty_without_stages() {
let dir = tempfile::tempdir().expect("tempdir");
let writer = DiskBuffer::single_file(dir.path().join("trace.bin")).expect("writer");
let recorder = recorder(writer)
.source(OnceSource {
emitted: false,
value: 5,
})
.build();
recorder.graceful_shutdown(Duration::from_secs(5));
let compressed = std::fs::read_dir(dir.path())
.expect("trace dir readable")
.filter_map(|e| e.ok().map(|e| e.path()))
.any(|p| p.to_string_lossy().ends_with(".gz"));
assert!(
!compressed,
"no stages means no worker, so nothing is gzipped"
);
let bytes = std::fs::read(sealed_segment(dir.path())).expect("read segment");
assert_eq!(decoded_test_values(&bytes), vec![5]);
}
#[cfg(feature = "pipeline")]
#[test]
fn terminal_processor_replaces_write_back() {
use crate::pipeline::{ProcessError, SegmentData, SegmentProcessor};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
struct Uploader(StdArc<AtomicUsize>);
impl SegmentProcessor for Uploader {
fn name(&self) -> &'static str {
"Uploader"
}
fn process(
&mut self,
data: SegmentData,
) -> Pin<Box<dyn Future<Output = Result<SegmentData, ProcessError>> + Send + '_>>
{
self.0.fetch_add(1, Ordering::SeqCst);
Box::pin(async move { Ok(data) })
}
}
let dir = tempfile::tempdir().expect("tempdir");
let writer = DiskBuffer::single_file(dir.path().join("trace.bin")).expect("writer");
let uploaded = StdArc::new(AtomicUsize::new(0));
let recorder = recorder(writer)
.source(OnceSource {
emitted: false,
value: 9,
})
.terminal_processor(Uploader(StdArc::clone(&uploaded)))
.build();
recorder.graceful_shutdown(Duration::from_secs(5));
assert!(
uploaded.load(Ordering::SeqCst) >= 1,
"the terminal stage runs on its own"
);
let written_back = std::fs::read_dir(dir.path())
.expect("trace dir readable")
.filter_map(|e| e.ok().map(|e| e.path()))
.any(|p| p.to_string_lossy().ends_with(".bin.gz"));
assert!(!written_back, "the terminal stage takes write-back's place");
}
#[test]
fn recording_thread_hook_runs_on_dial9_threads() {
use std::sync::Arc as StdArc;
use std::sync::atomic::{AtomicUsize, Ordering};
fn build_and_shut_down<M: BufferMode + Send + 'static>(
writer: crate::buffer::SegmentWriter<M>,
) -> (usize, usize) {
let started = StdArc::new(AtomicUsize::new(0));
let stopped = StdArc::new(AtomicUsize::new(0));
let (s, t) = (StdArc::clone(&started), StdArc::clone(&stopped));
let recorder =
RecorderSourceExt::on_recording_thread_start(recorder(writer), move || {
s.fetch_add(1, Ordering::SeqCst);
let t = StdArc::clone(&t);
move || {
t.fetch_add(1, Ordering::SeqCst);
}
})
.build();
recorder.graceful_shutdown(Duration::from_secs(5));
(
started.load(Ordering::SeqCst),
stopped.load(Ordering::SeqCst),
)
}
let (started, stopped) = build_and_shut_down(MemoryBuffer::new(64 * 1024).unwrap());
assert!(
started >= 1,
"the hook should run on the flush thread, ran {started} times"
);
assert_eq!(
started, stopped,
"every thread that ran the hook should run its teardown"
);
}
}
#[cfg(all(test, not(feature = "test-util")))]
mod single_recorder_tests {
use super::*;
use crate::buffer::MemoryBuffer;
#[test]
fn a_second_recorder_in_the_process_is_refused() {
let first = recorder(MemoryBuffer::new(1 << 20).unwrap()).build();
assert!(first.handle().is_connected(), "the first one records");
let second = recorder(MemoryBuffer::new(1 << 20).unwrap()).build();
assert!(
!second.handle().is_connected(),
"the process already has a recorder, so this one is inert"
);
drop(second);
drop(first);
let after = recorder(MemoryBuffer::new(1 << 20).unwrap()).build();
assert!(
after.handle().is_connected(),
"the slot frees up once the holder stops"
);
}
}