From acd08eb23d48287c6f5e5053f940204462ad7c1d Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Wed, 13 Mar 2019 10:28:45 -0700
Subject: [PATCH] tokio: Enable trace subscriber propagation in the runtime
 (#966)

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 Cargo.toml                        |  2 ++
 src/lib.rs                        |  1 +
 src/runtime/threadpool/builder.rs | 22 ++++++++++++++++------
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index a9a3e42f3..9ac5ecb26 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -70,6 +70,7 @@ rt-full = [
   "tokio-current-thread",
   "tokio-executor",
   "tokio-threadpool",
+  "tokio-trace-core",
 ]
 sync = ["tokio-sync"]
 tcp = ["tokio-tcp"]
@@ -105,6 +106,7 @@ tokio-threadpool = { version = "0.1.8", path = "tokio-threadpool", optional = tr
 tokio-tcp = { version = "0.1.0", path = "tokio-tcp", optional = true }
 tokio-udp = { version = "0.1.0", path = "tokio-udp", optional = true }
 tokio-timer = { version = "0.2.8", path = "tokio-timer", optional = true }
+tokio-trace-core = { version = "0.1", path = "tokio-trace/tokio-trace-core", optional = true }
 
 # Needed until `reactor` is removed from `tokio`.
 mio = { version = "0.6.14", optional = true }
diff --git a/src/lib.rs b/src/lib.rs
index b4f49c268..3692f2199 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -135,6 +135,7 @@ pub mod util;
 
 if_runtime! {
     extern crate tokio_executor;
+    extern crate tokio_trace_core;
     pub mod executor;
     pub mod runtime;
 
diff --git a/src/runtime/threadpool/builder.rs b/src/runtime/threadpool/builder.rs
index f7af4493e..d32251da0 100644
--- a/src/runtime/threadpool/builder.rs
+++ b/src/runtime/threadpool/builder.rs
@@ -11,6 +11,7 @@ use tokio_reactor;
 use tokio_threadpool::Builder as ThreadPoolBuilder;
 use tokio_timer::clock::{self, Clock};
 use tokio_timer::timer::{self, Timer};
+use tokio_trace_core as trace;
 
 /// Builds Tokio Runtime with custom configuration values.
 ///
@@ -90,10 +91,10 @@ impl Builder {
 
     /// Set builder to set up the thread pool instance.
     #[deprecated(
-        since="0.1.9",
-        note="use the `core_threads`, `blocking_threads`, `name_prefix`, \
-              `keep_alive`, and `stack_size` functions on `runtime::Builder`, \
-              instead")]
+        since = "0.1.9",
+        note = "use the `core_threads`, `blocking_threads`, `name_prefix`, \
+                `keep_alive`, and `stack_size` functions on `runtime::Builder`, \
+                instead")]
     #[doc(hidden)]
     pub fn threadpool_builder(&mut self, val: ThreadPoolBuilder) -> &mut Self {
         self.threadpool_builder = val;
@@ -330,14 +331,23 @@ impl Builder {
         // Get a handle to the clock for the runtime.
         let clock = self.clock.clone();
 
-        let pool = self.threadpool_builder
+        // Get the current trace dispatcher.
+        // TODO(eliza): when `tokio-trace-core` is stable enough to take a
+        // public API dependency, we should allow users to set a custom
+        // subscriber for the runtime.
+        let dispatch = trace::dispatcher::get_default(trace::Dispatch::clone);
+
+        let pool = self
+            .threadpool_builder
             .around_worker(move |w, enter| {
                 let index = w.id().to_usize();
 
                 tokio_reactor::with_default(&reactor_handles[index], enter, |enter| {
                     clock::with_default(&clock, enter, |enter| {
                         timer::with_default(&timer_handles[index], enter, |_| {
-                            w.run();
+                            trace::dispatcher::with_default(&dispatch, || {
+                                w.run();
+                            })
                         });
                     })
                 });