Fix tsan problem in env_test.

PiperOrigin-RevId: 268265314

1 files changed, 49 insertions, 37 deletions

diff --git a/util/env_test.cc b/util/env_test.cc
index 9e2ad1e..7db03fc 100644
--- a/util/env_test.cc
+++ b/util/env_test.cc
@@ -5,7 +5,6 @@
 #include "leveldb/env.h"
 
 #include <algorithm>
-#include <atomic>
 
 #include "port/port.h"
 #include "port/thread_annotations.h"
@@ -24,16 +23,6 @@ class EnvTest {
   Env* env_;
 };
 
-namespace {
-
-static void SetAtomicBool(void* atomic_bool_ptr) {
-  std::atomic<bool>* atomic_bool =
-      reinterpret_cast<std::atomic<bool>*>(atomic_bool_ptr);
-  atomic_bool->store(true, std::memory_order_relaxed);
-}
-
-}  // namespace
-
 TEST(EnvTest, ReadWrite) {
   Random rnd(test::RandomSeed());
 
@@ -82,45 +71,73 @@ TEST(EnvTest, ReadWrite) {
 }
 
 TEST(EnvTest, RunImmediately) {
-  std::atomic<bool> called(false);
-  env_->Schedule(&SetAtomicBool, &called);
-  env_->SleepForMicroseconds(kDelayMicros);
-  ASSERT_TRUE(called.load(std::memory_order_relaxed));
+  struct RunState {
+    port::Mutex mu;
+    port::CondVar cvar{&mu};
+    bool called = false;
+
+    static void Run(void* arg) {
+      RunState* state = reinterpret_cast<RunState*>(arg);
+      MutexLock l(&state->mu);
+      ASSERT_EQ(state->called, false);
+      state->called = true;
+      state->cvar.Signal();
+    }
+  };
+
+  RunState state;
+  env_->Schedule(&RunState::Run, &state);
+
+  MutexLock l(&state.mu);
+  while (!state.called) {
+    state.cvar.Wait();
+  }
 }
 
 TEST(EnvTest, RunMany) {
-  std::atomic<int> last_id(0);
+  struct RunState {
+    port::Mutex mu;
+    port::CondVar cvar{&mu};
+    int last_id = 0;
+  };
 
   struct Callback {
-    std::atomic<int>* const last_id_ptr_;  // Pointer to shared state.
+    RunState* state_;  // Pointer to shared state.
     const int id_;  // Order# for the execution of this callback.
 
-    Callback(std::atomic<int>* last_id_ptr, int id)
-        : last_id_ptr_(last_id_ptr), id_(id) {}
+    Callback(RunState* s, int id) : state_(s), id_(id) {}
 
     static void Run(void* arg) {
       Callback* callback = reinterpret_cast<Callback*>(arg);
-      int current_id = callback->last_id_ptr_->load(std::memory_order_relaxed);
-      ASSERT_EQ(callback->id_ - 1, current_id);
-      callback->last_id_ptr_->store(callback->id_, std::memory_order_relaxed);
+      RunState* state = callback->state_;
+
+      MutexLock l(&state->mu);
+      ASSERT_EQ(state->last_id, callback->id_ - 1);
+      state->last_id = callback->id_;
+      state->cvar.Signal();
     }
   };
 
-  Callback callback1(&last_id, 1);
-  Callback callback2(&last_id, 2);
-  Callback callback3(&last_id, 3);
-  Callback callback4(&last_id, 4);
+  RunState state;
+  Callback callback1(&state, 1);
+  Callback callback2(&state, 2);
+  Callback callback3(&state, 3);
+  Callback callback4(&state, 4);
   env_->Schedule(&Callback::Run, &callback1);
   env_->Schedule(&Callback::Run, &callback2);
   env_->Schedule(&Callback::Run, &callback3);
   env_->Schedule(&Callback::Run, &callback4);
 
-  env_->SleepForMicroseconds(kDelayMicros);
-  ASSERT_EQ(4, last_id.load(std::memory_order_relaxed));
+  MutexLock l(&state.mu);
+  while (state.last_id != 4) {
+    state.cvar.Wait();
+  }
 }
 
 struct State {
   port::Mutex mu;
+  port::CondVar cvar{&mu};
+
   int val GUARDED_BY(mu);
   int num_running GUARDED_BY(mu);
 
@@ -132,6 +149,7 @@ static void ThreadBody(void* arg) {
   s->mu.Lock();
   s->val += 1;
   s->num_running -= 1;
+  s->cvar.Signal();
   s->mu.Unlock();
 }
 
@@ -140,17 +158,11 @@ TEST(EnvTest, StartThread) {
   for (int i = 0; i < 3; i++) {
     env_->StartThread(&ThreadBody, &state);
   }
-  while (true) {
-    state.mu.Lock();
-    int num = state.num_running;
-    state.mu.Unlock();
-    if (num == 0) {
-      break;
-    }
-    env_->SleepForMicroseconds(kDelayMicros);
-  }
 
   MutexLock l(&state.mu);
+  while (state.num_running != 0) {
+    state.cvar.Wait();
+  }
   ASSERT_EQ(state.val, 3);
 }