Destroy tasks as they are run in the thread pool

[chandlerc]

Without this, any RAII objects held in the task's captures aren't destroyed in a similar fashion to the task being run. If those objects in turn interact with the thread pool itself, chaos ensues. This comes up quite naturally with RAII-objects used for synchronization such as RAII-powered latches or releasing a mutex, etc.

A unit test is crafted that tries to very directly test that the logic of the thread pool continues to hold even with an RAII object. This isn't the only type of failure mode (a deadlock due to mutexes in the captures can also occur), but seemed the easiest to test.

[llvmbot]

@llvm/pr-subscribers-llvm-support

Author: Chandler Carruth (chandlerc)

Changes

Without this, any RAII objects held in the task's captures aren't destroyed in a similar fashion to the task being run. If those objects in turn interact with the thread pool itself, chaos ensues. This comes up quite naturally with RAII-objects used for synchronization such as RAII-powered latches or releasing a mutex, etc.

A unit test is crafted that tries to very directly test that the logic of the thread pool continues to hold even with an RAII object. This isn't the only type of failure mode (a deadlock due to mutexes in the captures can also occur), but seemed the easiest to test.

---

Full diff: https://github.com/llvm/llvm-project/pull/167852.diff

2 Files Affected:

• (modified) llvm/lib/Support/ThreadPool.cpp (+7-2)
• (modified) llvm/unittests/Support/ThreadPool.cpp (+37)

diff --git a/llvm/lib/Support/ThreadPool.cpp b/llvm/lib/Support/ThreadPool.cpp
index 4779e673cc055..2f9ff13109b61 100644
--- a/llvm/lib/Support/ThreadPool.cpp
+++ b/llvm/lib/Support/ThreadPool.cpp
@@ -110,8 +110,13 @@ void StdThreadPool::processTasks(ThreadPoolTaskGroup *WaitingForGroup) {
     CurrentThreadTaskGroups->push_back(GroupOfTask);
 #endif
 
-    // Run the task we just grabbed
-    Task();
+    // Run the task we just grabbed. This also destroys the task once run to
+    // release any resources held by it through RAII captured objects.
+    //
+    // It is particularly important to do this here so that we're not holding
+    // any lock and any further operations on the thread or `ThreadPool` take
+    // place here, at the same point as the task itself is executed.
+    std::exchange(Task, {})();
 
 #ifndef NDEBUG
     CurrentThreadTaskGroups->pop_back();
diff --git a/llvm/unittests/Support/ThreadPool.cpp b/llvm/unittests/Support/ThreadPool.cpp
index b5268c82e4199..49565fb8969c1 100644
--- a/llvm/unittests/Support/ThreadPool.cpp
+++ b/llvm/unittests/Support/ThreadPool.cpp
@@ -8,6 +8,7 @@
 
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Config/llvm-config.h" // for LLVM_ENABLE_THREADS
@@ -197,6 +198,42 @@ TYPED_TEST(ThreadPoolTest, AsyncMoveOnly) {
   ASSERT_EQ(42, f.get());
 }
 
+TYPED_TEST(ThreadPoolTest, AsyncRAIICaptures) {
+  CHECK_UNSUPPORTED();
+  DefaultThreadPool Pool(hardware_concurrency(2));
+
+  // We use a task group and a non-atomic value to stress test that the chaining
+  // of tasks via a captured RAII object in fact chains and synchronizes within
+  // a group.
+  ThreadPoolTaskGroup Group(Pool);
+  int value = 0;
+
+  // Create an RAII object that when destroyed schedules more work. This makes
+  // it easy to check that the RAII is resolved at the same point as a task runs
+  // on the thread pool.
+  auto schedule_next = llvm::make_scope_exit([&Group, &value] {
+    // We sleep before scheduling the final task to make it much more likely
+    // that an incorrect implementation actually exbitits a bug. Without the
+    // sleep, we may get "lucky" and have the second task finish before the
+    // assertion below fails even with an incorrect implementaiton. The
+    // sleep is making _failures_ more reliable, it is not needed for
+    // correctness and this test should only flakily _pass_, never flakily
+    // fail.
+    sleep(1);
+    Group.async([&value] { value = 42; });
+  });
+
+  // Now schedule the initial task, moving the RAII object to schedule the final
+  // task into its captures.
+  Group.async([schedule_next = std::move(schedule_next)]() {
+    // Nothing to do here, the captured RAII object does the work.
+  });
+
+  // Both tasks should complete here, synchronizing with the read of value.
+  Group.wait();
+  ASSERT_EQ(42, value);
+}
+
 TYPED_TEST(ThreadPoolTest, GetFuture) {
   CHECK_UNSUPPORTED();
   DefaultThreadPool Pool(hardware_concurrency(2));

[dwblaikie]

Appreciate the detailed comments in the test case, thanks!