Refactor expected state in stress/crash test (#8913)

Summary:
This is a precursor refactoring to enable an upcoming feature: persistence failure correctness testing.

- Changed `--expected_values_path` to `--expected_values_dir` and migrated "db_crashtest.py" to use the new flag. For persistence failure correctness testing there are multiple possible correct states since unsynced data is allowed to be dropped. Making it possible to restore all these possible correct states will eventually involve files containing snapshots of expected values and DB trace files.
- The expected values directory is managed by an `ExpectedStateManager` instance. Managing expected state files is separated out of `SharedState` to prevent `SharedState` from becoming too complex when the new files and features (snapshotting, tracing, and restoring) are introduced.
- Migrated expected values file access/management out of `SharedState` into a separate class called `ExpectedState`. This is not exposed directly to the test but rather the `ExpectedState` for the latest values file is accessed via a pass-through API on `ExpectedStateManager`. This forces the test to always access the single latest `ExpectedState`.
- Changed the initialization of the latest expected values file to use a tempfile followed by rename, and also add cleanup logic for possible stranded tempfiles.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/8913

Test Plan:
run in several ways; try to make sure it's not obviously broken.

- crashtest blackbox without TEST_TMPDIR
```
$ python3 tools/db_crashtest.py blackbox --simple --write_buffer_size=1048576 --target_file_size_base=1048576 --max_bytes_for_level_base=4194304 --max_key=100000 --value_size_mult=33 --compression_type=none --duration=120 --interval=10 --compression_type=none --blob_compression_type=none
```
- crashtest blackbox with TEST_TMPDIR
```
$ TEST_TMPDIR=/dev/shm python3 tools/db_crashtest.py blackbox --simple --write_buffer_size=1048576 --target_file_size_base=1048576 --max_bytes_for_level_base=4194304 --max_key=100000 --value_size_mult=33 --compression_type=none --duration=120 --interval=10 --compression_type=none --blob_compression_type=none
```
- crashtest whitebox with TEST_TMPDIR
```
$ TEST_TMPDIR=/dev/shm python3 tools/db_crashtest.py whitebox --simple --write_buffer_size=1048576 --target_file_size_base=1048576 --max_bytes_for_level_base=4194304 --max_key=100000 --value_size_mult=33 --compression_type=none --duration=120 --interval=10 --compression_type=none --blob_compression_type=none --random_kill_odd=88887
```
- db_stress without expected_values_dir
```
$ ./db_stress --write_buffer_size=1048576 --target_file_size_base=1048576 --max_bytes_for_level_base=4194304 --max_key=100000 --value_size_mult=33 --compression_type=none --ops_per_thread=10000 --clear_column_family_one_in=0 --destroy_db_initially=true
```
- db_stress with expected_values_dir and manual corruption
```
$ ./db_stress --write_buffer_size=1048576 --target_file_size_base=1048576 --max_bytes_for_level_base=4194304 --max_key=100000 --value_size_mult=33 --compression_type=none --ops_per_thread=10000 --clear_column_family_one_in=0 --destroy_db_initially=true --expected_values_dir=./
// modify one byte in "./LATEST.state"
$ ./db_stress --write_buffer_size=1048576 --target_file_size_base=1048576 --max_bytes_for_level_base=4194304 --max_key=100000 --value_size_mult=33 --compression_type=none --ops_per_thread=10000 --clear_column_family_one_in=0 --destroy_db_initially=false --expected_values_dir=./
...
Verification failed for column family 0 key 0000000000000000 (0): Value not found: NotFound:
...
```

Reviewed By: riversand963

Differential Revision: D30921951

Pulled By: ajkr

fbshipit-source-id: babfe218062e55d018c9b046536c0289fb78f41c

TARGETS

@@ -847,6 +847,7 @@ cpp_library(
         "db_stress_tool/db_stress_shared_state.cc",
         "db_stress_tool/db_stress_test_base.cc",
         "db_stress_tool/db_stress_tool.cc",
+        "db_stress_tool/expected_state.cc",
         "db_stress_tool/no_batched_ops_stress.cc",
         "test_util/testutil.cc",
         "tools/block_cache_analyzer/block_cache_trace_analyzer.cc",

db_stress_tool/CMakeLists.txt

@@ -9,6 +9,7 @@ add_executable(db_stress${ARTIFACT_SUFFIX}
   db_stress_shared_state.cc
   db_stress_gflags.cc
   db_stress_tool.cc
+  expected_state.cc
   no_batched_ops_stress.cc)
 target_link_libraries(db_stress${ARTIFACT_SUFFIX} ${ROCKSDB_LIB} ${THIRDPARTY_LIBS})
 list(APPEND tool_deps db_stress)

db_stress_tool/db_stress_common.h

@@ -153,7 +153,7 @@ DECLARE_int32(index_type);
 DECLARE_string(db);
 DECLARE_string(secondaries_base);
 DECLARE_bool(test_secondary);
-DECLARE_string(expected_values_path);
+DECLARE_string(expected_values_dir);
 DECLARE_bool(verify_checksum);
 DECLARE_bool(mmap_read);
 DECLARE_bool(mmap_write);

db_stress_tool/db_stress_gflags.cc

@@ -448,12 +448,12 @@ DEFINE_string(secondaries_base, "",
 DEFINE_bool(test_secondary, false, "Test secondary instance.");
 
 DEFINE_string(
-    expected_values_path, "",
+    expected_values_dir, "",
-    "File where the array of expected uint32_t values will be stored. If "
+    "Dir where file with array of expected uint32_t values will be stored. If "
     "provided and non-empty, the DB state will be verified against these "
     "values after recovery. --max_key and --column_family must be kept the "
     "same across invocations of this program that use the same "
-    "--expected_values_path.");
+    "--expected_values_dir.");
 
 DEFINE_bool(verify_checksum, false,
             "Verify checksum for every block read from storage");

db_stress_tool/db_stress_shared_state.h

@@ -11,6 +11,7 @@
 #pragma once
 
 #include "db_stress_tool/db_stress_stat.h"
+#include "db_stress_tool/expected_state.h"
 // SyncPoint is not supported in Released Windows Mode.
 #if !(defined NDEBUG) || !defined(OS_WIN)
 #include "test_util/sync_point.h"
@@ -23,7 +24,7 @@ DECLARE_uint64(log2_keys_per_lock);
 DECLARE_int32(threads);
 DECLARE_int32(column_families);
 DECLARE_int32(nooverwritepercent);
-DECLARE_string(expected_values_path);
+DECLARE_string(expected_values_dir);
 DECLARE_int32(clear_column_family_one_in);
 DECLARE_bool(test_batches_snapshots);
 DECLARE_int32(compaction_thread_pool_adjust_interval);
@@ -81,7 +82,7 @@ class SharedState {
         verification_failure_(false),
         should_stop_test_(false),
         no_overwrite_ids_(FLAGS_column_families),
-        values_(nullptr),
+        expected_state_manager_(nullptr),
         printing_verification_results_(false) {
     // Pick random keys in each column family that will not experience
     // overwrite
@@ -110,65 +111,38 @@ class SharedState {
     }
     delete[] permutation;
 
-    size_t expected_values_size =
-        sizeof(std::atomic<uint32_t>) * FLAGS_column_families * max_key_;
-    bool values_init_needed = false;
     Status status;
-    if (!FLAGS_expected_values_path.empty()) {
+    // TODO: We should introduce a way to explicitly disable verification
+    // during shutdown. When that is disabled and FLAGS_expected_values_dir
+    // is empty (disabling verification at startup), we can skip tracking
+    // expected state. Only then should we permit bypassing the below feature
+    // compatibility checks.
+    if (!FLAGS_expected_values_dir.empty()) {
       if (!std::atomic<uint32_t>{}.is_lock_free()) {
         status = Status::InvalidArgument(
-            "Cannot use --expected_values_path on platforms without lock-free "
+            "Cannot use --expected_values_dir on platforms without lock-free "
             "std::atomic<uint32_t>");
       }
       if (status.ok() && FLAGS_clear_column_family_one_in > 0) {
         status = Status::InvalidArgument(
-            "Cannot use --expected_values_path on when "
+            "Cannot use --expected_values_dir on when "
             "--clear_column_family_one_in is greater than zero.");
       }
-      uint64_t size = 0;
+    }
-      Env* default_env = Env::Default();
+    if (status.ok()) {
-      if (status.ok()) {
+      if (FLAGS_expected_values_dir.empty()) {
-        status = default_env->GetFileSize(FLAGS_expected_values_path, &size);
+        expected_state_manager_.reset(
-      }
+            new AnonExpectedStateManager(FLAGS_max_key, FLAGS_column_families));
-      std::unique_ptr<WritableFile> wfile;
-      if (status.ok() && size == 0) {
-        const EnvOptions soptions;
-        status = default_env->NewWritableFile(FLAGS_expected_values_path,
-                                              &wfile, soptions);
-      }
-      if (status.ok() && size == 0) {
-        std::string buf(expected_values_size, '\0');
-        status = wfile->Append(buf);
-        values_init_needed = true;
-      }
-      if (status.ok()) {
-        status = default_env->NewMemoryMappedFileBuffer(
-            FLAGS_expected_values_path, &expected_mmap_buffer_);
-      }
-      if (status.ok()) {
-        assert(expected_mmap_buffer_->GetLen() == expected_values_size);
-        values_ = static_cast<std::atomic<uint32_t>*>(
-            expected_mmap_buffer_->GetBase());
-        assert(values_ != nullptr);
       } else {
-        fprintf(stderr, "Failed opening shared file '%s' with error: %s\n",
+        expected_state_manager_.reset(new FileExpectedStateManager(
-                FLAGS_expected_values_path.c_str(), status.ToString().c_str());
+            FLAGS_max_key, FLAGS_column_families, FLAGS_expected_values_dir));
-        assert(values_ == nullptr);
       }
+      status = expected_state_manager_->Open();
     }
-    if (values_ == nullptr) {
+    if (!status.ok()) {
-      values_allocation_.reset(
+      fprintf(stderr, "Failed setting up expected state with error: %s\n",
-          new std::atomic<uint32_t>[FLAGS_column_families * max_key_]);
+              status.ToString().c_str());
-      values_ = &values_allocation_[0];
+      exit(1);
-      values_init_needed = true;
-    }
-    assert(values_ != nullptr);
-    if (values_init_needed) {
-      for (int i = 0; i < FLAGS_column_families; ++i) {
-        for (int j = 0; j < max_key_; ++j) {
-          Delete(i, j, false /* pending */);
-        }
-      }
     }
 
     if (FLAGS_test_batches_snapshots) {
@@ -257,89 +231,76 @@ class SharedState {
 
   bool ShouldStopTest() const { return should_stop_test_.load(); }
 
+  // Returns a lock covering `key` in `cf`.
   port::Mutex* GetMutexForKey(int cf, int64_t key) {
     return key_locks_[cf][key >> log2_keys_per_lock_].get();
   }
 
+  // Acquires locks for all keys in `cf`.
   void LockColumnFamily(int cf) {
     for (auto& mutex : key_locks_[cf]) {
       mutex->Lock();
     }
   }
 
+  // Releases locks for all keys in `cf`.
   void UnlockColumnFamily(int cf) {
     for (auto& mutex : key_locks_[cf]) {
       mutex->Unlock();
     }
   }
 
-  std::atomic<uint32_t>& Value(int cf, int64_t key) const {
+  // Requires external locking covering all keys in `cf`.
-    return values_[cf * max_key_ + key];
-  }
-
   void ClearColumnFamily(int cf) {
-    std::fill(&Value(cf, 0 /* key */), &Value(cf + 1, 0 /* key */),
+    return expected_state_manager_->ClearColumnFamily(cf);
-              DELETION_SENTINEL);
   }
 
   // @param pending True if the update may have started but is not yet
   //    guaranteed finished. This is useful for crash-recovery testing when the
   //    process may crash before updating the expected values array.
+  //
+  // Requires external locking covering `key` in `cf`.
   void Put(int cf, int64_t key, uint32_t value_base, bool pending) {
-    if (!pending) {
+    return expected_state_manager_->Put(cf, key, value_base, pending);
-      // prevent expected-value update from reordering before Write
-      std::atomic_thread_fence(std::memory_order_release);
-    }
-    Value(cf, key).store(pending ? UNKNOWN_SENTINEL : value_base,
-                         std::memory_order_relaxed);
-    if (pending) {
-      // prevent Write from reordering before expected-value update
-      std::atomic_thread_fence(std::memory_order_release);
-    }
   }
 
-  uint32_t Get(int cf, int64_t key) const { return Value(cf, key); }
+  // Requires external locking covering `key` in `cf`.
+  uint32_t Get(int cf, int64_t key) const {
+    return expected_state_manager_->Get(cf, key);
+  }
 
   // @param pending See comment above Put()
   // Returns true if the key was not yet deleted.
+  //
+  // Requires external locking covering `key` in `cf`.
   bool Delete(int cf, int64_t key, bool pending) {
-    if (Value(cf, key) == DELETION_SENTINEL) {
+    return expected_state_manager_->Delete(cf, key, pending);
-      return false;
-    }
-    Put(cf, key, DELETION_SENTINEL, pending);
-    return true;
   }
 
   // @param pending See comment above Put()
   // Returns true if the key was not yet deleted.
+  //
+  // Requires external locking covering `key` in `cf`.
   bool SingleDelete(int cf, int64_t key, bool pending) {
-    return Delete(cf, key, pending);
+    return expected_state_manager_->Delete(cf, key, pending);
   }
 
   // @param pending See comment above Put()
   // Returns number of keys deleted by the call.
+  //
+  // Requires external locking covering keys in `[begin_key, end_key)` in `cf`.
   int DeleteRange(int cf, int64_t begin_key, int64_t end_key, bool pending) {
-    int covered = 0;
+    return expected_state_manager_->DeleteRange(cf, begin_key, end_key,
-    for (int64_t key = begin_key; key < end_key; ++key) {
+                                                pending);
-      if (Delete(cf, key, pending)) {
-        ++covered;
-      }
-    }
-    return covered;
   }
 
   bool AllowsOverwrite(int64_t key) {
     return no_overwrite_ids_.find(key) == no_overwrite_ids_.end();
   }
 
+  // Requires external locking covering `key` in `cf`.
   bool Exists(int cf, int64_t key) {
-    // UNKNOWN_SENTINEL counts as exists. That assures a key for which overwrite
+    return expected_state_manager_->Exists(cf, key);
-    // is disallowed can't be accidentally added a second time, in which case
-    // SingleDelete wouldn't be able to properly delete the key. It does allow
-    // the case where a SingleDelete might be added which covers nothing, but
-    // that's not a correctness issue.
-    uint32_t expected_value = Value(cf, key).load();
-    return expected_value != DELETION_SENTINEL;
   }
 
   uint32_t GetSeed() const { return seed_; }
@@ -355,7 +316,7 @@ class SharedState {
   }
 
   bool ShouldVerifyAtBeginning() const {
-    return expected_mmap_buffer_.get() != nullptr;
+    return !FLAGS_expected_values_dir.empty();
   }
 
   bool PrintingVerificationResults() {
@@ -395,12 +356,10 @@ class SharedState {
   // Keys that should not be overwritten
   std::unordered_set<size_t> no_overwrite_ids_;
 
-  std::atomic<uint32_t>* values_;
+  std::unique_ptr<ExpectedStateManager> expected_state_manager_;
-  std::unique_ptr<std::atomic<uint32_t>[]> values_allocation_;
   // Has to make it owned by a smart ptr as port::Mutex is not copyable
   // and storing it in the container may require copying depending on the impl.
   std::vector<std::vector<std::unique_ptr<port::Mutex>>> key_locks_;
-  std::unique_ptr<MemoryMappedFileBuffer> expected_mmap_buffer_;
   std::atomic<bool> printing_verification_results_;
 };
 

db_stress_tool/expected_state.cc

@@ -0,0 +1,244 @@
+//  Copyright (c) 2021-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).
+
+#ifdef GFLAGS
+
+#include "db_stress_tool/expected_state.h"
+
+#include "db_stress_tool/db_stress_shared_state.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+ExpectedState::ExpectedState(size_t max_key, size_t num_column_families)
+    : max_key_(max_key),
+      num_column_families_(num_column_families),
+      values_(nullptr) {}
+
+void ExpectedState::ClearColumnFamily(int cf) {
+  std::fill(&Value(cf, 0 /* key */), &Value(cf + 1, 0 /* key */),
+            SharedState::DELETION_SENTINEL);
+}
+
+void ExpectedState::Put(int cf, int64_t key, uint32_t value_base,
+                        bool pending) {
+  if (!pending) {
+    // prevent expected-value update from reordering before Write
+    std::atomic_thread_fence(std::memory_order_release);
+  }
+  Value(cf, key).store(pending ? SharedState::UNKNOWN_SENTINEL : value_base,
+                       std::memory_order_relaxed);
+  if (pending) {
+    // prevent Write from reordering before expected-value update
+    std::atomic_thread_fence(std::memory_order_release);
+  }
+}
+
+uint32_t ExpectedState::Get(int cf, int64_t key) const {
+  return Value(cf, key);
+}
+
+bool ExpectedState::Delete(int cf, int64_t key, bool pending) {
+  if (Value(cf, key) == SharedState::DELETION_SENTINEL) {
+    return false;
+  }
+  Put(cf, key, SharedState::DELETION_SENTINEL, pending);
+  return true;
+}
+
+bool ExpectedState::SingleDelete(int cf, int64_t key, bool pending) {
+  return Delete(cf, key, pending);
+}
+
+int ExpectedState::DeleteRange(int cf, int64_t begin_key, int64_t end_key,
+                               bool pending) {
+  int covered = 0;
+  for (int64_t key = begin_key; key < end_key; ++key) {
+    if (Delete(cf, key, pending)) {
+      ++covered;
+    }
+  }
+  return covered;
+}
+
+bool ExpectedState::Exists(int cf, int64_t key) {
+  // UNKNOWN_SENTINEL counts as exists. That assures a key for which overwrite
+  // is disallowed can't be accidentally added a second time, in which case
+  // SingleDelete wouldn't be able to properly delete the key. It does allow
+  // the case where a SingleDelete might be added which covers nothing, but
+  // that's not a correctness issue.
+  uint32_t expected_value = Value(cf, key).load();
+  return expected_value != SharedState::DELETION_SENTINEL;
+}
+
+void ExpectedState::Reset() {
+  for (size_t i = 0; i < num_column_families_; ++i) {
+    for (size_t j = 0; j < max_key_; ++j) {
+      Delete(static_cast<int>(i), j, false /* pending */);
+    }
+  }
+}
+
+FileExpectedState::FileExpectedState(std::string expected_state_file_path,
+                                     size_t max_key, size_t num_column_families)
+    : ExpectedState(max_key, num_column_families),
+      expected_state_file_path_(expected_state_file_path) {}
+
+Status FileExpectedState::Open(bool create) {
+  size_t expected_values_size = GetValuesLen();
+
+  Env* default_env = Env::Default();
+
+  Status status;
+  if (create) {
+    std::unique_ptr<WritableFile> wfile;
+    const EnvOptions soptions;
+    status = default_env->NewWritableFile(expected_state_file_path_, &wfile,
+                                          soptions);
+    if (status.ok()) {
+      std::string buf(expected_values_size, '\0');
+      status = wfile->Append(buf);
+    }
+  }
+  if (status.ok()) {
+    status = default_env->NewMemoryMappedFileBuffer(
+        expected_state_file_path_, &expected_state_mmap_buffer_);
+  }
+  if (status.ok()) {
+    assert(expected_state_mmap_buffer_->GetLen() == expected_values_size);
+    values_ = static_cast<std::atomic<uint32_t>*>(
+        expected_state_mmap_buffer_->GetBase());
+    assert(values_ != nullptr);
+    if (create) {
+      Reset();
+    }
+  } else {
+    assert(values_ == nullptr);
+  }
+  return status;
+}
+
+AnonExpectedState::AnonExpectedState(size_t max_key, size_t num_column_families)
+    : ExpectedState(max_key, num_column_families) {}
+
+#ifndef NDEBUG
+Status AnonExpectedState::Open(bool create) {
+#else
+Status AnonExpectedState::Open(bool /* create */) {
+#endif
+  // AnonExpectedState only supports being freshly created.
+  assert(create);
+  values_allocation_.reset(
+      new std::atomic<uint32_t>[GetValuesLen() /
+                                sizeof(std::atomic<uint32_t>)]);
+  values_ = &values_allocation_[0];
+  Reset();
+  return Status::OK();
+}
+
+ExpectedStateManager::ExpectedStateManager(size_t max_key,
+                                           size_t num_column_families)
+    : max_key_(max_key),
+      num_column_families_(num_column_families),
+      latest_(nullptr) {}
+
+ExpectedStateManager::~ExpectedStateManager() {}
+
+const std::string FileExpectedStateManager::kLatestFilename = "LATEST.state";
+
+FileExpectedStateManager::FileExpectedStateManager(
+    size_t max_key, size_t num_column_families,
+    std::string expected_state_dir_path)
+    : ExpectedStateManager(max_key, num_column_families),
+      expected_state_dir_path_(std::move(expected_state_dir_path)) {
+  assert(!expected_state_dir_path_.empty());
+}
+
+Status FileExpectedStateManager::Open() {
+  Status s = Clean();
+
+  std::string expected_state_file_path = GetPathForFilename(kLatestFilename);
+  bool found = false;
+  if (s.ok()) {
+    Status exists_status = Env::Default()->FileExists(expected_state_file_path);
+    if (exists_status.ok()) {
+      found = true;
+    } else if (exists_status.IsNotFound()) {
+      found = false;
+    } else {
+      s = exists_status;
+    }
+  }
+
+  if (!found) {
+    // Initialize the file in a temp path and then rename it. That way, in case
+    // this process is killed during setup, `Clean()` will take care of removing
+    // the incomplete expected values file.
+    std::string temp_expected_state_file_path =
+        GetTempPathForFilename(kLatestFilename);
+    FileExpectedState temp_expected_state(temp_expected_state_file_path,
+                                          max_key_, num_column_families_);
+    if (s.ok()) {
+      s = temp_expected_state.Open(true /* create */);
+    }
+    if (s.ok()) {
+      s = Env::Default()->RenameFile(temp_expected_state_file_path,
+                                     expected_state_file_path);
+    }
+  }
+
+  if (s.ok()) {
+    latest_.reset(new FileExpectedState(std::move(expected_state_file_path),
+                                        max_key_, num_column_families_));
+    s = latest_->Open(false /* create */);
+  }
+  return s;
+}
+
+Status FileExpectedStateManager::Clean() {
+  // An incomplete `Open()` could have left behind an invalid temporary file.
+  std::string temp_path = GetTempPathForFilename(kLatestFilename);
+  Status s = Env::Default()->FileExists(temp_path);
+  if (s.ok()) {
+    s = Env::Default()->DeleteFile(temp_path);
+  } else if (s.IsNotFound()) {
+    s = Status::OK();
+  }
+  return s;
+}
+
+std::string FileExpectedStateManager::GetTempPathForFilename(
+    const std::string& filename) {
+  static const std::string kTempFilenamePrefix = ".";
+  static const std::string kTempFilenameSuffix = ".tmp";
+
+  assert(!expected_state_dir_path_.empty());
+  std::string expected_state_dir_path_slash =
+      expected_state_dir_path_.back() == '/' ? expected_state_dir_path_
+                                             : expected_state_dir_path_ + "/";
+  return expected_state_dir_path_slash + kTempFilenamePrefix + filename +
+         kTempFilenameSuffix;
+}
+
+std::string FileExpectedStateManager::GetPathForFilename(
+    const std::string& filename) {
+  assert(!expected_state_dir_path_.empty());
+  std::string expected_state_dir_path_slash =
+      expected_state_dir_path_.back() == '/' ? expected_state_dir_path_
+                                             : expected_state_dir_path_ + "/";
+  return expected_state_dir_path_slash + filename;
+}
+
+AnonExpectedStateManager::AnonExpectedStateManager(size_t max_key,
+                                                   size_t num_column_families)
+    : ExpectedStateManager(max_key, num_column_families) {}
+
+Status AnonExpectedStateManager::Open() {
+  latest_.reset(new AnonExpectedState(max_key_, num_column_families_));
+  return latest_->Open(true /* create */);
+}
+
+}  // namespace ROCKSDB_NAMESPACE
+
+#endif  // GFLAGS

db_stress_tool/expected_state.h

@@ -0,0 +1,215 @@
+//  Copyright (c) 2021-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).
+
+#ifdef GFLAGS
+
+#pragma once
+
+#include <stdint.h>
+
+#include <atomic>
+#include <memory>
+
+#include "rocksdb/env.h"
+#include "rocksdb/rocksdb_namespace.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+// An `ExpectedState` provides read/write access to expected values for every
+// key.
+class ExpectedState {
+ public:
+  explicit ExpectedState(size_t max_key, size_t num_column_families);
+
+  virtual ~ExpectedState() {}
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  virtual Status Open(bool create) = 0;
+
+  // Requires external locking covering all keys in `cf`.
+  void ClearColumnFamily(int cf);
+
+  // @param pending True if the update may have started but is not yet
+  //    guaranteed finished. This is useful for crash-recovery testing when the
+  //    process may crash before updating the expected values array.
+  //
+  // Requires external locking covering `key` in `cf`.
+  void Put(int cf, int64_t key, uint32_t value_base, bool pending);
+
+  // Requires external locking covering `key` in `cf`.
+  uint32_t Get(int cf, int64_t key) const;
+
+  // @param pending See comment above Put()
+  // Returns true if the key was not yet deleted.
+  //
+  // Requires external locking covering `key` in `cf`.
+  bool Delete(int cf, int64_t key, bool pending);
+
+  // @param pending See comment above Put()
+  // Returns true if the key was not yet deleted.
+  //
+  // Requires external locking covering `key` in `cf`.
+  bool SingleDelete(int cf, int64_t key, bool pending);
+
+  // @param pending See comment above Put()
+  // Returns number of keys deleted by the call.
+  //
+  // Requires external locking covering keys in `[begin_key, end_key)` in `cf`.
+  int DeleteRange(int cf, int64_t begin_key, int64_t end_key, bool pending);
+
+  // Requires external locking covering `key` in `cf`.
+  bool Exists(int cf, int64_t key);
+
+ private:
+  // Requires external locking covering `key` in `cf`.
+  std::atomic<uint32_t>& Value(int cf, int64_t key) const {
+    return values_[cf * max_key_ + key];
+  }
+
+  const size_t max_key_;
+  const size_t num_column_families_;
+
+ protected:
+  size_t GetValuesLen() const {
+    return sizeof(std::atomic<uint32_t>) * num_column_families_ * max_key_;
+  }
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  void Reset();
+
+  std::atomic<uint32_t>* values_;
+};
+
+// A `FileExpectedState` implements `ExpectedState` backed by a file.
+class FileExpectedState : public ExpectedState {
+ public:
+  explicit FileExpectedState(std::string expected_state_file_path,
+                             size_t max_key, size_t num_column_families);
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  Status Open(bool create) override;
+
+ private:
+  const std::string expected_state_file_path_;
+  std::unique_ptr<MemoryMappedFileBuffer> expected_state_mmap_buffer_;
+};
+
+// An `AnonExpectedState` implements `ExpectedState` backed by a memory
+// allocation.
+class AnonExpectedState : public ExpectedState {
+ public:
+  explicit AnonExpectedState(size_t max_key, size_t num_column_families);
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  Status Open(bool create) override;
+
+ private:
+  std::unique_ptr<std::atomic<uint32_t>[]> values_allocation_;
+};
+
+// An `ExpectedStateManager` manages data about the expected state of the
+// database. It exposes operations for reading and modifying the latest
+// expected state.
+class ExpectedStateManager {
+ public:
+  explicit ExpectedStateManager(size_t max_key, size_t num_column_families);
+
+  virtual ~ExpectedStateManager();
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  virtual Status Open() = 0;
+
+  // Requires external locking covering all keys in `cf`.
+  void ClearColumnFamily(int cf) { return latest_->ClearColumnFamily(cf); }
+
+  // @param pending True if the update may have started but is not yet
+  //    guaranteed finished. This is useful for crash-recovery testing when the
+  //    process may crash before updating the expected values array.
+  //
+  // Requires external locking covering `key` in `cf`.
+  void Put(int cf, int64_t key, uint32_t value_base, bool pending) {
+    return latest_->Put(cf, key, value_base, pending);
+  }
+
+  // Requires external locking covering `key` in `cf`.
+  uint32_t Get(int cf, int64_t key) const { return latest_->Get(cf, key); }
+
+  // @param pending See comment above Put()
+  // Returns true if the key was not yet deleted.
+  //
+  // Requires external locking covering `key` in `cf`.
+  bool Delete(int cf, int64_t key, bool pending) {
+    return latest_->Delete(cf, key, pending);
+  }
+
+  // @param pending See comment above Put()
+  // Returns true if the key was not yet deleted.
+  //
+  // Requires external locking covering `key` in `cf`.
+  bool SingleDelete(int cf, int64_t key, bool pending) {
+    return latest_->SingleDelete(cf, key, pending);
+  }
+
+  // @param pending See comment above Put()
+  // Returns number of keys deleted by the call.
+  //
+  // Requires external locking covering keys in `[begin_key, end_key)` in `cf`.
+  int DeleteRange(int cf, int64_t begin_key, int64_t end_key, bool pending) {
+    return latest_->DeleteRange(cf, begin_key, end_key, pending);
+  }
+
+  // Requires external locking covering `key` in `cf`.
+  bool Exists(int cf, int64_t key) { return latest_->Exists(cf, key); }
+
+ protected:
+  const size_t max_key_;
+  const size_t num_column_families_;
+  std::unique_ptr<ExpectedState> latest_;
+};
+
+// A `FileExpectedStateManager` implements an `ExpectedStateManager` backed by
+// a directory of files containing data about the expected state of the
+// database.
+class FileExpectedStateManager : public ExpectedStateManager {
+ public:
+  explicit FileExpectedStateManager(size_t max_key, size_t num_column_families,
+                                    std::string expected_state_dir_path);
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  Status Open() override;
+
+ private:
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  Status Clean();
+
+  std::string GetTempPathForFilename(const std::string& filename);
+  std::string GetPathForFilename(const std::string& filename);
+
+  static const std::string kLatestFilename;
+
+  const std::string expected_state_dir_path_;
+};
+
+// An `AnonExpectedStateManager` implements an `ExpectedStateManager` backed by
+// a memory allocation containing data about the expected state of the database.
+class AnonExpectedStateManager : public ExpectedStateManager {
+ public:
+  explicit AnonExpectedStateManager(size_t max_key, size_t num_column_families);
+
+  // Requires external locking preventing concurrent execution with any other
+  // member function.
+  Status Open() override;
+};
+
+}  // namespace ROCKSDB_NAMESPACE
+
+#endif  // GFLAGS

src.mk

@@ -345,6 +345,7 @@ STRESS_LIB_SOURCES =                                            \
   db_stress_tool/db_stress_gflags.cc                           \
   db_stress_tool/db_stress_shared_state.cc                     \
   db_stress_tool/db_stress_tool.cc                             \
+  db_stress_tool/expected_state.cc                             \
   db_stress_tool/no_batched_ops_stress.cc                      \
 
 TEST_LIB_SOURCES =                                              \

tools/db_crashtest.py

@@ -60,7 +60,7 @@ default_params = {
     "destroy_db_initially": 0,
     "enable_pipelined_write": lambda: random.randint(0, 1),
     "enable_compaction_filter": lambda: random.choice([0, 0, 0, 1]),
-    "expected_values_path": lambda: setup_expected_values_file(),
+    "expected_values_dir": lambda: setup_expected_values_dir(),
     "fail_if_options_file_error": lambda: random.randint(0, 1),
     "flush_one_in": 1000000,
     "file_checksum_impl": lambda: random.choice(["none", "crc32c", "xxh64", "big"]),
@@ -172,23 +172,23 @@ def get_dbname(test_name):
         os.mkdir(dbname)
     return dbname
 
-expected_values_file = None
+expected_values_dir = None
-def setup_expected_values_file():
+def setup_expected_values_dir():
-    global expected_values_file
+    global expected_values_dir
-    if expected_values_file is not None:
+    if expected_values_dir is not None:
-        return expected_values_file
+        return expected_values_dir
-    expected_file_name = "rocksdb_crashtest_" + "expected"
+    expected_dir_prefix = "rocksdb_crashtest_expected_"
     test_tmpdir = os.environ.get(_TEST_DIR_ENV_VAR)
     if test_tmpdir is None or test_tmpdir == "":
-        expected_values_file = tempfile.NamedTemporaryFile(
+        expected_values_dir = tempfile.mkdtemp(
-            prefix=expected_file_name, delete=False).name
+            prefix=expected_dir_prefix)
     else:
-        # if tmpdir is specified, store the expected_values_file in the same dir
+        # if tmpdir is specified, store the expected_values_dir under that dir
-        expected_values_file = test_tmpdir + "/" + expected_file_name
+        expected_values_dir = test_tmpdir + "/rocksdb_crashtest_expected"
-        if os.path.exists(expected_values_file):
+        if os.path.exists(expected_values_dir):
-            os.remove(expected_values_file)
+            shutil.rmtree(expected_values_dir)
-        open(expected_values_file, 'a').close()
+        os.mkdir(expected_values_dir)
-    return expected_values_file
+    return expected_values_dir
 
 
 def is_direct_io_supported(dbname):
@@ -673,7 +673,7 @@ def whitebox_crash_main(args, unknown_args):
             # success
             shutil.rmtree(dbname, True)
             os.mkdir(dbname)
-            cmd_params.pop('expected_values_path', None)
+            cmd_params.pop('expected_values_dir', None)
             check_mode = (check_mode + 1) % total_check_mode
 
         time.sleep(1)  # time to stabilize after a kill
@@ -718,9 +718,9 @@ def main():
         blackbox_crash_main(args, unknown_args)
     if args.test_type == 'whitebox':
         whitebox_crash_main(args, unknown_args)
-    # Only delete the `expected_values_file` if test passes
+    # Only delete the `expected_values_dir` if test passes
-    if os.path.exists(expected_values_file):
+    if expected_values_dir is not None:
-        os.remove(expected_values_file)
+        shutil.rmtree(expected_values_dir)
 
 
 if __name__ == '__main__':