From 59494e0897b2acc034d4344f3b2e1c6a57db371c Mon Sep 17 00:00:00 2001 From: Jerome Gravel-Niquet Date: Mon, 10 Aug 2020 19:31:09 -0400 Subject: [PATCH] subscriber: use the thread_local crate to allow multiple registries (#901) ## Motivation Fixes part of #898 where creating multiple registries made them fight over the thread_local-stored span stacks. ## Solution Use the `thread_local` crate which has a tighter scope. I believe @hawkw has a test that fails without this PR. --- tracing-subscriber/Cargo.toml | 9 ++++--- tracing-subscriber/src/registry/sharded.rs | 26 +++++++++++++------ tracing-subscriber/src/registry/stack.rs | 18 +++---------- .../tests/registry_with_subscriber.rs | 24 +++++++++++++++++ 4 files changed, 52 insertions(+), 25 deletions(-) create mode 100644 tracing-subscriber/tests/registry_with_subscriber.rs diff --git a/tracing-subscriber/Cargo.toml b/tracing-subscriber/Cargo.toml index b70e1bbac4..be89d05086 100644 --- a/tracing-subscriber/Cargo.toml +++ b/tracing-subscriber/Cargo.toml @@ -4,7 +4,7 @@ version = "0.2.10" authors = [ "Eliza Weisman ", "David Barsky ", - "Tokio Contributors " + "Tokio Contributors ", ] 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" } diff --git a/tracing-subscriber/src/registry/sharded.rs b/tracing-subscriber/src/registry/sharded.rs index 40075c6e93..c55b8081ce 100644 --- a/tracing-subscriber/src/registry/sharded.rs +++ b/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, + current_spans: ThreadLocal>, } /// 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 = Cell::new(0); - static CURRENT_SPANS: RefCell = 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)?; diff --git a/tracing-subscriber/src/registry/stack.rs b/tracing-subscriber/src/registry/stack.rs index 50fdc5ef69..e85cb7c3cf 100644 --- a/tracing-subscriber/src/registry/stack.rs +++ b/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, ids: HashSet, } 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 = 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)); diff --git a/tracing-subscriber/tests/registry_with_subscriber.rs b/tracing-subscriber/tests/registry_with_subscriber.rs new file mode 100644 index 0000000000..e0d6544e0a --- /dev/null +++ b/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(); +}