[unittests] ThreadPool: Guard updates to MainThreadReady

llvm-svn: 256096

1 files changed, 16 insertions, 12 deletions

diff --git a/llvm/unittests/Support/ThreadPool.cpp b/llvm/unittests/Support/ThreadPool.cpp
index ca55237..80b89e3c 100644
--- a/llvm/unittests/Support/ThreadPool.cpp
+++ b/llvm/unittests/Support/ThreadPool.cpp
@@ -62,6 +62,14 @@ protected:
       WaitMainThread.wait(LockGuard, [&] { return MainThreadReady; });
     }
   }
+
+  /// Set the readiness of the main thread.
+  void setMainThreadReadyState(bool Ready) {
+    std::unique_lock<std::mutex> LockGuard(WaitMainThreadMutex);
+    MainThreadReady = Ready;
+    WaitMainThread.notify_all();
+  }
+
   std::condition_variable WaitMainThread;
   std::mutex WaitMainThreadMutex;
   bool MainThreadReady;
@@ -80,7 +88,7 @@ TEST_F(ThreadPoolTest, AsyncBarrier) {
 
   std::atomic_int checked_in{0};
 
-  MainThreadReady = false;
+  setMainThreadReadyState(false);
   ThreadPool Pool;
   for (size_t i = 0; i < 5; ++i) {
     Pool.async([this, &checked_in, i] {
@@ -89,8 +97,7 @@ TEST_F(ThreadPoolTest, AsyncBarrier) {
     });
   }
   ASSERT_EQ(0, checked_in);
-  MainThreadReady = true;
-  WaitMainThread.notify_all();
+  setMainThreadReadyState(true);
   Pool.wait();
   ASSERT_EQ(5, checked_in);
 }
@@ -114,15 +121,14 @@ TEST_F(ThreadPoolTest, Async) {
   CHECK_UNSUPPORTED();
   ThreadPool Pool;
   std::atomic_int i{0};
-  MainThreadReady = false;
+  setMainThreadReadyState(false);
   Pool.async([this, &i] {
     waitForMainThread();
     ++i;
   });
   Pool.async([&i] { ++i; });
   ASSERT_NE(2, i.load());
-  MainThreadReady = true;
-  WaitMainThread.notify_all();
+  setMainThreadReadyState(true);
   Pool.wait();
   ASSERT_EQ(2, i.load());
 }
@@ -131,7 +137,7 @@ TEST_F(ThreadPoolTest, GetFuture) {
   CHECK_UNSUPPORTED();
   ThreadPool Pool;
   std::atomic_int i{0};
-  MainThreadReady = false;
+  setMainThreadReadyState(false);
   Pool.async([this, &i] {
     waitForMainThread();
     ++i;
@@ -139,8 +145,7 @@ TEST_F(ThreadPoolTest, GetFuture) {
   // Force the future using get()
   Pool.async([&i] { ++i; }).get();
   ASSERT_NE(2, i.load());
-  MainThreadReady = true;
-  WaitMainThread.notify_all();
+  setMainThreadReadyState(true);
   Pool.wait();
   ASSERT_EQ(2, i.load());
 }
@@ -150,7 +155,7 @@ TEST_F(ThreadPoolTest, PoolDestruction) {
   // Test that we are waiting on destruction
   std::atomic_int checked_in{0};
   {
-    MainThreadReady = false;
+    setMainThreadReadyState(false);
     ThreadPool Pool;
     for (size_t i = 0; i < 5; ++i) {
       Pool.async([this, &checked_in, i] {
@@ -159,8 +164,7 @@ TEST_F(ThreadPoolTest, PoolDestruction) {
       });
     }
     ASSERT_EQ(0, checked_in);
-    MainThreadReady = true;
-    WaitMainThread.notify_all();
+    setMainThreadReadyState(true);
   }
   ASSERT_EQ(5, checked_in);
 }