From 72cb7cf201e17399226c111855eda567aaf26394 Mon Sep 17 00:00:00 2001
From: Lei Jin <lei@fb.com>
Date: Fri, 31 Oct 2014 15:16:31 -0700
Subject: [PATCH] Add fsync / corrupt simulation to env_mem

Summary: as title

Test Plan: env_mem_test

Reviewers: sdong, yhchiang, rven, igor

Reviewed By: igor

Subscribers: dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D28077
---
 util/mock_env.cc      | 39 ++++++++++++++++++++++++++++++++++++---
 util/mock_env.h       |  2 ++
 util/mock_env_test.cc | 39 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/util/mock_env.cc b/util/mock_env.cc
index 32c202beb..c44592314 100644
--- a/util/mock_env.cc
+++ b/util/mock_env.cc
@@ -12,13 +12,16 @@
 #include <algorithm>
 #include <chrono>
 #include "util/rate_limiter.h"
+#include "util/random.h"
+#include "util/murmurhash.h"
 
 namespace rocksdb {
 
 class MemFile {
  public:
   explicit MemFile(const std::string& fn) :
-    fn_(fn), refs_(0), size_(0), modified_time_(Now()) {}
+    fn_(fn), refs_(0), size_(0), modified_time_(Now()),
+    rnd_((uint32_t)MurmurHash(fn.data(), fn.size(), 0)), fsynced_bytes_(0) {}
 
   void Ref() {
     MutexLock lock(&mutex_);
@@ -53,6 +56,19 @@ class MemFile {
     }
   }
 
+  void CorruptBuffer() {
+    if (fsynced_bytes_ >= size_) {
+      return;
+    }
+    uint64_t buffered_bytes = size_ - fsynced_bytes_;
+    uint64_t start = fsynced_bytes_ + rnd_.Uniform(buffered_bytes);
+    uint64_t end = std::min(start + 512, size_.load());
+    MutexLock lock(&mutex_);
+    for (uint64_t pos = start; pos < end; ++pos) {
+      data_[pos] = static_cast<char>(rnd_.Uniform(256));
+    }
+  }
+
   Status Read(uint64_t offset, size_t n, Slice* result, char* scratch) const {
     MutexLock lock(&mutex_);
     if (offset > Size()) {
@@ -84,6 +100,7 @@ class MemFile {
   }
 
   Status Fsync() {
+    fsynced_bytes_ = size_.load();
     return Status::OK();
   }
 
@@ -110,9 +127,14 @@ class MemFile {
   mutable port::Mutex mutex_;
   int refs_;
 
+  // Data written into this file, all bytes before fsynced_bytes are
+  // persistent.
   std::string data_;
   std::atomic<uint64_t> size_;
   std::atomic<uint64_t> modified_time_;
+
+  Random rnd_;
+  std::atomic<uint64_t> fsynced_bytes_;
 };
 
 namespace {
@@ -197,7 +219,7 @@ class WritableFileImpl : public WritableFile {
   }
 
   virtual Status Close() {
-    return Status::OK();
+    return file_->Fsync();
   }
 
   virtual Status Flush() {
@@ -581,7 +603,7 @@ Status MockEnv::GetTestDirectory(std::string* path) {
   return Status::OK();
 }
 
-  // Non-virtual functions, specific to MockEnv
+// Non-virtual functions, specific to MockEnv
 Status MockEnv::Truncate(const std::string& fname, size_t size) {
   auto fn = NormalizePath(fname);
   MutexLock lock(&mutex_);
@@ -593,6 +615,17 @@ Status MockEnv::Truncate(const std::string& fname, size_t size) {
   return Status::OK();
 }
 
+Status MockEnv::CorruptBuffer(const std::string& fname) {
+  auto fn = NormalizePath(fname);
+  MutexLock lock(&mutex_);
+  auto iter = file_map_.find(fn);
+  if (iter == file_map_.end()) {
+    return Status::IOError(fn, "File not found");
+  }
+  iter->second->CorruptBuffer();
+  return Status::OK();
+}
+
 std::string MockEnv::NormalizePath(const std::string path) {
   std::string dst;
   for (auto c : path) {
diff --git a/util/mock_env.h b/util/mock_env.h
index d128c75b0..b92caa5cf 100644
--- a/util/mock_env.h
+++ b/util/mock_env.h
@@ -81,6 +81,8 @@ class MockEnv : public EnvWrapper {
   // Non-virtual functions, specific to MockEnv
   Status Truncate(const std::string& fname, size_t size);
 
+  Status CorruptBuffer(const std::string& fname);
+
  private:
   std::string NormalizePath(const std::string path);
 
diff --git a/util/mock_env_test.cc b/util/mock_env_test.cc
index 51ae8e296..521f0fb1c 100644
--- a/util/mock_env_test.cc
+++ b/util/mock_env_test.cc
@@ -182,6 +182,45 @@ TEST(MockEnvTest, LargeWrite) {
   delete [] scratch;
 }
 
+TEST(MockEnvTest, Corrupt) {
+  const std::string kGood = "this is a good string, synced to disk";
+  const std::string kCorrupted = "this part may be corrupted";
+  const std::string kFileName = "/dir/f";
+  unique_ptr<WritableFile> writable_file;
+  ASSERT_OK(env_->NewWritableFile(kFileName, &writable_file, soptions_));
+  ASSERT_OK(writable_file->Append(kGood));
+  ASSERT_TRUE(writable_file->GetFileSize() == kGood.size());
+
+  std::string scratch;
+  scratch.resize(kGood.size() + kCorrupted.size() + 16);
+  Slice result;
+  unique_ptr<RandomAccessFile> rand_file;
+  ASSERT_OK(env_->NewRandomAccessFile(kFileName, &rand_file, soptions_));
+  ASSERT_OK(rand_file->Read(0, kGood.size(), &result, &(scratch[0])));
+  ASSERT_EQ(result.compare(kGood), 0);
+
+  // Sync + corrupt => no change
+  ASSERT_OK(writable_file->Fsync());
+  ASSERT_OK(dynamic_cast<MockEnv*>(env_)->CorruptBuffer(kFileName));
+  result.clear();
+  ASSERT_OK(rand_file->Read(0, kGood.size(), &result, &(scratch[0])));
+  ASSERT_EQ(result.compare(kGood), 0);
+
+  // Add new data and corrupt it
+  ASSERT_OK(writable_file->Append(kCorrupted));
+  ASSERT_TRUE(writable_file->GetFileSize() == kGood.size() + kCorrupted.size());
+  result.clear();
+  ASSERT_OK(rand_file->Read(kGood.size(), kCorrupted.size(),
+            &result, &(scratch[0])));
+  ASSERT_EQ(result.compare(kCorrupted), 0);
+  // Corrupted
+  ASSERT_OK(dynamic_cast<MockEnv*>(env_)->CorruptBuffer(kFileName));
+  result.clear();
+  ASSERT_OK(rand_file->Read(kGood.size(), kCorrupted.size(),
+            &result, &(scratch[0])));
+  ASSERT_NE(result.compare(kCorrupted), 0);
+}
+
 TEST(MockEnvTest, DBTest) {
   Options options;
   options.create_if_missing = true;