[unittests] ThreadPool: Guard updates to MainThreadReady

[oota-llvm.git] / unittests / Support / ThreadPool.cpp

diff --git a/unittests/Support/ThreadPool.cpp b/unittests/Support/ThreadPool.cpp
index ca55237d49102bcda5e61188262b5439b9564a7f..80b89e3c0ab30c80347e434c58ad7f63a324cdde 100644 (file)
--- a/unittests/Support/ThreadPool.cpp
+++ b/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);
 }