graft_tracing/
lib.rs

1//! Tracing utilities for the Graft project.
2//!
3//! This crate provides functionality for initializing and configuring
4//! [tracing](https://docs.rs/tracing) in different environments (test, server, tool).
5
6use parking_lot::Once;
7use std::time::Instant;
8use tracing_subscriber::{
9    fmt::{MakeWriter, time::SystemTime},
10    util::SubscriberInitExt,
11};
12
13use tracing::level_filters::LevelFilter;
14use tracing_subscriber::{
15    EnvFilter,
16    fmt::{
17        format::{FmtSpan, Writer},
18        time::FormatTime,
19    },
20};
21
22/// Checks if the application is running in the Antithesis testing environment.
23pub fn running_in_antithesis() -> bool {
24    std::env::var("ANTITHESIS_OUTPUT_DIR").is_ok()
25}
26
27/// Specifies the type of application consuming the tracing output.
28#[derive(PartialEq, Eq)]
29pub enum TracingConsumer {
30    /// Test environment consumer
31    Test,
32    /// Server application consumer
33    Server,
34    /// Command-line tool consumer
35    Tool,
36}
37
38/// Initializes tracing with stdout as the output.
39pub fn init_tracing(consumer: TracingConsumer, process_id: Option<String>) {
40    init_tracing_with_writer(consumer, process_id, std::io::stdout);
41}
42
43/// Initializes tracing with a custom writer for output.
44///
45/// # Parameters
46/// * `consumer` - The type of application consuming the tracing output
47/// * `process_id` - Optional identifier for the process, randomly generated if None
48/// * `writer` - Custom writer implementation for tracing output
49///
50/// # Type Parameters
51/// * `W` - Writer type that implements the [`tracing_subscriber::fmt::MakeWriter`] trait
52pub fn init_tracing_with_writer<W>(consumer: TracingConsumer, process_id: Option<String>, writer: W)
53where
54    W: for<'writer> MakeWriter<'writer> + 'static + Send + Sync,
55{
56    static INIT: Once = Once::new();
57    INIT.call_once(move || {
58        let process_id = process_id
59            .unwrap_or_else(|| bs58::encode(rand::random::<u64>().to_le_bytes()).into_string());
60
61        let antithesis = running_in_antithesis();
62        let testing = consumer == TracingConsumer::Test;
63        let color = !antithesis && !std::env::var("NO_COLOR").is_ok_and(|s| !s.is_empty());
64
65        let default_level = if consumer == TracingConsumer::Tool {
66            LevelFilter::WARN
67        } else {
68            LevelFilter::INFO
69        };
70
71        let mut filter = EnvFilter::builder()
72            .with_default_directive(default_level.into())
73            .from_env()
74            .unwrap();
75
76        let mut span_events = FmtSpan::NONE;
77
78        if antithesis || testing {
79            span_events = FmtSpan::NEW | FmtSpan::CLOSE;
80            filter = filter
81                .add_directive("graft_client=trace".parse().unwrap())
82                .add_directive("graft_core=trace".parse().unwrap())
83                .add_directive("graft_server=trace".parse().unwrap())
84                .add_directive("graft_test=trace".parse().unwrap())
85                .add_directive("graft_sqlite=debug".parse().unwrap())
86        }
87
88        let prefix = if antithesis || testing {
89            Some(process_id.clone())
90        } else {
91            None
92        };
93
94        let time = if antithesis {
95            TimeFormat::None
96        } else if consumer == TracingConsumer::Server {
97            TimeFormat::Long(SystemTime)
98        } else {
99            TimeFormat::Offset { start: Instant::now() }
100        };
101
102        tracing_subscriber::fmt()
103            .with_env_filter(filter)
104            .with_thread_names(true)
105            .with_span_events(span_events)
106            .with_ansi(color)
107            .with_timer(TimeAndPrefix::new(prefix, time))
108            .with_writer(writer)
109            .finish()
110            .try_init()
111            .expect("failed to setup tracing subscriber");
112    });
113}
114
115enum TimeFormat {
116    None,
117    Long(SystemTime),
118    Offset { start: Instant },
119}
120
121struct TimeAndPrefix {
122    prefix: Option<String>,
123    time: TimeFormat,
124}
125
126impl TimeAndPrefix {
127    fn new(prefix: Option<String>, time: TimeFormat) -> Self {
128        Self { prefix, time }
129    }
130    fn write_time(&self, w: &mut Writer<'_>) -> std::fmt::Result {
131        match self.time {
132            TimeFormat::None => Ok(()),
133            TimeFormat::Long(inner) => inner.format_time(w),
134            TimeFormat::Offset { start } => {
135                let e = start.elapsed();
136                let nanos = e.subsec_nanos();
137                // round nanos to the nearest millisecond
138                let millis = (nanos as f64 / 1_000_000.0).round();
139                write!(w, "{:03}.{:03}s", e.as_secs(), millis)
140            }
141        }
142    }
143}
144
145impl FormatTime for TimeAndPrefix {
146    fn format_time(&self, w: &mut Writer<'_>) -> std::fmt::Result {
147        match (&self.prefix, &self.time) {
148            (None, _) => self.write_time(w),
149            (Some(prefix), TimeFormat::None) => write!(w, "{prefix}"),
150            (Some(prefix), _) => {
151                write!(w, "{prefix} ")?;
152                self.write_time(w)
153            }
154        }
155    }
156}
graft_tracing/lib.rs

graft_tracing/
lib.rs
