From d3dca4552b300ad7b43491bcb67f8ddbebab503d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <bkchr@users.noreply.github.com>
Date: Mon, 19 Nov 2018 18:04:58 +0100
Subject: [PATCH] Expose `after_start` and `before_stop` in `runtime::Builder`
 (#756)

Closes #705
---
 src/runtime/builder.rs | 53 ++++++++++++++++++++++++++++++++++++++++++
 tests/runtime.rs       | 23 +++++++++++++++++-
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/src/runtime/builder.rs b/src/runtime/builder.rs
index 2ba3ea98e..8a5dca772 100644
--- a/src/runtime/builder.rs
+++ b/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.
diff --git a/tests/runtime.rs b/tests/runtime.rs
index 39e8f12f5..ed77dba5b 100644
--- a/tests/runtime.rs
+++ b/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);
+}