Merge pull request #82 from mladedav/dm/recursive

feat: detect recursive calls that can deadlock `HierarchicalLayer`

src/lib.rs

@@ -10,7 +10,11 @@ use std::{
     io::{self, IsTerminal},
     iter::Fuse,
     mem,
-    sync::Mutex,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Mutex,
+    },
+    thread::LocalKey,
     time::Instant,
 };
 use tracing_core::{
@@ -463,6 +467,28 @@ where
             unit = self.styled(Style::new().dimmed(), unit),
         )
     }
+
+    fn is_recursive() -> Option<RecursiveGuard> {
+        thread_local! {
+            pub static IS_EMPTY: AtomicBool = const { AtomicBool::new(true) };
+        }
+
+        IS_EMPTY.with(|is_empty| {
+            is_empty
+                .compare_exchange(true, false, Ordering::Relaxed, Ordering::Relaxed)
+                .ok()
+                .map(|_| RecursiveGuard(&IS_EMPTY))
+        })
+    }
+}
+
+struct RecursiveGuard(&'static LocalKey<AtomicBool>);
+
+impl Drop for RecursiveGuard {
+    fn drop(&mut self) {
+        self.0
+            .with(|is_empty| is_empty.store(true, Ordering::Relaxed));
+    }
 }
 
 impl<S, W, FT> Layer<S> for HierarchicalLayer<W, FT>
@@ -472,6 +498,10 @@ where
     FT: FormatTime + 'static,
 {
     fn on_new_span(&self, attrs: &Attributes, id: &Id, ctx: Context<S>) {
+        let Some(_guard) = Self::is_recursive() else {
+            return;
+        };
+
         let span = ctx.span(id).expect("in new_span but span does not exist");
 
         if span.extensions().get::<Data>().is_none() {
@@ -507,6 +537,10 @@ where
     }
 
     fn on_event(&self, event: &Event<'_>, ctx: Context<S>) {
+        let Some(_guard) = Self::is_recursive() else {
+            return;
+        };
+
         let span = ctx.current_span();
         let span_id = span.id();
         let span = span_id.and_then(|id| ctx.span(id));
@@ -588,6 +622,10 @@ where
     }
 
     fn on_close(&self, id: Id, ctx: Context<S>) {
+        let Some(_guard) = Self::is_recursive() else {
+            return;
+        };
+
         let bufs = &mut *self.bufs.lock().unwrap();
 
         let span = ctx.span(&id).expect("invalid span in on_close");

tests/recursive_event.rs

@@ -0,0 +1,47 @@
+use std::{io, str, sync::Mutex};
+
+use tracing::subscriber::set_global_default;
+use tracing_subscriber::{layer::SubscriberExt, registry};
+
+use tracing_tree::HierarchicalLayer;
+
+struct RecursiveWriter(Mutex<Vec<u8>>);
+
+impl io::Write for &RecursiveWriter {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.0.lock().unwrap().extend(buf);
+
+        tracing::error!("Nobody expects the Spanish Inquisition");
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        tracing::error!("Nobody expects the Spanish Inquisition");
+        Ok(())
+    }
+}
+
+/// This test checks that if `tracing` events happen during processing of
+/// `on_event`, the library does not deadlock.
+#[test]
+fn recursive_event() {
+    static WRITER: RecursiveWriter = RecursiveWriter(Mutex::new(Vec::new()));
+
+    let subscriber = registry().with(HierarchicalLayer::new(2).with_writer(|| &WRITER));
+    // This has to be its own integration test because we can't just set a
+    // global default like this otherwise and not expect everything else to
+    // break.
+    set_global_default(subscriber).unwrap();
+
+    tracing::error!("We can never expect the unexpected.");
+
+    let output = WRITER.0.lock().unwrap();
+    let output = str::from_utf8(&output).unwrap();
+
+    // If this test finished we're happy. Let's just also check that we did
+    // in fact log _something_ and that the logs from within the writer did
+    // not actually go through.
+    assert!(output.contains("We can never expect the unexpected."));
+    assert!(!output.contains("Nobody expects the Spanish Inquisition"));
+}