Fix flakiness of WalManagerTest

Summary: We should use mocked-out env for these tests to make it more realiable. Added benefit is that instead of actually sleeping for 3 seconds, we can instead pretend to sleep and just increase time counters.

Test Plan: for i in `seq 100`; do ./wal_manager_test --gtest_filter=WalManagerTest.WALArchivalTtl ;done

Reviewers: rven, meyering

Reviewed By: meyering

Subscribers: meyering, dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D36951

db/wal_manager_test.cc

@@ -14,6 +14,7 @@
 #include "db/column_family.h"
 #include "db/version_set.h"
 #include "db/writebuffer.h"
+#include "util/mock_env.h"
 #include "util/string_util.h"
 #include "util/testharness.h"
 #include "util/testutil.h"
@@ -27,7 +28,7 @@ namespace rocksdb {
 class WalManagerTest : public testing::Test {
  public:
   WalManagerTest()
-      : env_(Env::Default()),
+      : env_(new MockEnv(Env::Default())),
         dbname_(test::TmpDir() + "/wal_manager_test"),
         table_cache_(NewLRUCache(50000, 16)),
         write_buffer_(db_options_.db_write_buffer_size),
@@ -41,6 +42,7 @@ class WalManagerTest : public testing::Test {
     db_options_.db_paths.emplace_back(dbname_,
                                       std::numeric_limits<uint64_t>::max());
     db_options_.wal_dir = dbname_;
+    db_options_.env = env_.get();
 
     versions_.reset(new VersionSet(dbname_, &db_options_, env_options_,
                                    table_cache_.get(), &write_buffer_,
@@ -91,7 +93,7 @@ class WalManagerTest : public testing::Test {
     return std::move(iter);
   }
 
-  Env* env_;
+  std::unique_ptr<MockEnv> env_;
   std::string dbname_;
   WriteController write_controller_;
   EnvOptions env_options_;
@@ -212,22 +214,22 @@ TEST_F(WalManagerTest, WALArchivalSizeLimit) {
   CreateArchiveLogs(20, 5000);
 
   std::vector<std::uint64_t> log_files =
-      ListSpecificFiles(env_, archive_dir, kLogFile);
+      ListSpecificFiles(env_.get(), archive_dir, kLogFile);
   ASSERT_EQ(log_files.size(), 20U);
 
   db_options_.WAL_size_limit_MB = 8;
   Reopen();
   wal_manager_->PurgeObsoleteWALFiles();
 
-  uint64_t archive_size = GetLogDirSize(archive_dir, env_);
+  uint64_t archive_size = GetLogDirSize(archive_dir, env_.get());
   ASSERT_TRUE(archive_size <= db_options_.WAL_size_limit_MB * 1024 * 1024);
 
   db_options_.WAL_ttl_seconds = 1;
-  env_->SleepForMicroseconds(2 * 1000 * 1000);
+  env_->FakeSleepForMicroseconds(2 * 1000 * 1000);
   Reopen();
   wal_manager_->PurgeObsoleteWALFiles();
 
-  log_files = ListSpecificFiles(env_, archive_dir, kLogFile);
+  log_files = ListSpecificFiles(env_.get(), archive_dir, kLogFile);
   ASSERT_TRUE(log_files.empty());
 }
 
@@ -245,15 +247,15 @@ TEST_F(WalManagerTest, WALArchivalTtl) {
   CreateArchiveLogs(20, 5000);
 
   std::vector<uint64_t> log_files =
-      ListSpecificFiles(env_, archive_dir, kLogFile);
+      ListSpecificFiles(env_.get(), archive_dir, kLogFile);
   ASSERT_GT(log_files.size(), 0U);
 
   db_options_.WAL_ttl_seconds = 1;
-  env_->SleepForMicroseconds(3 * 1000 * 1000);
+  env_->FakeSleepForMicroseconds(3 * 1000 * 1000);
   Reopen();
   wal_manager_->PurgeObsoleteWALFiles();
 
-  log_files = ListSpecificFiles(env_, archive_dir, kLogFile);
+  log_files = ListSpecificFiles(env_.get(), archive_dir, kLogFile);
   ASSERT_TRUE(log_files.empty());
 }
 

util/mock_env.cc

@@ -19,8 +19,9 @@ namespace rocksdb {
 
 class MemFile {
  public:
-  explicit MemFile(const std::string& fn, bool _is_lock_file = false)
-      : fn_(fn),
+  explicit MemFile(Env* env, const std::string& fn, bool _is_lock_file = false)
+      : env_(env),
+        fn_(fn),
         refs_(0),
         is_lock_file_(_is_lock_file),
         locked_(false),
@@ -137,8 +138,10 @@ class MemFile {
 
  private:
   uint64_t Now() {
-    return std::chrono::duration_cast<std::chrono::seconds>(
-      std::chrono::system_clock::now().time_since_epoch()).count();
+    int64_t unix_time;
+    auto s = env_->GetCurrentTime(&unix_time);
+    assert(s.ok());
+    return static_cast<uint64_t>(unix_time);
   }
 
   // Private since only Unref() should be used to delete it.
@@ -150,6 +153,7 @@ class MemFile {
   MemFile(const MemFile&);
   void operator=(const MemFile&);
 
+  Env* env_;
   const std::string fn_;
   mutable port::Mutex mutex_;
   int refs_;
@@ -393,8 +397,7 @@ class TestMemLogger : public Logger {
 
 }  // Anonymous namespace
 
-MockEnv::MockEnv(Env* base_env)
-  : EnvWrapper(base_env) {}
+MockEnv::MockEnv(Env* base_env) : EnvWrapper(base_env), fake_sleep_micros_(0) {}
 
 MockEnv::~MockEnv() {
   for (FileSystem::iterator i = file_map_.begin(); i != file_map_.end(); ++i) {
@@ -445,7 +448,7 @@ Status MockEnv::NewWritableFile(const std::string& fname,
   if (file_map_.find(fn) != file_map_.end()) {
     DeleteFileInternal(fn);
   }
-  MemFile* file = new MemFile(fn, false);
+  MemFile* file = new MemFile(this, fn, false);
   file->Ref();
   file_map_[fn] = file;
 
@@ -599,7 +602,7 @@ Status MockEnv::NewLogger(const std::string& fname,
   auto iter = file_map_.find(fn);
   MemFile* file = nullptr;
   if (iter == file_map_.end()) {
-    file = new MemFile(fn, false);
+    file = new MemFile(this, fn, false);
     file->Ref();
     file_map_[fn] = file;
   } else {
@@ -622,7 +625,7 @@ Status MockEnv::LockFile(const std::string& fname, FileLock** flock) {
         return Status::IOError(fn, "Lock is already held.");
       }
     } else {
-      auto* file = new MemFile(fn, true);
+      auto* file = new MemFile(this, fn, true);
       file->Ref();
       file->Lock();
       file_map_[fn] = file;
@@ -652,6 +655,20 @@ Status MockEnv::GetTestDirectory(std::string* path) {
   return Status::OK();
 }
 
+Status MockEnv::GetCurrentTime(int64_t* unix_time) {
+  auto s = EnvWrapper::GetCurrentTime(unix_time);
+  *unix_time += fake_sleep_micros_.load() / (1000 * 1000);
+  return s;
+}
+
+uint64_t MockEnv::NowMicros() {
+  return EnvWrapper::NowMicros() + fake_sleep_micros_.load();
+}
+
+uint64_t MockEnv::NowNanos() {
+  return EnvWrapper::NowNanos() + fake_sleep_micros_.load() * 1000;
+}
+
 // Non-virtual functions, specific to MockEnv
 Status MockEnv::Truncate(const std::string& fname, size_t size) {
   auto fn = NormalizePath(fname);
@@ -686,4 +703,8 @@ std::string MockEnv::NormalizePath(const std::string path) {
   return dst;
 }
 
+void MockEnv::FakeSleepForMicroseconds(int64_t micros) {
+  fake_sleep_micros_.fetch_add(micros);
+}
+
 }  // namespace rocksdb

util/mock_env.h

@@ -82,11 +82,20 @@ class MockEnv : public EnvWrapper {
 
   virtual Status GetTestDirectory(std::string* path) override;
 
+  // Results of these can be affected by FakeSleepForMicroseconds()
+  virtual Status GetCurrentTime(int64_t* unix_time) override;
+  virtual uint64_t NowMicros() override;
+  virtual uint64_t NowNanos() override;
+
   // Non-virtual functions, specific to MockEnv
   Status Truncate(const std::string& fname, size_t size);
 
   Status CorruptBuffer(const std::string& fname);
 
+  // Doesn't really sleep, just affects output of GetCurrentTime(), NowMicros()
+  // and NowNanos()
+  void FakeSleepForMicroseconds(int64_t micros);
+
  private:
   std::string NormalizePath(const std::string path);
 
@@ -94,6 +103,8 @@ class MockEnv : public EnvWrapper {
   typedef std::map<std::string, MemFile*> FileSystem;
   port::Mutex mutex_;
   FileSystem file_map_;  // Protected by mutex_.
+
+  std::atomic<int64_t> fake_sleep_micros_;
 };
 
 }  // namespace rocksdb

util/mock_env_test.cc

@@ -15,7 +15,7 @@ namespace rocksdb {
 
 class MockEnvTest : public testing::Test {
  public:
-  Env* env_;
+  MockEnv* env_;
   const EnvOptions soptions_;
 
   MockEnvTest()
@@ -264,6 +264,19 @@ TEST_F(MockEnvTest, DBTest) {
   delete db;
 }
 
+TEST_F(MockEnvTest, FakeSleeping) {
+  int64_t now = 0;
+  auto s = env_->GetCurrentTime(&now);
+  ASSERT_OK(s);
+  env_->FakeSleepForMicroseconds(3 * 1000 * 1000);
+  int64_t after_sleep = 0;
+  s = env_->GetCurrentTime(&after_sleep);
+  ASSERT_OK(s);
+  auto delta = after_sleep - now;
+  // this will be true unless test runs for 2 seconds
+  ASSERT_TRUE(delta == 3 || delta == 4);
+}
+
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {