From b8913ec7c0dbb83592d7f6a500568b1a7d051469 Mon Sep 17 00:00:00 2001
From: Jonathan Bastien-Filiatrault <jonathan@zerospam.ca>
Date: Mon, 7 Oct 2019 17:04:28 -0400
Subject: [PATCH] executor: accurate idle thread tracking for the blocking pool
 (#1621)

Use a counter to count notifications. This protects against spurious
wakeups by pthreads and other libraries. The state transitions now
track num_idle precisely.
---
 tokio-executor/src/blocking.rs | 52 ++++++++++++++++++++++++----------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/tokio-executor/src/blocking.rs b/tokio-executor/src/blocking.rs
index 424d46767..724872466 100644
--- a/tokio-executor/src/blocking.rs
+++ b/tokio-executor/src/blocking.rs
@@ -20,6 +20,7 @@ struct Shared {
     queue: VecDeque<Box<dyn FnOnce() + Send>>,
     num_th: u32,
     num_idle: u32,
+    num_notify: u32,
 }
 
 lazy_static! {
@@ -47,12 +48,12 @@ where
         let mut shared = POOL.shared.lock().unwrap();
 
         shared.queue.push_back(Box::new(move || {
-            // The receiver may have dropped
+            // The receiver may have been dropped.
             let _ = tx.send(f());
         }));
 
         if shared.num_idle == 0 {
-            // No threads are able to process the task
+            // No threads are able to process the task.
 
             if shared.num_th == MAX_THREADS {
                 // At max number of threads
@@ -62,7 +63,13 @@ where
                 true
             }
         } else {
+            // Notify an idle worker thread. The notification counter
+            // is used to count the needed amount of notifications
+            // exactly. Thread libraries may generate spurious
+            // wakeups, this counter is used to keep us in a
+            // consistent state.
             shared.num_idle -= 1;
+            shared.num_notify += 1;
             POOL.condvar.notify_one();
             false
         }
@@ -96,13 +103,13 @@ fn spawn_thread() {
     thread::Builder::new()
         .name("tokio-blocking-driver".to_string())
         .spawn(|| {
-            'outer: loop {
-                let mut shared = POOL.shared.lock().unwrap();
-
-                if let Some(task) = shared.queue.pop_front() {
+            let mut shared = POOL.shared.lock().unwrap();
+            loop {
+                // BUSY
+                while let Some(task) = shared.queue.pop_front() {
                     drop(shared);
                     run_task(task);
-                    continue;
+                    shared = POOL.shared.lock().unwrap();
                 }
 
                 // IDLE
@@ -113,15 +120,29 @@ fn spawn_thread() {
                     shared = lock_result.0;
                     let timeout_result = lock_result.1;
 
-                    if let Some(task) = shared.queue.pop_front() {
-                        drop(shared);
-                        run_task(task);
-                        continue 'outer;
-                    } else if timeout_result.timed_out() {
-                        shared.num_idle = shared.num_idle.saturating_sub(1);
-                        shared.num_th -= 1;
-                        break 'outer;
+                    if shared.num_notify != 0 {
+                        // We have received a legitimate wakeup,
+                        // acknowledge it by decrementing the counter
+                        // and transition to the BUSY state.
+                        shared.num_notify -= 1;
+                        break;
                     }
+
+                    if timeout_result.timed_out() {
+                        // Thread exit
+                        shared.num_th -= 1;
+
+                        // num_idle should now be tracked exactly, panic
+                        // with a descriptive message if it is not the
+                        // case.
+                        shared.num_idle = shared
+                            .num_idle
+                            .checked_sub(1)
+                            .expect("num_idle underflowed on thread exit");
+                        return;
+                    }
+
+                    // Spurious wakeup detected, go back to sleep.
                 }
             }
         })
@@ -141,6 +162,7 @@ impl Pool {
                 queue: VecDeque::new(),
                 num_th: 0,
                 num_idle: 0,
+                num_notify: 0,
             }),
             condvar: Condvar::new(),
         }