Switch to tokio channel

Cargo.lock

@@ -5716,6 +5716,7 @@ dependencies = [
  "serde",
  "serde_json",
  "stats_alloc",
+ "tokio",
  "tracing",
  "tracing-error",
  "tracing-subscriber",

tracing-trace/Cargo.toml

@@ -18,4 +18,5 @@ byte-unit = { version = "4.0.19", default-features = false, features = [
     "std",
     "serde",
 ] }
+tokio = { version = "1.35.1", features = ["sync"] }
 

tracing-trace/src/layer.rs

@@ -19,7 +19,7 @@ use crate::{Error, Trace};
 
 /// Layer that measures the time spent in spans.
 pub struct TraceLayer<A: GlobalAlloc + 'static = System> {
-    sender: std::sync::mpsc::Sender<Entry>,
+    sender: tokio::sync::mpsc::UnboundedSender<Entry>,
     callsites: RwLock<HashMap<OpaqueIdentifier, ResourceId>>,
     start_time: std::time::Instant,
     memory_allocator: Option<&'static StatsAlloc<A>>,
@@ -27,7 +27,7 @@ pub struct TraceLayer<A: GlobalAlloc + 'static = System> {
 
 impl<W: Write> Trace<W> {
     pub fn new(writer: W) -> (Self, TraceLayer<System>) {
-        let (sender, receiver) = std::sync::mpsc::channel();
+        let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
         let trace = Trace { writer, receiver };
         let layer = TraceLayer {
             sender,
@@ -42,7 +42,7 @@ impl<W: Write> Trace<W> {
         writer: W,
         stats_alloc: &'static StatsAlloc<A>,
     ) -> (Self, TraceLayer<A>) {
-        let (sender, receiver) = std::sync::mpsc::channel();
+        let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
         let trace = Trace { writer, receiver };
         let layer = TraceLayer {
             sender,
@@ -53,8 +53,17 @@ impl<W: Write> Trace<W> {
         (trace, layer)
     }
 
-    pub fn receive(&mut self) -> Result<ControlFlow<(), ()>, Error> {
+    pub async fn receive(&mut self) -> Result<ControlFlow<(), ()>, Error> {
-        let Ok(entry) = self.receiver.recv() else {
+        let Some(entry) = self.receiver.recv().await else {
+            return Ok(ControlFlow::Break(()));
+        };
+        self.write(entry)?;
+        Ok(ControlFlow::Continue(()))
+    }
+
+    /// Panics if called from an asynchronous context
+    pub fn blocking_receive(&mut self) -> Result<ControlFlow<(), ()>, Error> {
+        let Some(entry) = self.receiver.blocking_recv() else {
             return Ok(ControlFlow::Break(()));
         };
         self.write(entry)?;

tracing-trace/src/lib.rs

@@ -11,7 +11,7 @@ pub use error::Error;
 
 pub struct Trace<W: Write> {
     writer: W,
-    receiver: std::sync::mpsc::Receiver<Entry>,
+    receiver: tokio::sync::mpsc::UnboundedReceiver<Entry>,
 }
 
 pub struct TraceReader<R: Read> {