runtime: add custom keep_alive functionality (#2809)

Co-authored-by: Eliza Weisman <eliza@buoyant.io>
Fixes: #2585

2 files changed, 43 insertions, 1 deletions

diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs
index 47895fcf..633021ed 100644
--- a/tokio/src/runtime/blocking/pool.rs
+++ b/tokio/src/runtime/blocking/pool.rs
@@ -47,6 +47,9 @@ struct Inner {
 
     // Maximum number of threads
     thread_cap: usize,
+
+    // Customizable wait timeout
+    keep_alive: Duration,
 }
 
 struct Shared {
@@ -91,6 +94,10 @@ where
 impl BlockingPool {
     pub(crate) fn new(builder: &Builder, thread_cap: usize) -> BlockingPool {
         let (shutdown_tx, shutdown_rx) = shutdown::channel();
+        #[cfg(feature = "blocking")]
+        let keep_alive = builder.keep_alive.unwrap_or(KEEP_ALIVE);
+        #[cfg(not(feature = "blocking"))]
+        let keep_alive = KEEP_ALIVE;
 
         BlockingPool {
             spawner: Spawner {
@@ -110,6 +117,7 @@ impl BlockingPool {
                     after_start: builder.after_start.clone(),
                     before_stop: builder.before_stop.clone(),
                     thread_cap,
+                    keep_alive,
                 }),
             },
             shutdown_rx,
@@ -258,7 +266,7 @@ impl Inner {
             shared.num_idle += 1;
 
             while !shared.shutdown {
-                let lock_result = self.condvar.wait_timeout(shared, KEEP_ALIVE).unwrap();
+                let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap();
 
                 shared = lock_result.0;
                 let timeout_result = lock_result.1;
diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs
index db01cf58..ed2cd251 100644
--- a/tokio/src/runtime/builder.rs
+++ b/tokio/src/runtime/builder.rs
@@ -4,6 +4,8 @@ use crate::runtime::shell::Shell;
 use crate::runtime::{blocking, io, time, Callback, Runtime, Spawner};
 
 use std::fmt;
+#[cfg(feature = "blocking")]
+use std::time::Duration;
 
 /// Builds Tokio Runtime with custom configuration values.
 ///
@@ -65,6 +67,11 @@ pub struct Builder {
 
     /// To run before each worker thread stops
     pub(super) before_stop: Option<Callback>,
+
+    #[cfg(feature = "blocking")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
+    /// Customizable keep alive timeout for BlockingPool
+    pub(super) keep_alive: Option<Duration>,
 }
 
 pub(crate) type ThreadNameFn = std::sync::Arc<dyn Fn() -> String + Send + Sync + 'static>;
@@ -108,6 +115,9 @@ impl Builder {
             // No worker thread callbacks
             after_start: None,
             before_stop: None,
+
+            #[cfg(feature = "blocking")]
+            keep_alive: None,
         }
     }
 
@@ -375,6 +385,30 @@ impl Builder {
             blocking_pool,
         })
     }
+
+    #[cfg(feature = "blocking")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
+    /// Sets a custom timeout for a thread in the blocking pool.
+    ///
+    /// By default, the timeout for a thread is set to 10 seconds. This can
+    /// be overriden using .thread_keep_alive().
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use tokio::runtime;
+    /// # use std::time::Duration;
+    ///
+    /// # pub fn main() {
+    /// let rt = runtime::Builder::new()
+    ///     .thread_keep_alive(Duration::from_millis(100))
+    ///     .build();
+    /// # }
+    /// ```
+    pub fn thread_keep_alive(&mut self, duration: Duration) -> &mut Self {
+        self.keep_alive = Some(duration);
+        self
+    }
 }
 
 cfg_io_driver! {