diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 47895fcf477..633021ededf 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 db01cf5871e..ed2cd251c35 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, + + #[cfg(feature = "blocking")] + #[cfg_attr(docsrs, doc(cfg(feature = "blocking")))] + /// Customizable keep alive timeout for BlockingPool + pub(super) keep_alive: Option, } pub(crate) type ThreadNameFn = std::sync::Arc 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! {