From efe9659cbfae128b6b3ee866590f1bfde9b6f56f Mon Sep 17 00:00:00 2001
From: Chandler Carruth <chandlerc@gmail.com>
Date: Thu, 13 Nov 2025 08:47:59 +0000
Subject: [PATCH] Destroy tasks as they are run in the thread pool

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.
---
 llvm/lib/Support/ThreadPool.cpp       |  9 +++++--
 llvm/unittests/Support/ThreadPool.cpp | 37 +++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

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..7f7274740db7d 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.
+    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));