Add support for waiting on other threads task counts in scheduler

2 files changed, 34 insertions, 13 deletions

diff --git a/build2/scheduler b/build2/scheduler
index 563ec31..42f82e0 100644
--- a/build2/scheduler
+++ b/build2/scheduler
@@ -73,8 +73,11 @@ namespace build2
     void
     async (atomic_count& task_count, F&&, A&&...);
 
-    // Wait until the task count reaches 0. If the scheduler is shutdown
-    // while waiting, throw system_error(ECANCELED).
+    // Wait until the task count reaches 0. If the scheduler is shutdown while
+    // waiting, throw system_error(ECANCELED).
+    //
+    // Note that it is valid to wait on another thread's task count (that is,
+    // without making any async() calls in this thread).
     //
     void
     wait (atomic_count& task_count);
@@ -288,11 +291,15 @@ namespace build2
     // depends on the number of waiters that we can have which cannot be
     // greater than the total number of threads.
     //
+    // The pointer to the task count is used to identify the already waiting
+    // group of threads for collision statistics.
+    //
     struct wait_slot
     {
       std::mutex mutex;
       std::condition_variable condv;
       size_t waiters = 0;
+      const atomic_count* tcount;
       bool shutdown = true;
     };
 
diff --git a/build2/scheduler.cxx b/build2/scheduler.cxx
index 47d38e4..c151847 100644
--- a/build2/scheduler.cxx
+++ b/build2/scheduler.cxx
@@ -18,17 +18,20 @@ namespace build2
 
     // See if we can run some of our own tasks.
     //
-    task_queue& tq (*task_queue_); // Must have been set by async() or task
-                                   // would have been 0.
-
-    for (lock ql (tq.mutex); !tq.shutdown && !empty_back (tq); )
-      pop_back (tq, ql);
-
-    // Note that empty task queue doesn't automatically mean the task count
-    // is zero (some might still be executing asynchronously).
+    // If we are waiting on someone else's task count then there migh still
+    // be no queue which is set by async().
     //
-    if (task_count == 0)
-      return;
+    if (task_queue* tq = task_queue_)
+    {
+      for (lock ql (tq->mutex); !tq->shutdown && !empty_back (*tq); )
+        pop_back (*tq, ql);
+
+      // Note that empty task queue doesn't automatically mean the task count
+      // is zero (some might still be executing asynchronously).
+      //
+      if (task_count == 0)
+        return;
+    }
 
     suspend (task_count);
   }
@@ -66,7 +69,18 @@ namespace build2
     bool collision;
     {
       lock l (s.mutex);
-      collision = (s.waiters++ != 0);
+
+      // We have a collision if there is already a waiter for a different
+      // task count.
+      //
+      collision = (s.waiters++ != 0 && s.tcount != &tc);
+
+      // This is nuanced: we want to always have the task count of the last
+      // thread to join the queue. Otherwise, if threads are leaving and
+      // joining the queue simultaneously, we may end up with a task count of
+      // a thread group that is no longer waiting.
+      //
+      s.tcount = &tc;
 
       // Since we use a mutex for synchronization, we can relax the atomic
       // access.