[Support] Fix fragile ThreadPool test

[aganea]

Summary

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.

[mehdi_amini]

llvm/unittests/Support/ThreadPool.cpp

200

This still leaves the test sensitive to a timing effect, can we make it entirely deterministic?

For example can we increment a counter with a condition variable at the beginning of the lambda before waitForMainThread() and wait on the variable in the main thread until it reaches llvm::get_cpus() before calling into setMainThreadReady();?

[aganea]

As suggested by @mehdi_amini

[MaskRay]

LGTM as well.

llvm/unittests/Support/ThreadPool.cpp

188

Supposed this is deterministic now because of the llvm::ThreadPoolStrategy::apply_thread_strategy magic on Windows.

llvm/unittests/Support/ThreadPool.cpp

void setMainThreadReady() {
       std::unique_lock<std::mutex> LockGuard(WaitMainThreadMutex);
       MainThreadReady = true;
     }
     WaitMainThread.notify_all();
   }
 
   void SetUp() override { MainThreadReady = false; }
 
-  void TestAllThreads(ThreadPoolStrategy S);
+  void RunOnAllSockets(ThreadPoolStrategy S);
 
   std::condition_variable WaitMainThread;
   std::mutex WaitMainThreadMutex;
   bool MainThreadReady = false;
 };
 
 #define CHECK_UNSUPPORTED() \
   do { \
std::atomic_int checked_in{0};
     ASSERT_EQ(0, checked_in);
     setMainThreadReady();
   }
   ASSERT_EQ(5, checked_in);
 }
 
 #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.
   Triple Host(Triple::normalize(sys::getProcessTriple()));
   if (!Host.isOSWindows())
     return;
 
   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();
         ThreadsUsed.insert(Mask);
       });
     }
     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());
 }
 
 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