Thread: Make Barrier reusable

1 files changed, 5 insertions, 5 deletions

diff --git a/src/common/thread.h b/src/common/thread.h
index 9bd0e8e7f..bbfa8befa 100644
--- a/src/common/thread.h
+++ b/src/common/thread.h
@@ -69,20 +69,19 @@ private:
 
 class Barrier {
 public:
-    explicit Barrier(size_t count_) : count(count_), waiting(0) {}
+    explicit Barrier(size_t count_) : count(count_), waiting(0), generation(0) {}
 
     /// Blocks until all "count" threads have called Sync()
     void Sync() {
         std::unique_lock<std::mutex> lk(mutex);
-
-        // TODO: broken when next round of Sync()s
-        // is entered before all waiting threads return from the notify_all
+        const size_t current_generation = generation;
 
         if (++waiting == count) {
+            generation++;
             waiting = 0;
             condvar.notify_all();
         } else {
-            condvar.wait(lk, [&]{ return waiting == 0; });
+            condvar.wait(lk, [this, current_generation]{ return current_generation != generation; });
         }
     }
 
@@ -91,6 +90,7 @@ private:
     std::mutex mutex;
     const size_t count;
     size_t waiting;
+    size_t generation; // Incremented once each time the barrier is used
 };
 
 void SleepCurrentThread(int ms);