tracing_samply/
lib.rs

1#![doc = include_str!("../README.md")]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3
4use smallvec::{SmallVec, smallvec};
5use std::{
6    cell::RefCell,
7    fs::{File, OpenOptions},
8    io::{self, BufWriter, Write},
9    path::{Path, PathBuf},
10};
11use tracing_core::{Subscriber, span};
12use tracing_subscriber::{Layer, layer::Context, registry::LookupSpan};
13
14thread_local! {
15    static MARKER_FILE: RefCell<Option<MarkerFile>> = const { RefCell::new(None) };
16}
17
18/// [`SamplyLayer`] builder.
19///
20/// See the [crate docs](crate) for more information.
21pub struct SamplyLayerBuilder {
22    output_dir: Option<PathBuf>,
23}
24
25impl Default for SamplyLayerBuilder {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31impl SamplyLayerBuilder {
32    /// Creates a new [`SamplyLayerBuilder`].
33    pub fn new() -> Self {
34        Self { output_dir: None }
35    }
36
37    /// Sets the output directory for intermediate files.
38    ///
39    /// If unset, a temporary directory will be created and used.
40    pub fn output_dir(mut self, dir: impl Into<PathBuf>) -> Self {
41        self.output_dir = Some(dir.into());
42        self
43    }
44
45    /// Builds a new [`SamplyLayer`].
46    pub fn build(self) -> io::Result<SamplyLayer> {
47        let Self { output_dir } = self;
48        let dir = match &output_dir {
49            Some(dir) => dir,
50            None => &*std::env::temp_dir().join("tracing-samply"),
51        };
52        let dir = dir.join(std::process::id().to_string());
53        if cfg!(unix) {
54            std::fs::create_dir_all(&dir)
55                .map_err(map_io_err("could not create perf markers dir", &dir))?;
56        }
57        Ok(SamplyLayer { dir: dir.into_boxed_path() })
58    }
59}
60
61/// A tracing layer that bridges `tracing` events and spans with `samply`.
62///
63/// See the [crate docs](crate) for more information.
64pub struct SamplyLayer {
65    dir: Box<Path>,
66}
67
68struct SpanDataStack {
69    stack: SmallVec<[SpanData; 1]>,
70}
71struct SpanData {
72    start_ts: u64,
73}
74
75impl SamplyLayer {
76    /// Creates a new [`SamplyLayer`].
77    ///
78    /// This is the same as `SamplyLayer::builder().build()`.
79    pub fn new() -> io::Result<Self> {
80        Self::builder().build()
81    }
82
83    /// Creates a new [`SamplyLayer`] builder.
84    pub fn builder() -> SamplyLayerBuilder {
85        SamplyLayerBuilder::new()
86    }
87
88    fn create_marker_file(&self) -> MarkerFile {
89        match self.try_create_marker_file() {
90            Ok(file) => file,
91            Err(err) => panic!("{err}"),
92        }
93    }
94
95    fn try_create_marker_file(&self) -> io::Result<MarkerFile> {
96        let pid = std::process::id();
97        let fname = match gettid() {
98            Some(tid) => format!("marker-{pid}-{tid}.txt"),
99            None => format!("marker-{pid}.txt"),
100        };
101        let path = &*self.dir.join(fname);
102        let file = OpenOptions::new()
103            .create_new(true)
104            .read(true)
105            .write(true)
106            .open(path)
107            .map_err(map_io_err("could not create perf markers file", path))?;
108        // mmap the file to notify samply.
109        // Linux perf needs `exec` permission to record it in perf.data.
110        // On macOS, samply only needs the file to be opened, not mmap'ed.
111        #[cfg(all(unix, not(target_vendor = "apple")))]
112        let _ = unsafe {
113            memmap2::MmapOptions::new()
114                .map_exec(&file)
115                .map_err(map_io_err("could not mmap perf markers file", path))?
116        };
117        Ok(MarkerFile::new(file))
118    }
119}
120
121impl<S> Layer<S> for SamplyLayer
122where
123    S: Subscriber + for<'a> LookupSpan<'a>,
124{
125    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
126        if !cfg!(unix) {
127            return;
128        }
129        let Some(span) = ctx.span(id) else { return };
130        let data = SpanData { start_ts: now_timestamp() };
131        let mut extensions = span.extensions_mut();
132        if let Some(stack) = extensions.get_mut::<SpanDataStack>() {
133            stack.stack.push(data);
134        } else {
135            extensions.insert(SpanDataStack { stack: smallvec![data] });
136        }
137    }
138
139    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
140        if !cfg!(unix) {
141            return;
142        }
143        let Some(span) = ctx.span(id) else { return };
144        let mut extensions = span.extensions_mut();
145        let Some(data) = extensions.get_mut::<SpanDataStack>() else { return };
146        let Some(SpanData { start_ts }) = data.stack.pop() else { return };
147        let end_ts = now_timestamp();
148        MARKER_FILE.with_borrow_mut(|file| {
149            let file = file.get_or_insert_with(|| self.create_marker_file());
150            file.write_entry(start_ts, end_ts, span.name());
151        });
152    }
153}
154
155struct MarkerFile {
156    file: BufWriter<File>,
157}
158
159impl MarkerFile {
160    fn new(file: File) -> Self {
161        Self { file: BufWriter::new(file) }
162    }
163
164    fn write_entry(&mut self, start_ts: u64, end_ts: u64, name: &str) {
165        let _ = self.file.write_all(itoa::Buffer::new().format(start_ts).as_bytes());
166        let _ = self.file.write_all(b" ");
167        let _ = self.file.write_all(itoa::Buffer::new().format(end_ts).as_bytes());
168        let _ = self.file.write_all(b" ");
169        let _ = self.file.write_all(name.as_bytes());
170        let _ = self.file.write_all(b"\n");
171    }
172}
173
174fn now_timestamp() -> u64 {
175    cfg_if::cfg_if! {
176        if #[cfg(target_vendor = "apple")] {
177            // https://github.com/mstange/samply/blob/2041b956f650bb92d912990052967d03fef66b75/samply/src/mac/time.rs#L7
178            use std::sync::OnceLock;
179            use mach2::mach_time;
180
181            static NANOS_PER_TICK: OnceLock<mach_time::mach_timebase_info> = OnceLock::new();
182
183            let nanos_per_tick = NANOS_PER_TICK.get_or_init(|| unsafe {
184                let mut info = mach_time::mach_timebase_info::default();
185                let errno = mach_time::mach_timebase_info(&mut info as *mut _);
186                if errno != 0 || info.denom == 0 {
187                    info.numer = 1;
188                    info.denom = 1;
189                };
190                info
191            });
192
193            let time = unsafe { mach_time::mach_absolute_time() };
194
195            time * nanos_per_tick.numer as u64 / nanos_per_tick.denom as u64
196        } else if #[cfg(unix)] {
197            let mut ts = unsafe { std::mem::zeroed() };
198            if unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC, &mut ts) } != 0 {
199                return u64::MAX;
200            }
201            std::time::Duration::new(ts.tv_sec as _, ts.tv_nsec as _)
202                .as_nanos()
203                .try_into()
204                .unwrap_or(u64::MAX)
205        } else {
206            0
207        }
208    }
209}
210
211fn gettid() -> Option<u64> {
212    // https://github.com/rust-lang/rust/blob/9044e98b66d074e7f88b1d4cea58bb0538f2eda6/library/std/src/sys/thread/unix.rs#L325
213    cfg_if::cfg_if! {
214        if #[cfg(target_vendor = "apple")] {
215            let mut tid = 0u64;
216            let status = unsafe { libc::pthread_threadid_np(0, &mut tid) };
217            (status == 0).then_some(tid)
218        } else if #[cfg(unix)] {
219            Some(unsafe { libc::gettid() } as u64)
220        // } else if #[cfg(windows)] {
221        //     let tid = unsafe { c::GetCurrentThreadId() } as u64;
222        //     if tid == 0 { None } else { Some(tid as _) }
223        } else {
224            None
225        }
226    }
227}
228
229fn map_io_err(s: &str, p: &Path) -> impl FnOnce(io::Error) -> io::Error {
230    move |e| io::Error::new(e.kind(), format!("{s} {p:?}: {e}"))
231}
232
233// Not public API. Only for testing purposes.
234#[doc(hidden)]
235pub mod __private {
236    use super::*;
237
238    pub fn flush_marker_file() {
239        MARKER_FILE.with_borrow_mut(|file| {
240            if let Some(file) = file {
241                let _ = file.file.flush();
242            }
243        });
244    }
245}
tracing_samply/lib.rs

tracing_samply/
lib.rs
