Fix some bugs in MockEnv

Summary:
Fixing some bugs in MockEnv so it be actually used.
Closes https://github.com/facebook/rocksdb/pull/1914

Differential Revision: D4609923

Pulled By: maysamyabandeh

fbshipit-source-id: ca25735

db/db_bloom_filter_test.cc

@@ -606,6 +606,7 @@ class BloomStatsTestWithParam
       PlainTableOptions table_options;
       options_.table_factory.reset(NewPlainTableFactory(table_options));
     }
+    options_.env = env_;
 
     perf_context.Reset();
     DestroyAndReopen(options_);

db/db_dynamic_level_test.cc

@@ -75,6 +75,7 @@ TEST_F(DBTestDynamicLevel, DynamicLevelMaxBytesBase) {
       options.compression_per_level[0] = kNoCompression;
       options.compression_per_level[1] = kLZ4Compression;
       options.compression_per_level[2] = kSnappyCompression;
+      options.env = env_;
 
       DestroyAndReopen(options);
 

db/db_flush_test.cc

@@ -25,6 +25,7 @@ TEST_F(DBFlushTest, FlushWhileWritingManifest) {
   Options options;
   options.disable_auto_compactions = true;
   options.max_background_flushes = 2;
+  options.env = env_;
   Reopen(options);
   FlushOptions no_wait;
   no_wait.wait = false;
@@ -50,7 +51,7 @@ TEST_F(DBFlushTest, FlushWhileWritingManifest) {
 
 TEST_F(DBFlushTest, SyncFail) {
   std::unique_ptr<FaultInjectionTestEnv> fault_injection_env(
-      new FaultInjectionTestEnv(Env::Default()));
+      new FaultInjectionTestEnv(env_));
   Options options;
   options.disable_auto_compactions = true;
   options.env = fault_injection_env.get();

db/db_iterator_test.cc

@@ -981,6 +981,7 @@ TEST_F(DBIteratorTest, PrevAfterAndNextAfterMerge) {
   Options options;
   options.create_if_missing = true;
   options.merge_operator = MergeOperators::CreatePutOperator();
+  options.env = env_;
   DestroyAndReopen(options);
 
   // write three entries with different keys using Merge()

db/db_memtable_test.cc

@@ -123,6 +123,7 @@ TEST_F(DBMemTableTest, InsertWithHint) {
   options.memtable_factory.reset(new MockMemTableRepFactory());
   options.memtable_insert_with_hint_prefix_extractor.reset(
       new TestPrefixExtractor());
+  options.env = env_;
   Reopen(options);
   MockMemTableRep* rep =
       reinterpret_cast<MockMemTableRepFactory*>(options.memtable_factory.get())

db/db_merge_operator_test.cc

@@ -44,6 +44,7 @@ TEST_F(DBMergeOperatorTest, MergeErrorOnRead) {
   Options options;
   options.create_if_missing = true;
   options.merge_operator.reset(new TestPutOperator());
+  options.env = env_;
   Reopen(options);
   ASSERT_OK(Merge("k1", "v1"));
   ASSERT_OK(Merge("k1", "corrupted"));
@@ -57,6 +58,7 @@ TEST_F(DBMergeOperatorTest, MergeErrorOnWrite) {
   options.create_if_missing = true;
   options.merge_operator.reset(new TestPutOperator());
   options.max_successive_merges = 3;
+  options.env = env_;
   Reopen(options);
   ASSERT_OK(Merge("k1", "v1"));
   ASSERT_OK(Merge("k1", "v2"));
@@ -71,6 +73,7 @@ TEST_F(DBMergeOperatorTest, MergeErrorOnIteration) {
   Options options;
   options.create_if_missing = true;
   options.merge_operator.reset(new TestPutOperator());
+  options.env = env_;
 
   DestroyAndReopen(options);
   ASSERT_OK(Merge("k1", "v1"));

db/db_options_test.cc

@@ -62,6 +62,7 @@ class DBOptionsTest : public DBTestBase {
   std::unordered_map<std::string, std::string> GetRandomizedMutableCFOptionsMap(
       Random* rnd) {
     Options options;
+    options.env = env_;
     ImmutableDBOptions db_options(options);
     test::RandomInitCFOptions(&options, rnd);
     auto sanitized_options = SanitizeOptions(db_options, options);
@@ -87,6 +88,7 @@ TEST_F(DBOptionsTest, GetLatestDBOptions) {
   // GetOptions should be able to get latest option changed by SetOptions.
   Options options;
   options.create_if_missing = true;
+  options.env = env_;
   Random rnd(228);
   Reopen(options);
   auto new_options = GetRandomizedMutableDBOptionsMap(&rnd);
@@ -98,6 +100,7 @@ TEST_F(DBOptionsTest, GetLatestCFOptions) {
   // GetOptions should be able to get latest option changed by SetOptions.
   Options options;
   options.create_if_missing = true;
+  options.env = env_;
   Random rnd(228);
   Reopen(options);
   CreateColumnFamilies({"foo"}, options);
@@ -118,6 +121,7 @@ TEST_F(DBOptionsTest, SetOptionsAndReopen) {
   ASSERT_OK(dbfull()->SetOptions(rand_opts));
   // Verify if DB can be reopen after setting options.
   Options options;
+  options.env = env_;
   ASSERT_OK(TryReopen(options));
 }
 
@@ -137,6 +141,7 @@ TEST_F(DBOptionsTest, EnableAutoCompactionAndTriggerStall) {
           std::numeric_limits<uint64_t>::max();
       options.soft_pending_compaction_bytes_limit =
           std::numeric_limits<uint64_t>::max();
+      options.env = env_;
 
       DestroyAndReopen(options);
       int i = 0;
@@ -228,6 +233,7 @@ TEST_F(DBOptionsTest, SetOptionsMayTriggerCompaction) {
   Options options;
   options.create_if_missing = true;
   options.level0_file_num_compaction_trigger = 1000;
+  options.env = env_;
   Reopen(options);
   for (int i = 0; i < 3; i++) {
     // Need to insert two keys to avoid trivial move.
@@ -247,6 +253,7 @@ TEST_F(DBOptionsTest, SetBackgroundCompactionThreads) {
   options.create_if_missing = true;
   options.base_background_compactions = 1;  // default value
   options.max_background_compactions = 1;   // default value
+  options.env = env_;
   Reopen(options);
   ASSERT_EQ(1, dbfull()->TEST_BGCompactionsAllowed());
   ASSERT_OK(dbfull()->SetDBOptions({{"base_background_compactions", "2"},
@@ -260,6 +267,7 @@ TEST_F(DBOptionsTest, AvoidFlushDuringShutdown) {
   Options options;
   options.create_if_missing = true;
   options.disable_auto_compactions = true;
+  options.env = env_;
   WriteOptions write_without_wal;
   write_without_wal.disableWAL = true;
 
@@ -282,6 +290,7 @@ TEST_F(DBOptionsTest, SetDelayedWriteRateOption) {
   Options options;
   options.create_if_missing = true;
   options.delayed_write_rate = 2 * 1024U * 1024U;
+  options.env = env_;
   Reopen(options);
   ASSERT_EQ(2 * 1024U * 1024U, dbfull()->TEST_write_controler().max_delayed_write_rate());
 
@@ -297,6 +306,7 @@ TEST_F(DBOptionsTest, MaxTotalWalSizeChange) {
 
   Options options;
   options.create_if_missing = true;
+  options.env = env_;
   CreateColumnFamilies({"1", "2", "3"}, options);
   ReopenWithColumnFamilies({"default", "1", "2", "3"}, options);
 
@@ -326,7 +336,7 @@ static void assert_candidate_files_empty(DBImpl* dbfull, const bool empty) {
 }
 
 TEST_F(DBOptionsTest, DeleteObsoleteFilesPeriodChange) {
-  SpecialEnv env(Env::Default());
+  SpecialEnv env(env_);
   env.time_elapse_only_sleep_ = true;
   Options options;
   options.env = &env;

db/db_properties_test.cc

@@ -1156,6 +1156,7 @@ TEST_F(DBPropertiesTest, TablePropertiesNeedCompactTest) {
   options.max_bytes_for_level_multiplier = 4;
   options.soft_pending_compaction_bytes_limit = 1024 * 1024;
   options.num_levels = 8;
+  options.env = env_;
 
   std::shared_ptr<TablePropertiesCollectorFactory> collector_factory =
       std::make_shared<CountingDeleteTabPropCollectorFactory>();
@@ -1227,6 +1228,7 @@ TEST_F(DBPropertiesTest, NeedCompactHintPersistentTest) {
   options.level0_slowdown_writes_trigger = 10;
   options.level0_stop_writes_trigger = 10;
   options.disable_auto_compactions = true;
+  options.env = env_;
 
   std::shared_ptr<TablePropertiesCollectorFactory> collector_factory =
       std::make_shared<CountingDeleteTabPropCollectorFactory>();

db/db_table_properties_test.cc

@@ -147,6 +147,7 @@ TEST_F(DBTablePropertiesTest, GetPropertiesOfTablesInRange) {
   options.max_bytes_for_level_multiplier = 4;
   options.hard_pending_compaction_bytes_limit = 16 * 1024;
   options.num_levels = 8;
+  options.env = env_;
 
   DestroyAndReopen(options);
 

db/db_tailing_iter_test.cc

@@ -126,6 +126,7 @@ TEST_F(DBTestTailingIterator, TailingIteratorTrimSeekToNext) {
   options.write_buffer_size = k150KB;
   options.max_write_buffer_number = 3;
   options.min_write_buffer_number_to_merge = 2;
+  options.env = env_;
   CreateAndReopenWithCF({"pikachu"}, options);
   ReadOptions read_options;
   read_options.tailing = true;
@@ -298,7 +299,6 @@ TEST_F(DBTestTailingIterator, TailingIteratorPrefixSeek) {
   read_options.tailing = true;
 
   Options options = CurrentOptions();
-  options.env = env_;
   options.create_if_missing = true;
   options.disable_auto_compactions = true;
   options.prefix_extractor.reset(NewFixedPrefixTransform(2));
@@ -616,7 +616,6 @@ TEST_F(DBTestTailingIterator, ManagedTailingIteratorPrefixSeek) {
   read_options.managed = true;
 
   Options options = CurrentOptions();
-  options.env = env_;
   options.create_if_missing = true;
   options.disable_auto_compactions = true;
   options.prefix_extractor.reset(NewFixedPrefixTransform(2));

db/db_test_util.cc

@@ -52,6 +52,7 @@ DBTestBase::DBTestBase(const std::string path)
   alternative_wal_dir_ = dbname_ + "/wal";
   alternative_db_log_dir_ = dbname_ + "/db_log_dir";
   auto options = CurrentOptions();
+  options.env = env_;
   auto delete_options = options;
   delete_options.wal_dir = alternative_wal_dir_;
   EXPECT_OK(DestroyDB(dbname_, delete_options));
@@ -72,6 +73,7 @@ DBTestBase::~DBTestBase() {
   options.db_paths.emplace_back(dbname_ + "_2", 0);
   options.db_paths.emplace_back(dbname_ + "_3", 0);
   options.db_paths.emplace_back(dbname_ + "_4", 0);
+  options.env = env_;
 
   if (getenv("KEEP_DB")) {
     printf("DB is still at %s\n", dbname_.c_str());

db/db_test_util.h

@@ -14,9 +14,6 @@
 
 #include <fcntl.h>
 #include <inttypes.h>
-#ifndef OS_WIN
-#include <unistd.h>
-#endif
 
 #include <algorithm>
 #include <map>

util/mock_env.cc

@@ -116,6 +116,17 @@ class MemFile {
     return Status::OK();
   }
 
+  Status Write(uint64_t offset, const Slice& data) {
+    MutexLock lock(&mutex_);
+    if (offset + data.size() > data_.size()) {
+      data_.resize(offset + data.size());
+    }
+    data_.replace(offset, data.size(), data.data(), data.size());
+    size_ = data_.size();
+    modified_time_ = Now();
+    return Status::OK();
+  }
+
   Status Append(const Slice& data) {
     MutexLock lock(&mutex_);
     data_.append(data.data(), data.size());
@@ -223,6 +234,31 @@ class MockRandomAccessFile : public RandomAccessFile {
   MemFile* file_;
 };
 
+class MockRandomRWFile : public RandomRWFile {
+ public:
+  explicit MockRandomRWFile(MemFile* file) : file_(file) { file_->Ref(); }
+
+  ~MockRandomRWFile() { file_->Unref(); }
+
+  virtual Status Write(uint64_t offset, const Slice& data) override {
+    return file_->Write(offset, data);
+  }
+
+  virtual Status Read(uint64_t offset, size_t n, Slice* result,
+                      char* scratch) const override {
+    return file_->Read(offset, n, result, scratch);
+  }
+
+  virtual Status Close() override { return file_->Fsync(); }
+
+  virtual Status Flush() override { return Status::OK(); }
+
+  virtual Status Sync() override { return file_->Fsync(); }
+
+ private:
+  MemFile* file_;
+};
+
 class MockWritableFile : public WritableFile {
  public:
   MockWritableFile(MemFile* file, RateLimiter* rate_limiter)
@@ -248,6 +284,7 @@ class MockWritableFile : public WritableFile {
     return Status::OK();
   }
   virtual Status Truncate(uint64_t size) override {
+    file_->Truncate(size);
     return Status::OK();
   }
   virtual Status Close() override { return file_->Fsync(); }
@@ -439,6 +476,35 @@ Status MockEnv::NewRandomAccessFile(const std::string& fname,
   return Status::OK();
 }
 
+Status MockEnv::NewRandomRWFile(const std::string& fname,
+                                unique_ptr<RandomRWFile>* result,
+                                const EnvOptions& soptions) {
+  auto fn = NormalizePath(fname);
+  MutexLock lock(&mutex_);
+  if (file_map_.find(fn) == file_map_.end()) {
+    *result = NULL;
+    return Status::IOError(fn, "File not found");
+  }
+  auto* f = file_map_[fn];
+  if (f->is_lock_file()) {
+    return Status::InvalidArgument(fn, "Cannot open a lock file.");
+  }
+  result->reset(new MockRandomRWFile(f));
+  return Status::OK();
+}
+
+Status MockEnv::ReuseWritableFile(const std::string& fname,
+                                  const std::string& old_fname,
+                                  unique_ptr<WritableFile>* result,
+                                  const EnvOptions& options) {
+  auto s = RenameFile(old_fname, fname);
+  if (!s.ok()) {
+    return s;
+  }
+  result->reset();
+  return NewWritableFile(fname, result, options);
+}
+
 Status MockEnv::NewWritableFile(const std::string& fname,
                                 unique_ptr<WritableFile>* result,
                                 const EnvOptions& env_options) {
@@ -598,6 +664,7 @@ Status MockEnv::LinkFile(const std::string& src, const std::string& dest) {
 
   DeleteFileInternal(t);
   file_map_[t] = file_map_[s];
+  file_map_[t]->Ref();  // Otherwise it might get deleted when noone uses s
   return Status::OK();
 }
 

util/mock_env.h

@@ -35,6 +35,15 @@ class MockEnv : public EnvWrapper {
                                      unique_ptr<RandomAccessFile>* result,
                                      const EnvOptions& soptions) override;
 
+  virtual Status NewRandomRWFile(const std::string& fname,
+                                 unique_ptr<RandomRWFile>* result,
+                                 const EnvOptions& options) override;
+
+  virtual Status ReuseWritableFile(const std::string& fname,
+                                   const std::string& old_fname,
+                                   unique_ptr<WritableFile>* result,
+                                   const EnvOptions& options) override;
+
   virtual Status NewWritableFile(const std::string& fname,
                                  unique_ptr<WritableFile>* result,
                                  const EnvOptions& env_options) override;