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