trace: add program-wide default dispatcher (#1152)

## Motivation

I was just trying to use tokio-trace for a greenfield project, but I was frustrated to discover that I couldn't really use it easily.

I was using the [`runtime`](https://docs.rs/runtime/0.3.0-alpha.4/runtime/) crate, which transparently spawns a thread pool executor for futures. In that thread pool, there's no way to set a tokio-trace subscriber for the duration of each thread, since you don't control the thread initialization. You *might* be able to wrap every future you spawn with a subscriber call, but that's a lot of work.

I was also confused because the documentation said that setting a subscriber in the main thread would use that subscriber for the rest of the program. That isn't the case, though -- the subscriber will be used only on the main thread, and not on worker threads, etc.

## Solution

I added a function `set_global_default`, which works similarly to the `log` crate:

```rust
tokio_trace::subscriber::set_global_default(FooSubscriber::new());
```

The global subscriber (actually a global `Dispatch`) is a `static mut` protected by an atomic; implementation is copied from the `log` crate. It is used as a fallback if a thread has no `Dispatch` currently set. This is extremely simple to use, and doesn't break any existing functionality.

Performance-wise, thread-local `Dispatch` lookup goes from ~4.5ns to ~5ns, according to the benchmarks. So, barely any runtime overhead. (Presumably there's a little compile-time overhead but idk how to measure that.) Since the atomic guard is only ever written once, it will be shared among a CPU's cores and read very cheaply.

I added some docs to partially address #1151. I also switched the tokio-trace benchmarks to criterion because the nightly benchmarks weren't compiling (missing `dyn` flags?)

tokio-trace/Cargo.toml

@@ -11,6 +11,7 @@ name = "tokio-trace"
 version = "0.1.0"
 authors = ["Tokio Contributors <team@tokio.rs>"]
 license = "MIT"
+readme = "README.md"
 repository = "https://github.com/tokio-rs/tokio"
 homepage = "https://tokio.rs"
 documentation = "https://docs.rs/tokio-trace/0.1.0/tokio_trace"
@@ -21,7 +22,7 @@ categories = ["development-tools::debugging", "asynchronous"]
 keywords = ["logging", "tracing"]
 
 [dependencies]
-tokio-trace-core = "0.2"
+tokio-trace-core = { version = "0.2", path = "tokio-trace-core" }
 log = { version = "0.4", optional = true }
 cfg-if = "0.1.7"
 

tokio-trace/README.md

@@ -128,9 +128,28 @@ In order to record trace events, executables have to use a `Subscriber`
 implementation compatible with `tokio-trace`. A `Subscriber` implements a way of
 collecting trace data, such as by logging it to standard output.
 
-Unlike the `log` crate, `tokio-trace` does *not* use a global `Subscriber` which
+There currently aren't too many subscribers to choose from. The best one to use right now
-is initialized once. Instead, it follows the `tokio` pattern of executing code
+is probably [`tokio-trace-fmt`], which logs to the terminal.
-in a context. For example:
+
+The simplest way to use a subscriber is to call the `set_global_default` function:
+
+```rust
+#[macro_use]
+extern crate tokio_trace;
+
+let my_subscriber = FooSubscriber::new();
+
+tokio_trace::subscriber::set_global_default(my_subscriber).expect("setting tokio_trace default failed");
+```
+
+This subscriber will be used as the default in all threads for the remainder of the duration
+of the program, similar to how loggers work in the `log` crate.
+
+Note: Libraries should *NOT* call `set_global_default()`! That will cause conflicts when
+executables try to set the default later.
+
+In addition, you can locally override the default subscriber, using the `tokio` pattern
+of executing code in a context. For example:
 
 ```rust
 #[macro_use]
@@ -145,9 +164,10 @@ tokio_trace::subscriber::with_default(subscriber, || {
 ```
 
 This approach allows trace data to be collected by multiple subscribers within
-different contexts in the program. Alternatively, a single subscriber may be
+different contexts in the program. Note that the override only applies to the
-constructed by the `main` function and all subsequent code executed with that
+currently executing thread; other threads will not see the change from with_default.
-subscriber as the default. Any trace events generated outside the context of a
+
+Any trace events generated outside the context of a
 subscriber will not be collected.
 
 The executable itself may use the `tokio-trace` crate to instrument itself as
@@ -159,6 +179,7 @@ be used with the `tokio-trace` ecosystem. It includes a collection of
 
 [`log`]: https://docs.rs/log/0.4.6/log/
 [`tokio-trace-nursery`]: https://github.com/tokio-rs/tokio-trace-nursery
+[`tokio-trace-fmt`]: https://github.com/tokio-rs/tokio-trace-nursery/tree/master/tokio-trace-fmt
 
 ## License
 

tokio-trace/src/lib.rs

@@ -366,6 +366,26 @@
 //! implementation compatible with `tokio-trace`. A `Subscriber` implements a
 //! way of collecting trace data, such as by logging it to standard output.
 //!
+//! There currently aren't too many subscribers to choose from. The best one to use right now
+//!  is probably [`tokio-trace-fmt`], which logs to the terminal.
+//! The simplest way to use a subscriber is to call the `set_global_default` function:
+//!
+//! ```no_build
+//! #[macro_use]
+//! extern crate tokio_trace;
+//! let my_subscriber = FooSubscriber::new();
+//! tokio_trace::subscriber::set_global_default(my_subscriber).expect("setting tokio_trace default failed");
+//! ```
+//!
+//! Note: Libraries should *NOT* call `set_global_default()`! That will cause conflicts when
+//! executables try to set the default later.
+//!
+//! This subscriber will be used as the default in all threads for the remainder of the duration
+//! of the program, similar to how loggers work in the `log` crate.
+//!
+//! In addition, you can locally override the default subscriber, using the `tokio` pattern
+//! of executing code in a context. For example:
+//!
 //! Unlike the `log` crate, `tokio-trace` does *not* use a global `Subscriber`
 //! which is initialized once. Instead, it follows the `tokio` pattern of
 //! executing code in a context. For example:
@@ -399,9 +419,11 @@
 //! ```
 //!
 //! This approach allows trace data to be collected by multiple subscribers
-//! within different contexts in the program. Alternatively, a single subscriber
+//! within different contexts in the program. Note that the override only applies to the
-//! may be constructed by the `main` function and all subsequent code executed
+//! currently executing thread; other threads will not see the change from with_default.
-//! with that subscriber as the default. Any trace events generated outside the
+//! with that subscriber as the default.
+//!
+//! Any trace events generated outside the
 //! context of a subscriber will not be collected.
 //!
 //! The executable itself may use the `tokio-trace` crate to instrument itself

tokio-trace/src/subscriber.rs

@@ -1,9 +1,9 @@
 //! Collects and records trace data.
 pub use tokio_trace_core::subscriber::*;
 
-/// Sets this dispatch as the default for the duration of a closure.
+/// Sets this subscriber as the default for the duration of a closure.
 ///
-/// The default dispatcher is used when creating a new [`Span`] or
+/// The default subscriber is used when creating a new [`Span`] or
 /// [`Event`], _if no span is currently executing_. If a span is currently
 /// executing, new spans or events are dispatched to the subscriber that
 /// tagged that span, instead.
@@ -17,3 +17,24 @@ where
 {
     ::dispatcher::with_default(&::Dispatch::new(subscriber), f)
 }
+
+/// Sets this subscriber as the global default for the duration of the entire program.
+/// Will be used as a fallback if no thread-local subscriber has been set in a thread (using `with_default`.)
+///
+/// Can only be set once; subsequent attempts to set the global default will fail.
+/// Returns whether the initialization was successful.
+///
+/// Note: Libraries should *NOT* call `set_global_default()`! That will cause conflicts when
+/// executables try to set them later.
+///
+/// [span]: ../span/index.html
+/// [`Subscriber`]: ../subscriber/trait.Subscriber.html
+/// [`Event`]: ../event/struct.Event.html
+pub fn set_global_default<S>(subscriber: S) -> Result<(), SetGlobalDefaultError>
+where
+    S: Subscriber + Send + Sync + 'static,
+{
+    ::dispatcher::set_global_default(::Dispatch::new(subscriber))
+}
+
+pub use tokio_trace_core::dispatcher::SetGlobalDefaultError;

tokio-trace/tokio-trace-core/src/dispatcher.rs

@@ -8,8 +8,11 @@ use {
 use std::{
     any::Any,
     cell::{Cell, RefCell},
-    fmt,
+    error, fmt,
-    sync::{Arc, Weak},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc, Weak,
+    },
 };
 
 /// `Dispatch` trace data to a [`Subscriber`].
@@ -27,6 +30,13 @@ thread_local! {
     };
 }
 
+static GLOBAL_INIT: AtomicUsize = AtomicUsize::new(UNINITIALIZED);
+const UNINITIALIZED: usize = 0;
+const INITIALIZING: usize = 1;
+const INITIALIZED: usize = 2;
+
+static mut GLOBAL_DISPATCH: Option<Dispatch> = None;
+
 /// The dispatch state of a thread.
 struct State {
     /// This thread's current default dispatcher.
@@ -63,6 +73,47 @@ pub fn with_default<T>(dispatcher: &Dispatch, f: impl FnOnce() -> T) -> T {
     let _guard = State::set_default(dispatcher.clone());
     f()
 }
+
+/// Sets this dispatch as the global default for the duration of the entire program.
+/// Will be used as a fallback if no thread-local dispatch has been set in a thread
+/// (using `with_default`.)
+///
+/// Can only be set once; subsequent attempts to set the global default will fail.
+/// Returns `Err` if the global default has already been set.
+///
+/// Note: Libraries should *NOT* call `set_global_default()`! That will cause conflicts when
+/// executables try to set them later.
+///
+/// [span]: ../span/index.html
+/// [`Subscriber`]: ../subscriber/trait.Subscriber.html
+/// [`Event`]: ../event/struct.Event.html
+pub fn set_global_default(dispatcher: Dispatch) -> Result<(), SetGlobalDefaultError> {
+    if GLOBAL_INIT.compare_and_swap(UNINITIALIZED, INITIALIZING, Ordering::SeqCst) == UNINITIALIZED
+    {
+        unsafe {
+            GLOBAL_DISPATCH = Some(dispatcher.clone());
+        }
+        GLOBAL_INIT.store(INITIALIZED, Ordering::SeqCst);
+        Ok(())
+    } else {
+        Err(SetGlobalDefaultError { _no_construct: () })
+    }
+}
+
+/// Returned if setting the global dispatcher fails.
+#[derive(Debug)]
+pub struct SetGlobalDefaultError {
+    _no_construct: (),
+}
+
+impl fmt::Display for SetGlobalDefaultError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.pad("a global default trace dispatcher has already been set")
+    }
+}
+
+impl error::Error for SetGlobalDefaultError {}
+
 /// Executes a closure with a reference to this thread's current [dispatcher].
 ///
 /// Note that calls to `get_default` should not be nested; if this function is
@@ -89,7 +140,20 @@ where
         .try_with(|state| {
             if state.can_enter.replace(false) {
                 let _guard = Entered(&state.can_enter);
-                f(&state.default.borrow())
+
+                let mut default = state.default.borrow_mut();
+
+                if default.is::<NoSubscriber>() && GLOBAL_INIT.load(Ordering::SeqCst) == INITIALIZED
+                {
+                    // don't redo this call on the next check
+                    unsafe {
+                        *default = GLOBAL_DISPATCH
+                            .as_ref()
+                            .expect("invariant violated: GLOBAL_DISPATCH must be initialized before GLOBAL_INIT is set")
+                            .clone()
+                    }
+                }
+                f(&*default)
             } else {
                 f(&Dispatch::none())
             }
@@ -484,4 +548,65 @@ mod test {
 
         with_default(&Dispatch::new(TestSubscriber), || mk_span())
     }
+
+    #[test]
+    fn global_dispatch() {
+        struct TestSubscriberA;
+        impl Subscriber for TestSubscriberA {
+            fn enabled(&self, _: &Metadata) -> bool {
+                true
+            }
+            fn new_span(&self, _: &span::Attributes) -> span::Id {
+                span::Id::from_u64(1)
+            }
+            fn record(&self, _: &span::Id, _: &span::Record) {}
+            fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
+            fn event(&self, _: &Event) {}
+            fn enter(&self, _: &span::Id) {}
+            fn exit(&self, _: &span::Id) {}
+        }
+        struct TestSubscriberB;
+        impl Subscriber for TestSubscriberB {
+            fn enabled(&self, _: &Metadata) -> bool {
+                true
+            }
+            fn new_span(&self, _: &span::Attributes) -> span::Id {
+                span::Id::from_u64(1)
+            }
+            fn record(&self, _: &span::Id, _: &span::Record) {}
+            fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
+            fn event(&self, _: &Event) {}
+            fn enter(&self, _: &span::Id) {}
+            fn exit(&self, _: &span::Id) {}
+        }
+
+        // NOTE: if you want to have other tests that set the default dispatch you'll need to
+        // write them as integration tests in ../tests/
+        set_global_default(Dispatch::new(TestSubscriberA)).expect("global dispatch set failed");
+        get_default(|current| {
+            assert!(
+                current.is::<TestSubscriberA>(),
+                "global dispatch get failed"
+            )
+        });
+
+        with_default(&Dispatch::new(TestSubscriberB), || {
+            get_default(|current| {
+                assert!(
+                    current.is::<TestSubscriberB>(),
+                    "thread-local override of global dispatch failed"
+                )
+            });
+        });
+
+        get_default(|current| {
+            assert!(
+                current.is::<TestSubscriberA>(),
+                "reset to global override failed"
+            )
+        });
+
+        set_global_default(Dispatch::new(TestSubscriberA))
+            .expect_err("double global dispatch set succeeded");
+    }
 }