[llvm] [ORC] Fix synchronization in CoreAPIsTest. (PR #144556)

[Jameson Nash via llvm-commits]

https://github.com/vtjnash created https://github.com/llvm/llvm-project/pull/144556

The code previously appeared to have a (benign?) race condition on `WorkThreads.size`, since it was being accessed outside of the mutex lock that protected it on the threads. This is usually okay since 1a1d6e6f98738be249b20994bcfed48dccac59e3, but doesn't seem reliable in general, so fix this code to express the intent more accurately. This instead relies on the same general principles as ref-counting, where each existing reference (thread) can add new references (threads) because they already have a reference themselves (until joined).

>From cd954867e05a416101cac8b7d404f6fee700d070 Mon Sep 17 00:00:00 2001
From: Jameson Nash <vtjnash at gmail.com>
Date: Tue, 17 Jun 2025 15:54:29 +0000
Subject: [PATCH] [ORC] Fix synchronization in CoreAPIsTest.

The code previously appeared to have a (benign?) race condition on
`WorkThreads.size`, since it was being accessed outside of the mutex
lock that protected it on the threads. This is usually okay since
1a1d6e6f98738be249b20994bcfed48dccac59e3, but doesn't seem reliable in
general, so fix this code to express the intent more accurately. This
relies on the same general principles as ref-counting, where each
existing reference (thread) can add new references (threads) because
they already have a reference themselves (until joined).
---
 .../ExecutionEngine/Orc/CoreAPIsTest.cpp       | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 4a05a293fa51f..980dff86b3346 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -1559,16 +1559,13 @@ TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) {
 #if LLVM_ENABLE_THREADS
 
   std::mutex WorkThreadsMutex;
-  std::vector<std::thread> WorkThreads;
+  SmallVector<std::thread,0> WorkThreads;
   DispatchOverride = [&](std::unique_ptr<Task> T) {
-    std::promise<void> WaitP;
-    std::lock_guard<std::mutex> Lock(WorkThreadsMutex);
+    std::unique_lock Lock(WorkThreadsMutex);
     WorkThreads.push_back(
-        std::thread([T = std::move(T), WaitF = WaitP.get_future()]() mutable {
-          WaitF.get();
+        std::thread([T = std::move(T)]() mutable {
           T->run();
         }));
-    WaitP.set_value();
   };
 
   cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
@@ -1580,8 +1577,15 @@ TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) {
   EXPECT_EQ(FooLookupResult.getFlags(), FooSym.getFlags())
       << "lookup returned incorrect flags";
 
-  for (auto &WT : WorkThreads)
+  std::unique_lock Lock(WorkThreadsMutex);
+  // This works because every child thread that is allowed to use WorkThreads
+  // must either be in WorkThreads or its parent must be in WorkThreads.
+  while (!WorkThreads.empty()) {
+    auto WT = WorkThreads.pop_back_val();
+    Lock.unlock();
     WT.join();
+    Lock.lock();
+  }
 #endif
 }