From 9352249c3e1282ef99e844d970b6996ba804ede4 Mon Sep 17 00:00:00 2001
From: Stjepan Glavina <stjepang@gmail.com>
Date: Tue, 31 Jul 2018 05:48:53 +0200
Subject: [PATCH] Terminate backup threads when idle (#489)

---
 tokio-threadpool/src/pool/backup.rs | 58 ++++++++++++++++-------------
 tokio-threadpool/src/pool/mod.rs    | 12 ++----
 2 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/tokio-threadpool/src/pool/backup.rs b/tokio-threadpool/src/pool/backup.rs
index 4f0c295a5..5380a32b2 100644
--- a/tokio-threadpool/src/pool/backup.rs
+++ b/tokio-threadpool/src/pool/backup.rs
@@ -5,6 +5,7 @@ use std::cell::UnsafeCell;
 use std::fmt;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::{self, Acquire, AcqRel, Relaxed};
+use std::time::{Duration, Instant};
 
 /// State associated with a thread in the thread pool.
 ///
@@ -155,7 +156,8 @@ impl Backup {
     }
 
     /// Wait for a worker handoff
-    pub fn wait_for_handoff(&self, sleep: bool) -> Handoff {
+    pub fn wait_for_handoff(&self, timeout: Option<Duration>) -> Handoff {
+        let sleep_until = timeout.map(|dur| Instant::now() + dur);
         let mut state: State = self.state.load(Acquire).into();
 
         // Run in a loop since there can be spurious wakeups
@@ -169,36 +171,40 @@ impl Backup {
                     (*self.handoff.get()).take()
                         .expect("no worker handoff")
                 };
-
                 return Handoff::Worker(worker_id);
             }
 
-            if sleep {
-                // TODO: Park with a timeout
-                self.park.park_sync(None);
-
-                // Reload the state
-                state = self.state.load(Acquire).into();
-                debug_assert!(state.is_running());
-            } else {
-                debug_assert!(state.is_running());
-
-                // Transition out of running
-                let mut next = state;
-                next.unset_running();
-
-                let actual = self.state.compare_and_swap(
-                    state.into(),
-                    next.into(),
-                    AcqRel).into();
-
-                if actual == state {
-                    debug_assert!(!next.is_running());
-
-                    return Handoff::Idle;
+            match sleep_until {
+                None => {
+                    self.park.park_sync(None);
+                    state = self.state.load(Acquire).into();
                 }
+                Some(when) => {
+                    let now = Instant::now();
 
-                state = actual;
+                    if now < when {
+                        self.park.park_sync(Some(when - now));
+                        state = self.state.load(Acquire).into();
+                    } else {
+                        debug_assert!(state.is_running());
+
+                        // Transition out of running
+                        let mut next = state;
+                        next.unset_running();
+
+                        let actual = self.state.compare_and_swap(
+                            state.into(),
+                            next.into(),
+                            AcqRel).into();
+
+                        if actual == state {
+                            debug_assert!(!next.is_running());
+                            return Handoff::Idle;
+                        }
+
+                        state = actual;
+                    }
+                }
             }
         }
     }
diff --git a/tokio-threadpool/src/pool/mod.rs b/tokio-threadpool/src/pool/mod.rs
index 7fbcf94ad..0c080b1c7 100644
--- a/tokio-threadpool/src/pool/mod.rs
+++ b/tokio-threadpool/src/pool/mod.rs
@@ -417,6 +417,8 @@ impl Pool {
                     break;
                 }
 
+                debug_assert!(!inner.backup[backup_id.0].is_pushed());
+
                 // Push the thread back onto the backup stack. This makes it
                 // available for future handoffs.
                 //
@@ -437,20 +439,14 @@ impl Pool {
 
                 // Wait for a handoff
                 let handoff = inner.backup[backup_id.0]
-                    .wait_for_handoff(true);
+                    .wait_for_handoff(inner.config.keep_alive);
 
                 match handoff {
                     Handoff::Worker(id) => {
                         debug_assert!(inner.backup[backup_id.0].is_running());
                         worker_id = id;
                     }
-                    Handoff::Idle => {
-                        // Worker is idle
-                        break;
-                    }
-                    Handoff::Terminated => {
-                        // TODO: When wait_for_handoff supports blocking with a
-                        // timeout, this will have to be smarter
+                    Handoff::Idle | Handoff::Terminated => {
                         break;
                     }
                 }