Use the thread_local crate allowing multiple registries

tracing-subscriber/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.2.10"
 authors = [
     "Eliza Weisman <[email protected]>",
     "David Barsky <[email protected]>",
-    "Tokio Contributors <[email protected]>"
+    "Tokio Contributors <[email protected]>",
 ]
 edition = "2018"
 license = "MIT"
@@ -27,7 +27,7 @@ default = ["env-filter", "smallvec", "fmt", "ansi", "chrono", "tracing-log", "js
 env-filter = ["matchers", "regex", "lazy_static"]
 fmt = ["registry"]
 ansi = ["fmt", "ansi_term"]
-registry = ["sharded-slab"]
+registry = ["sharded-slab", "thread_local"]
 json = ["tracing-serde", "serde", "serde_json"]
 
 [dependencies]
@@ -54,13 +54,16 @@ parking_lot = { version = ">= 0.7, <= 0.11", optional = true }
 
 # registry
 sharded-slab = { version = "^0.0.9", optional = true }
+thread_local = { version = "1.0.1", optional = true }
 
 [dev-dependencies]
-tracing = { path = "../tracing", version = "0.1"}
+tracing = { path = "../tracing", version = "0.1" }
 log = "0.4"
 tracing-log = { path = "../tracing-log", version = "0.1" }
 criterion = { version = "0.3", default_features = false }
 regex = { version = "1", default-features = false, features = ["std"] }
+tracing-futures = { path = "../tracing-futures", version = "0.2", default-features = false, features = ["std-future", "std"] }
+tokio = { version = "0.2", features = ["rt-core", "macros"] }
 
 [badges]
 maintenance = { status = "experimental" }

tracing-subscriber/src/registry/sharded.rs

@@ -1,4 +1,5 @@
 use sharded_slab::{Guard, Slab};
+use thread_local::ThreadLocal;
 
 use super::stack::SpanStack;
 use crate::{
@@ -47,6 +48,7 @@ use tracing_core::{
 #[derive(Debug)]
 pub struct Registry {
     spans: Slab<DataInner>,
+    current_spans: ThreadLocal<RefCell<SpanStack>>,
 }
 
 /// Span data stored in a [`Registry`].
@@ -78,7 +80,10 @@ struct DataInner {
 
 impl Default for Registry {
     fn default() -> Self {
-        Self { spans: Slab::new() }
+        Self {
+            spans: Slab::new(),
+            current_spans: ThreadLocal::new(),
+        }
     }
 }
 
@@ -154,7 +159,6 @@ thread_local! {
     ///
     /// [`CloseGuard`]: ./struct.CloseGuard.html
     static CLOSE_COUNT: Cell<usize> = Cell::new(0);
-    static CURRENT_SPANS: RefCell<SpanStack> = RefCell::new(SpanStack::new());
 }
 
 impl Subscriber for Registry {
@@ -198,13 +202,18 @@ impl Subscriber for Registry {
     fn event(&self, _: &Event<'_>) {}
 
     fn enter(&self, id: &span::Id) {
-        CURRENT_SPANS.with(|spans| {
-            spans.borrow_mut().push(self.clone_span(id));
-        })
+        self.current_spans
+            .get_or_default()
+            .borrow_mut()
+            .push(self.clone_span(id));
     }
 
     fn exit(&self, id: &span::Id) {
-        if let Some(id) = CURRENT_SPANS.with(|spans| spans.borrow_mut().pop(id)) {
+        if let Some(id) = self
+            .current_spans
+            .get()
+            .and_then(|spans| spans.borrow_mut().pop(id))
+        {
             dispatcher::get_default(|dispatch| dispatch.try_close(id.clone()));
         }
     }
@@ -224,8 +233,9 @@ impl Subscriber for Registry {
     }
 
     fn current_span(&self) -> Current {
-        CURRENT_SPANS
-            .with(|spans| {
+        self.current_spans
+            .get()
+            .and_then(|spans| {
                 let spans = spans.borrow();
                 let id = spans.current()?;
                 let span = self.get(id)?;

tracing-subscriber/src/registry/stack.rs

@@ -1,8 +1,8 @@
-use std::cell::RefCell;
 use std::collections::HashSet;
 
 pub(crate) use tracing_core::span::Id;
 
+#[derive(Debug)]
 struct ContextId {
     id: Id,
     duplicate: bool,
@@ -12,19 +12,13 @@ struct ContextId {
 ///
 /// A "separate current span" for each thread is a semantic choice, as each span
 /// can be executing in a different thread.
+#[derive(Debug, Default)]
 pub(crate) struct SpanStack {
     stack: Vec<ContextId>,
     ids: HashSet<Id>,
 }
 
 impl SpanStack {
-    pub(crate) fn new() -> Self {
-        SpanStack {
-            stack: vec![],
-            ids: HashSet::new(),
-        }
-    }
-
     pub(crate) fn push(&mut self, id: Id) {
         let duplicate = self.ids.contains(&id);
         if !duplicate {
@@ -61,17 +55,13 @@ impl SpanStack {
     }
 }
 
-thread_local! {
-    static CONTEXT: RefCell<SpanStack> = RefCell::new(SpanStack::new());
-}
-
 #[cfg(test)]
 mod tests {
     use super::{Id, SpanStack};
 
     #[test]
     fn pop_last_span() {
-        let mut stack = SpanStack::new();
+        let mut stack = SpanStack::default();
         let id = Id::from_u64(1);
         stack.push(id.clone());
 
@@ -80,7 +70,7 @@ mod tests {
 
     #[test]
     fn pop_first_span() {
-        let mut stack = SpanStack::new();
+        let mut stack = SpanStack::default();
         stack.push(Id::from_u64(1));
         stack.push(Id::from_u64(2));
 

tracing-subscriber/tests/registry_with_subscriber.rs

@@ -0,0 +1,24 @@
+use tracing_futures::{Instrument, WithSubscriber};
+use tracing_subscriber::prelude::*;
+
+#[tokio::test]
+async fn future_with_subscriber() {
+    let _default = tracing_subscriber::registry().init();
+    let span = tracing::info_span!("foo");
+    let _e = span.enter();
+    let span = tracing::info_span!("bar");
+    let _e = span.enter();
+    tokio::spawn(
+        async {
+            async {
+                let span = tracing::Span::current();
+                println!("{:?}", span);
+            }
+            .instrument(tracing::info_span!("hi"))
+            .await
+        }
+        .with_subscriber(tracing_subscriber::registry()),
+    )
+    .await
+    .unwrap();
+}