Destroy tasks as they are run in the thread pool

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();

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.
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    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));