From 0e7a411d4d8e081c6e9b6e4e3e4b28f1f2f811fb Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Wed, 31 Jan 2024 17:44:54 +0100
Subject: [PATCH] tracing-trace: introduce TraceWriter, trace now only exposes
 the channel

---
 tracing-trace/src/layer.rs | 26 ++++++++++++++++++++------
 tracing-trace/src/lib.rs   | 16 +++++++++++++++-
 tracing-trace/src/main.rs  |  2 +-
 3 files changed, 36 insertions(+), 8 deletions(-)

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<A: GlobalAlloc + 'static = System> {
@@ -25,10 +25,10 @@ pub struct TraceLayer<A: GlobalAlloc + 'static = System> {
     memory_allocator: Option<&'static StatsAlloc<A>>,
 }
 
-impl<W: Write> Trace<W> {
-    pub fn new(writer: W) -> (Self, TraceLayer<System>) {
+impl Trace {
+    pub fn new() -> (Self, TraceLayer<System>) {
         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<W: Write> Trace<W> {
     }
 
     pub fn with_stats_alloc<A: GlobalAlloc>(
-        writer: W,
         stats_alloc: &'static StatsAlloc<A>,
     ) -> (Self, TraceLayer<A>) {
         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<W: Write> Trace<W> {
         };
         (trace, layer)
     }
+}
+
+impl<W: Write> TraceWriter<W> {
+    pub fn new(writer: W) -> (Self, TraceLayer<System>) {
+        let (trace, layer) = Trace::new();
+        (trace.into_writer(writer), layer)
+    }
+
+    pub fn with_stats_alloc<A: GlobalAlloc>(
+        writer: W,
+        stats_alloc: &'static StatsAlloc<A>,
+    ) -> (Self, TraceLayer<A>) {
+        let (trace, layer) = Trace::with_stats_alloc(stats_alloc);
+        (trace.into_writer(writer), layer)
+    }
 
     pub async fn receive(&mut self) -> Result<ControlFlow<(), ()>, 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<W: Write> {
+pub struct TraceWriter<W: Write> {
     writer: W,
     receiver: tokio::sync::mpsc::UnboundedReceiver<Entry>,
 }
 
+pub struct Trace {
+    receiver: tokio::sync::mpsc::UnboundedReceiver<Entry>,
+}
+
+impl Trace {
+    pub fn into_receiver(self) -> tokio::sync::mpsc::UnboundedReceiver<Entry> {
+        self.receiver
+    }
+
+    pub fn into_writer<W: Write>(self, writer: W) -> TraceWriter<W> {
+        TraceWriter { writer, receiver: self.receiver }
+    }
+}
+
 pub struct TraceReader<R: Read> {
     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