diff --git a/tracing-trace/src/layer.rs b/tracing-trace/src/layer.rs index 1b9aadfa7..aa2908304 100644 --- a/tracing-trace/src/layer.rs +++ b/tracing-trace/src/layer.rs @@ -15,7 +15,7 @@ use crate::entry::{ Entry, Event, MemoryStats, NewCallsite, NewSpan, NewThread, ResourceId, SpanClose, SpanEnter, SpanExit, SpanId, }; -use crate::{Error, Trace}; +use crate::{Error, Trace, TraceWriter}; /// Layer that measures the time spent in spans. pub struct TraceLayer { @@ -25,10 +25,10 @@ pub struct TraceLayer { memory_allocator: Option<&'static StatsAlloc>, } -impl Trace { - pub fn new(writer: W) -> (Self, TraceLayer) { +impl Trace { + pub fn new() -> (Self, TraceLayer) { let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - let trace = Trace { writer, receiver }; + let trace = Trace { receiver }; let layer = TraceLayer { sender, callsites: Default::default(), @@ -39,11 +39,10 @@ impl Trace { } pub fn with_stats_alloc( - writer: W, stats_alloc: &'static StatsAlloc, ) -> (Self, TraceLayer) { let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - let trace = Trace { writer, receiver }; + let trace = Trace { receiver }; let layer = TraceLayer { sender, callsites: Default::default(), @@ -52,6 +51,21 @@ impl Trace { }; (trace, layer) } +} + +impl TraceWriter { + pub fn new(writer: W) -> (Self, TraceLayer) { + let (trace, layer) = Trace::new(); + (trace.into_writer(writer), layer) + } + + pub fn with_stats_alloc( + writer: W, + stats_alloc: &'static StatsAlloc, + ) -> (Self, TraceLayer) { + let (trace, layer) = Trace::with_stats_alloc(stats_alloc); + (trace.into_writer(writer), layer) + } pub async fn receive(&mut self) -> Result, Error> { let Some(entry) = self.receiver.recv().await else { diff --git a/tracing-trace/src/lib.rs b/tracing-trace/src/lib.rs index 3d00eef10..77eb61d47 100644 --- a/tracing-trace/src/lib.rs +++ b/tracing-trace/src/lib.rs @@ -9,11 +9,25 @@ pub mod processor; pub use error::Error; -pub struct Trace { +pub struct TraceWriter { writer: W, receiver: tokio::sync::mpsc::UnboundedReceiver, } +pub struct Trace { + receiver: tokio::sync::mpsc::UnboundedReceiver, +} + +impl Trace { + pub fn into_receiver(self) -> tokio::sync::mpsc::UnboundedReceiver { + self.receiver + } + + pub fn into_writer(self, writer: W) -> TraceWriter { + TraceWriter { writer, receiver: self.receiver } + } +} + pub struct TraceReader { reader: R, } diff --git a/tracing-trace/src/main.rs b/tracing-trace/src/main.rs index f9f665861..c74bf1fb7 100644 --- a/tracing-trace/src/main.rs +++ b/tracing-trace/src/main.rs @@ -74,7 +74,7 @@ fn on_panic(info: &std::panic::PanicInfo) { fn main() { let (mut trace, profiling_layer) = - tracing_trace::Trace::new(std::fs::File::create("trace.json").unwrap()); + tracing_trace::TraceWriter::new(std::fs::File::create("trace.json").unwrap()); let subscriber = tracing_subscriber::registry() // any number of other subscriber layers may be added before or