[Support] Fix fragile ThreadPool test

The test ThreadPoolTest.AllThreads_UseAllRessources occasionally fails on the bots: http://lab.llvm.org:8011/builders/llvm-clang-lld-x86_64-scei-ps4-windows10pro-fast/builds/32015/steps/test-check-all/logs/FAIL%3A%20LLVM-Unit%3A%3AThreadPoolTest.AllThreads_UseAllRessources

This is because jobs were executed too fast on the first CPU socket, and never manage to reach the second CPU socket.

Differential Revision: https://reviews.llvm.org/D78832

diff --git a/llvm/unittests/Support/ThreadPool.cpp b/llvm/unittests/Support/ThreadPool.cpp
index 25e71ae..43882d0 100644 (file)
--- a/llvm/unittests/Support/ThreadPool.cpp
+++ b/llvm/unittests/Support/ThreadPool.cpp
@@ -71,7 +71,7 @@ protected:
 
   void SetUp() override { MainThreadReady = false; }
 
-  void TestAllThreads(ThreadPoolStrategy S);
+  void RunOnAllSockets(ThreadPoolStrategy S);
 
   std::condition_variable WaitMainThread;
   std::mutex WaitMainThreadMutex;
@@ -169,7 +169,7 @@ TEST_F(ThreadPoolTest, PoolDestruction) {
 
 #if LLVM_ENABLE_THREADS == 1
 
-void ThreadPoolTest::TestAllThreads(ThreadPoolStrategy S) {
+void ThreadPoolTest::RunOnAllSockets(ThreadPoolStrategy S) {
   // FIXME: Skip these tests on non-Windows because multi-socket system were not
   // tested on Unix yet, and llvm::get_thread_affinity_mask() isn't implemented
   // for Unix.
@@ -180,9 +180,18 @@ void ThreadPoolTest::TestAllThreads(ThreadPoolStrategy S) {
   llvm::DenseSet<llvm::BitVector> ThreadsUsed;
   std::mutex Lock;
   {
+    std::condition_variable AllThreads;
+    std::mutex AllThreadsLock;
+    unsigned Active = 0;
+
     ThreadPool Pool(S);
-    for (size_t I = 0; I < 10000; ++I) {
+    for (size_t I = 0; I < S.compute_thread_count(); ++I) {
       Pool.async([&] {
+        {
+          std::lock_guard<std::mutex> Guard(AllThreadsLock);
+          ++Active;
+          AllThreads.notify_one();
+        }
         waitForMainThread();
         std::lock_guard<std::mutex> Guard(Lock);
         auto Mask = llvm::get_thread_affinity_mask();
@@ -190,6 +199,11 @@ void ThreadPoolTest::TestAllThreads(ThreadPoolStrategy S) {
       });
     }
     ASSERT_EQ(true, ThreadsUsed.empty());
+    {
+      std::unique_lock<std::mutex> Guard(AllThreadsLock);
+      AllThreads.wait(Guard,
+                      [&]() { return Active == S.compute_thread_count(); });
+    }
     setMainThreadReady();
   }
   ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size());
@@ -197,12 +211,12 @@ void ThreadPoolTest::TestAllThreads(ThreadPoolStrategy S) {
 
 TEST_F(ThreadPoolTest, AllThreads_UseAllRessources) {
   CHECK_UNSUPPORTED();
-  TestAllThreads({});
+  RunOnAllSockets({});
 }
 
 TEST_F(ThreadPoolTest, AllThreads_OneThreadPerCore) {
   CHECK_UNSUPPORTED();
-  TestAllThreads(llvm::heavyweight_hardware_concurrency());
+  RunOnAllSockets(llvm::heavyweight_hardware_concurrency());
 }
 
 #endif