Expose after_start and before_stop in runtime::Builder (#756)

Closes #705

src/runtime/builder.rs

@@ -240,6 +240,59 @@ impl Builder {
         self
     }
 
+    /// Execute function `f` after each thread is started but before it starts
+    /// doing work.
+    ///
+    /// This is intended for bookkeeping and monitoring use cases.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # extern crate tokio;
+    /// # extern crate futures;
+    /// # use tokio::runtime;
+    ///
+    /// # pub fn main() {
+    /// let thread_pool = runtime::Builder::new()
+    ///     .after_start(|| {
+    ///         println!("thread started");
+    ///     })
+    ///     .build();
+    /// # }
+    /// ```
+    pub fn after_start<F>(&mut self, f: F) -> &mut Self
+        where F: Fn() + Send + Sync + 'static
+    {
+        self.threadpool_builder.after_start(f);
+        self
+    }
+
+    /// Execute function `f` before each thread stops.
+    ///
+    /// This is intended for bookkeeping and monitoring use cases.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # extern crate tokio;
+    /// # extern crate futures;
+    /// # use tokio::runtime;
+    ///
+    /// # pub fn main() {
+    /// let thread_pool = runtime::Builder::new()
+    ///     .before_stop(|| {
+    ///         println!("thread stopping");
+    ///     })
+    ///     .build();
+    /// # }
+    /// ```
+    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
+        where F: Fn() + Send + Sync + 'static
+    {
+        self.threadpool_builder.before_stop(f);
+        self
+    }
+
     /// Create the configured `Runtime`.
     ///
     /// The returned `ThreadPool` instance is ready to spawn tasks.

tests/runtime.rs

@@ -3,7 +3,7 @@ extern crate env_logger;
 extern crate futures;
 
 use futures::sync::oneshot;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, atomic};
 use std::thread;
 use tokio::io;
 use tokio::net::{TcpStream, TcpListener};
@@ -493,3 +493,24 @@ fn runtime_reactor_handle() {
 
     th.join().unwrap();
 }
+
+#[test]
+fn after_start_and_before_stop_is_called() {
+    let _ = env_logger::try_init();
+
+    let after_start = Arc::new(atomic::AtomicUsize::new(0));
+    let before_stop = Arc::new(atomic::AtomicUsize::new(0));
+
+    let after_inner = after_start.clone();
+    let before_inner = before_stop.clone();
+    let runtime = tokio::runtime::Builder::new()
+            .after_start(move || { after_inner.clone().fetch_add(1, atomic::Ordering::Relaxed); })
+            .before_stop(move || { before_inner.clone().fetch_add(1, atomic::Ordering::Relaxed); })
+            .build()
+            .unwrap();
+
+    runtime.block_on_all(create_client_server_future()).unwrap();
+
+    assert!(after_start.load(atomic::Ordering::Relaxed) > 0);
+    assert!(before_stop.load(atomic::Ordering::Relaxed) > 0);
+}