3 files changed, 49 insertions, 6 deletions

diff --git a/src/buildtool/multithreading/task_system.cpp b/src/buildtool/multithreading/task_system.cpp
index 4af57d0d..ee5be910 100644
--- a/src/buildtool/multithreading/task_system.cpp
+++ b/src/buildtool/multithreading/task_system.cpp
@@ -17,12 +17,7 @@ TaskSystem::TaskSystem(std::size_t number_of_threads)
 }
 
 TaskSystem::~TaskSystem() {
-    // When starting a new task system all spawned threads will immediately go
-    // to sleep and wait for tasks. Even after adding some tasks, it can take a
-    // while until the first thread wakes up. Therefore, we need to wait for the
-    // total workload (number of active threads _and_ total number of queued
-    // tasks) to become zero.
-    total_workload_.WaitForZero();
+    Finish();
     for (auto& q : queues_) {
         q.done();
     }
@@ -31,6 +26,15 @@ TaskSystem::~TaskSystem() {
     }
 }
 
+void TaskSystem::Finish() noexcept {
+    // When starting a new task system all spawned threads will immediately go
+    // to sleep and wait for tasks. Even after adding some tasks, it can take a
+    // while until the first thread wakes up. Therefore, we need to wait for the
+    // total workload (number of active threads _and_ total number of queued
+    // tasks) to become zero.
+    total_workload_.WaitForZero();
+}
+
 void TaskSystem::Run(std::size_t idx) {
     gsl_Expects(thread_count_ > 0);
 
diff --git a/src/buildtool/multithreading/task_system.hpp b/src/buildtool/multithreading/task_system.hpp
index 6387988a..888ba090 100644
--- a/src/buildtool/multithreading/task_system.hpp
+++ b/src/buildtool/multithreading/task_system.hpp
@@ -48,6 +48,9 @@ class TaskSystem {
         return thread_count_;
     }
 
+    // Wait for all queues to become empty _and_ all tasks to finish.
+    void Finish() noexcept;
+
   private:
     std::size_t const thread_count_{
         std::max(1U, std::thread::hardware_concurrency())};
diff --git a/test/buildtool/multithreading/task_system.test.cpp b/test/buildtool/multithreading/task_system.test.cpp
index d34e3e4b..abcae8f4 100644
--- a/test/buildtool/multithreading/task_system.test.cpp
+++ b/test/buildtool/multithreading/task_system.test.cpp
@@ -222,3 +222,39 @@ TEST_CASE("All threads run until work is done", "[task_system]") {
         CHECK(tids.size() == kNumThreads);
     }
 }
+
+TEST_CASE("Use finish as system-wide barrier", "[task_system]") {
+    using namespace std::chrono_literals;
+    static auto const kNumThreads = std::thread::hardware_concurrency();
+
+    std::vector<int> vec(kNumThreads, 0);
+    std::vector<int> exp0(kNumThreads, 0);
+    std::vector<int> exp1(kNumThreads, 1);
+    std::vector<int> exp2(kNumThreads, 2);
+
+    {
+        TaskSystem ts{kNumThreads};
+
+        // Wait for all threads to go to sleep.
+        ts.Finish();
+        CHECK(vec == exp0);
+
+        for (std::size_t i{}; i < ts.NumberOfThreads(); ++i) {
+            ts.QueueTask([&vec, i] {
+                std::this_thread::sleep_for(1s);
+                vec[i] = 1;
+            });
+        }
+
+        ts.Finish();
+        CHECK(vec == exp1);
+
+        for (std::size_t i{}; i < ts.NumberOfThreads(); ++i) {
+            ts.QueueTask([&vec, i] {
+                std::this_thread::sleep_for(1s);
+                vec[i] = 2;
+            });
+        }
+    }
+    CHECK(vec == exp2);
+}