From 933250e355180d2d2c7815f79a186584e1963427 Mon Sep 17 00:00:00 2001
From: Yi Wu <yiwu@fb.com>
Date: Sun, 21 Oct 2018 20:14:58 -0700
Subject: [PATCH] Fix RepeatableThreadTest::MockEnvTest hang (#4560)

Summary:
When `MockTimeEnv` is used in test to mock time methods, we cannot use `CondVar::TimedWait` because it is using real time, not the mocked time for wait timeout. On Mac the method can return immediately without awaking other waiting threads, if the real time is larger than `wait_until` (which is a mocked time). When that happen, the `wait()` method will fall into an infinite loop.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4560

Differential Revision: D10472851

Pulled By: yiwu-arbug

fbshipit-source-id: 898902546ace7db7ac509337dd8677a527209d19
---
 db/db_test_util.h        | 32 +-----------------------------
 util/mock_time_env.h     | 43 ++++++++++++++++++++++++++++++++++++++++
 util/repeatable_thread.h | 11 ++++++++++
 3 files changed, 55 insertions(+), 31 deletions(-)
 create mode 100644 util/mock_time_env.h

diff --git a/db/db_test_util.h b/db/db_test_util.h
index c4edff83e..7f26e708a 100644
--- a/db/db_test_util.h
+++ b/db/db_test_util.h
@@ -46,6 +46,7 @@
 #include "table/scoped_arena_iterator.h"
 #include "util/compression.h"
 #include "util/filename.h"
+#include "util/mock_time_env.h"
 #include "util/mutexlock.h"
 
 #include "util/string_util.h"
@@ -582,37 +583,6 @@ class SpecialEnv : public EnvWrapper {
   std::atomic<size_t> compaction_readahead_size_{};
 };
 
-class MockTimeEnv : public EnvWrapper {
- public:
-  explicit MockTimeEnv(Env* base) : EnvWrapper(base) {}
-
-  virtual Status GetCurrentTime(int64_t* time) override {
-    assert(time != nullptr);
-    assert(current_time_ <=
-           static_cast<uint64_t>(std::numeric_limits<int64_t>::max()));
-    *time = static_cast<int64_t>(current_time_);
-    return Status::OK();
-  }
-
-  virtual uint64_t NowMicros() override {
-    assert(current_time_ <= std::numeric_limits<uint64_t>::max() / 1000000);
-    return current_time_ * 1000000;
-  }
-
-  virtual uint64_t NowNanos() override {
-    assert(current_time_ <= std::numeric_limits<uint64_t>::max() / 1000000000);
-    return current_time_ * 1000000000;
-  }
-
-  void set_current_time(uint64_t time) {
-    assert(time >= current_time_);
-    current_time_ = time;
-  }
-
- private:
-  std::atomic<uint64_t> current_time_{0};
-};
-
 #ifndef ROCKSDB_LITE
 class OnFileDeletionListener : public EventListener {
  public:
diff --git a/util/mock_time_env.h b/util/mock_time_env.h
new file mode 100644
index 000000000..c6ab8a748
--- /dev/null
+++ b/util/mock_time_env.h
@@ -0,0 +1,43 @@
+// Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#include "rocksdb/env.h"
+
+namespace rocksdb {
+
+class MockTimeEnv : public EnvWrapper {
+ public:
+  explicit MockTimeEnv(Env* base) : EnvWrapper(base) {}
+
+  virtual Status GetCurrentTime(int64_t* time) override {
+    assert(time != nullptr);
+    assert(current_time_ <=
+           static_cast<uint64_t>(std::numeric_limits<int64_t>::max()));
+    *time = static_cast<int64_t>(current_time_);
+    return Status::OK();
+  }
+
+  virtual uint64_t NowMicros() override {
+    assert(current_time_ <= std::numeric_limits<uint64_t>::max() / 1000000);
+    return current_time_ * 1000000;
+  }
+
+  virtual uint64_t NowNanos() override {
+    assert(current_time_ <= std::numeric_limits<uint64_t>::max() / 1000000000);
+    return current_time_ * 1000000000;
+  }
+
+  void set_current_time(uint64_t time) {
+    assert(time >= current_time_);
+    current_time_ = time;
+  }
+
+ private:
+  std::atomic<uint64_t> current_time_{0};
+};
+
+}  // namespace rocksdb
diff --git a/util/repeatable_thread.h b/util/repeatable_thread.h
index 34164ca56..3506234f9 100644
--- a/util/repeatable_thread.h
+++ b/util/repeatable_thread.h
@@ -10,6 +10,7 @@
 
 #include "port/port.h"
 #include "rocksdb/env.h"
+#include "util/mock_time_env.h"
 #include "util/mutexlock.h"
 
 namespace rocksdb {
@@ -80,7 +81,17 @@ class RepeatableThread {
       cond_var_.SignalAll();
 #endif
       while (running_) {
+#ifndef NDEBUG
+        if (dynamic_cast<MockTimeEnv*>(env_) != nullptr) {
+          // MockTimeEnv is used. Since it is not easy to mock TimedWait,
+          // we wait without timeout to wait for TEST_WaitForRun to wake us up.
+          cond_var_.Wait();
+        } else {
+          cond_var_.TimedWait(wait_until);
+        }
+#else
         cond_var_.TimedWait(wait_until);
+#endif
         if (env_->NowMicros() >= wait_until) {
           break;
         }