Expose around worker in Runtime builder

tokio/src/runtime/threadpool/builder.rs

@@ -6,11 +6,12 @@ use tokio_timer::clock::{self, Clock};
 use tokio_timer::timer::{self, Timer};
 
 use num_cpus;
-use tracing_core as trace;
+use std::any::Any;
+use std::fmt;
 use std::io;
-use std::sync::Mutex;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
-use std::any::Any;
+use tracing_core as trace;
 
 /// Builds Tokio Runtime with custom configuration values.
 ///
@@ -48,7 +49,6 @@ use std::any::Any;
 ///     // use runtime ...
 /// }
 /// ```
-#[derive(Debug)]
 pub struct Builder {
     /// Thread pool specific builder
     threadpool_builder: threadpool::Builder,
@@ -58,6 +58,8 @@ pub struct Builder {
 
     /// The clock to use
     clock: Clock,
+
+    around_worker: Option<Arc<dyn Fn(&threadpool::Worker) + Send + Sync + 'static>>,
 }
 
 impl Builder {
@@ -76,6 +78,7 @@ impl Builder {
             threadpool_builder,
             core_threads,
             clock: Clock::new(),
+            around_worker: None,
         }
     }
 
@@ -113,7 +116,6 @@ impl Builder {
         self
     }
 
-
     /// Set the maximum number of worker threads for the `Runtime`'s thread pool.
     ///
     /// This must be a number between 1 and 32,768 though it is advised to keep
@@ -244,6 +246,35 @@ impl Builder {
         self
     }
 
+    /// Execute function `f` on each worker thread.
+    ///
+    /// This function is provided a handle to the worker and is expected to call
+    /// [`Worker::run`], otherwise the worker thread will shutdown without doing
+    /// any work.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::runtime::Builder;
+    ///
+    /// let rt = Builder::new()
+    ///     .around_worker(|worker| {
+    ///         println!("worker is starting up");
+    ///         worker.run();
+    ///         println!("worker is shutting down");
+    ///     })
+    ///     .build();
+    /// ```
+    ///
+    /// [`Worker::run`]: struct.Worker.html#method.run
+    pub fn around_worker<F>(&mut self, f: F) -> &mut Self
+    where
+        F: Fn(&threadpool::Worker) + Send + Sync + 'static,
+    {
+        self.threadpool_builder.around_worker(f);
+        self
+    }
+
     /// Execute function `f` after each thread is started but before it starts
     /// doing work.
     ///
@@ -263,7 +294,8 @@ impl Builder {
     /// # }
     /// ```
     pub fn after_start<F>(&mut self, f: F) -> &mut Self
-        where F: Fn() + Send + Sync + 'static
+    where
+        F: Fn() + Send + Sync + 'static,
     {
         self.threadpool_builder.after_start(f);
         self
@@ -287,7 +319,8 @@ impl Builder {
     /// # }
     /// ```
     pub fn before_stop<F>(&mut self, f: F) -> &mut Self
-        where F: Fn() + Send + Sync + 'static
+    where
+        F: Fn() + Send + Sync + 'static,
     {
         self.threadpool_builder.before_stop(f);
         self
@@ -337,6 +370,7 @@ impl Builder {
         let trace = dispatch.clone();
 
         let background = background::spawn(&clock)?;
+        let around_worker = self.around_worker.clone();
 
         let pool = self
             .threadpool_builder
@@ -347,18 +381,18 @@ impl Builder {
                 clock::with_default(&clock, || {
                     let _timer = timer::set_default(&timer_handles[index]);
                     trace::dispatcher::with_default(&dispatch, || {
-                        w.run();
+                        if let Some(around) = around_worker.clone() {
+                            around(w)
+                        } else {
+                            w.run();
+                        }
                     })
                 })
             })
             .custom_park(move |worker_id| {
                 let index = worker_id.to_usize();
 
-                timers[index]
-                    .lock()
-                    .unwrap()
-                    .take()
-                    .unwrap()
+                timers[index].lock().unwrap().take().unwrap()
             })
             .build();
 
@@ -377,3 +411,13 @@ impl Default for Builder {
         Self::new()
     }
 }
+
+impl fmt::Debug for Builder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Builder")
+            .field("threadpool_builder", &self.threadpool_builder)
+            .field("core_threads", &self.core_threads)
+            .field("clock", &self.clock)
+            .finish()
+    }
+}