micromegas_tracing/
runtime.rs

1//! Runtime integration utilities for micromegas tracing
2//!
3//! This module provides helper functions and utilities for integrating
4//! micromegas tracing with async runtimes, particularly tokio.
5
6#[cfg(feature = "tokio")]
7use crate::dispatch::{flush_thread_buffer, init_thread_stream, unregister_thread_stream};
8
9/// Extension trait for `tokio::runtime::Builder` that adds tracing lifecycle callbacks.
10///
11/// This trait provides a convenient way to configure tokio runtimes with the proper
12/// thread lifecycle callbacks for micromegas tracing, ensuring that:
13/// - Thread streams are initialized when worker threads start
14/// - Event buffers are flushed when threads park (become idle)
15/// - Thread streams are properly unregistered when threads stop
16///
17/// This is useful in both production applications (like ingestion servers) and tests
18/// where you need proper tracing integration with tokio's thread pool.
19///
20/// # Examples
21///
22/// ```no_run
23/// use micromegas_tracing::runtime::TracingRuntimeExt;
24///
25/// let runtime = tokio::runtime::Builder::new_multi_thread()
26///     .enable_all()
27///     .thread_name("my-service")
28///     .with_tracing_callbacks()
29///     .build()
30///     .expect("Failed to build runtime");
31/// ```
32#[cfg(feature = "tokio")]
33pub trait TracingRuntimeExt {
34    /// Configures the runtime builder with standard tracing lifecycle callbacks.
35    ///
36    /// This method adds the following callbacks:
37    /// - `on_thread_start`: Initializes thread-local tracing stream
38    /// - `on_thread_park`: Flushes event buffer when thread becomes idle
39    /// - `on_thread_stop`: Unregisters thread stream to prevent dangling pointers
40    fn with_tracing_callbacks(&mut self) -> &mut Self;
41
42    /// Configures the runtime builder with tracing callbacks and custom thread start logic.
43    ///
44    /// This is useful when you need to perform additional setup during thread start
45    /// while still maintaining the standard tracing lifecycle.
46    ///
47    /// # Arguments
48    ///
49    /// * `on_start` - Custom function to call in addition to `init_thread_stream()`
50    ///
51    /// # Examples
52    ///
53    /// ```no_run
54    /// use micromegas_tracing::runtime::TracingRuntimeExt;
55    /// use std::sync::atomic::{AtomicUsize, Ordering};
56    /// use std::sync::Arc;
57    ///
58    /// let counter = Arc::new(AtomicUsize::new(0));
59    /// let counter_clone = counter.clone();
60    ///
61    /// let runtime = tokio::runtime::Builder::new_multi_thread()
62    ///     .enable_all()
63    ///     .with_tracing_callbacks_and_custom_start(move || {
64    ///         let id = counter_clone.fetch_add(1, Ordering::Relaxed);
65    ///         eprintln!("Worker thread {} starting", id);
66    ///     })
67    ///     .build()
68    ///     .expect("Failed to build runtime");
69    /// ```
70    fn with_tracing_callbacks_and_custom_start<F>(&mut self, on_start: F) -> &mut Self
71    where
72        F: Fn() + Send + Sync + 'static;
73}
74
75#[cfg(feature = "tokio")]
76impl TracingRuntimeExt for tokio::runtime::Builder {
77    fn with_tracing_callbacks(&mut self) -> &mut Self {
78        self.on_thread_start(|| {
79            init_thread_stream();
80        })
81        .on_thread_stop(|| {
82            flush_thread_buffer();
83            unregister_thread_stream();
84        })
85    }
86
87    fn with_tracing_callbacks_and_custom_start<F>(&mut self, on_start: F) -> &mut Self
88    where
89        F: Fn() + Send + Sync + 'static,
90    {
91        self.on_thread_start(move || {
92            init_thread_stream();
93            on_start();
94        })
95        .on_thread_stop(|| {
96            flush_thread_buffer();
97            unregister_thread_stream();
98        })
99    }
100}
micromegas_tracing/runtime.rs

micromegas_tracing/
runtime.rs
