Skip to main content

micromegas_tracing/
guards.rs

1//! RAII-style guards
2use std::{collections::HashMap, marker::PhantomData, sync::Arc};
3
4use crate::{
5    dispatch::{
6        flush_log_buffer, flush_metrics_buffer, flush_thread_buffer, init_event_dispatch,
7        init_thread_stream, on_begin_scope, on_end_scope, shutdown_dispatch,
8    },
9    errors::Result,
10    event::EventSink,
11    panic_hook::init_panic_hook,
12    spans::{SpanLocation, SpanMetadata, SpanName, StaticSpanName},
13};
14
15pub struct TracingSystemGuard {}
16
17impl TracingSystemGuard {
18    /// Instantiates a new system guard that initializes the telemetry system.
19    pub fn new(
20        logs_buffer_size: usize,
21        metrics_buffer_size: usize,
22        threads_buffer_size: usize,
23        sink: Arc<dyn EventSink>,
24        process_properties: HashMap<String, String>,
25        cpu_tracing_enabled: bool,
26    ) -> Result<Self> {
27        init_telemetry(
28            logs_buffer_size,
29            metrics_buffer_size,
30            threads_buffer_size,
31            sink,
32            process_properties,
33            cpu_tracing_enabled,
34        )?;
35        Ok(Self {})
36    }
37}
38
39impl Drop for TracingSystemGuard {
40    fn drop(&mut self) {
41        shutdown_telemetry();
42    }
43}
44
45pub fn init_telemetry(
46    logs_buffer_size: usize,
47    metrics_buffer_size: usize,
48    threads_buffer_size: usize,
49    sink: Arc<dyn EventSink>,
50    process_properties: HashMap<String, String>,
51    cpu_tracing_enabled: bool,
52) -> Result<()> {
53    init_event_dispatch(
54        logs_buffer_size,
55        metrics_buffer_size,
56        threads_buffer_size,
57        sink,
58        process_properties,
59        cpu_tracing_enabled,
60    )?;
61    init_panic_hook();
62    Ok(())
63}
64
65pub fn shutdown_telemetry() {
66    flush_log_buffer();
67    flush_metrics_buffer();
68    shutdown_dispatch();
69}
70
71pub struct TracingThreadGuard {
72    _dummy_ptr: *mut u8,
73}
74
75impl TracingThreadGuard {
76    pub fn new() -> Self {
77        init_thread_stream();
78        Self {
79            _dummy_ptr: std::ptr::null_mut(),
80        }
81    }
82}
83
84impl Drop for TracingThreadGuard {
85    fn drop(&mut self) {
86        flush_thread_buffer();
87    }
88}
89
90//not used at the time of writing, but clippy wants it
91impl Default for TracingThreadGuard {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97// sync scope guard
98pub struct ThreadSpanGuard {
99    thread_span_desc: &'static SpanMetadata,
100    _dummy_ptr: PhantomData<*mut u8>, // to mark the object as !Send
101}
102
103impl ThreadSpanGuard {
104    pub fn new(thread_span_desc: &'static SpanMetadata) -> Self {
105        let guard = Self {
106            thread_span_desc,
107            _dummy_ptr: std::marker::PhantomData {},
108        };
109        on_begin_scope(thread_span_desc);
110        guard
111    }
112}
113
114impl Drop for ThreadSpanGuard {
115    fn drop(&mut self) {
116        on_end_scope(self.thread_span_desc);
117    }
118}
119
120pub struct ThreadNamedSpanGuard<N: SpanName = StaticSpanName> {
121    thread_span_location: &'static SpanLocation,
122    name: N,
123    _dummy_ptr: PhantomData<*mut u8>, // to mark the object as !Send
124}
125
126impl ThreadNamedSpanGuard {
127    pub fn new(thread_span_location: &'static SpanLocation, name: &'static str) -> Self {
128        Self::with_name(thread_span_location, StaticSpanName(name))
129    }
130}
131
132impl<N: SpanName> ThreadNamedSpanGuard<N> {
133    pub fn with_name(thread_span_location: &'static SpanLocation, name: N) -> Self {
134        let guard = Self {
135            thread_span_location,
136            name,
137            _dummy_ptr: std::marker::PhantomData {},
138        };
139        guard.name.begin_thread(guard.thread_span_location);
140        guard
141    }
142}
143
144impl<N: SpanName> Drop for ThreadNamedSpanGuard<N> {
145    fn drop(&mut self) {
146        self.name.end_thread(self.thread_span_location);
147    }
148}