Add option to validate key during seek in SkipList Memtable (#13902)

Summary:
Add a new CF immutable option `paranoid_memory_check_key_checksum_on_seek` that allows additional data integrity validations during seek on SkipList Memtable. When this option is enabled and memtable_protection_bytes_per_key is non zero, skiplist-based memtable will validate the checksum of each key visited during seek operation. The option is opt-in due to performance overhead. This is an enhancement on top of paranoid_memory_checks option.

Pull Request resolved: #13902

Test Plan:
* new unit test added for paranoid_memory_check_key_checksum_on_seek=true.
    * existing unit test for paranoid_memory_check_key_checksum_on_seek=false.
    * enable in stress test.

Performance Benchmark: we check for performance regression in read path where data is in memtable only. For each benchmark, the script was run at the same time for main and this PR:

### Memtable-only randomread ops/sec:

* Value size = 100 Bytes
```
for B in 0 1 2 4 8; do (for I in $(seq 1 50);do  ./db_bench --benchmarks=fillseq,readrandom --write_buffer_size=268435456 --writes=250000 --value_size=100 --num=250000 --reads=500000  --seed=1723056275 --paranoid_memory_check_key_checksum_on_seek=true --memtable_protection_bytes_per_key=$B 2>&1 | grep "readrandom"; done;) | awk '{ t += $5; c++; print } END { print 1.0 * t / c }'; done;
```

1. Main: 928999
2. PR with paranoid_memory_check_key_checksum_on_seek=false: 930993 (+0.2%)
3. PR with paranoid_memory_check_key_checksum_on_seek=true:
3.1 memtable_protection_bytes_per_key=1: 464577 (-50%)
3.2 memtable_protection_bytes_per_key=2: 470319 (-49%)
3.3 memtable_protection_bytes_per_key=4: 468457 (-50%)
3.4 memtable_protection_bytes_per_key=8: 465061 (-50%)

* Value size = 1000 Bytes
```
for B in 0 1 2 4 8; do (for I in $(seq 1 50);do  ./db_bench --benchmarks=fillseq,readrandom --write_buffer_size=268435456 --writes=250000 --value_size=1000 --num=250000 --reads=500000  --seed=1723056275 --paranoid_memory_check_key_checksum_on_seek=true --memtable_protection_bytes_per_key=$B 2>&1 | grep "readrandom"; done;) | awk '{ t += $5; c++; print } END { print 1.0 * t / c }'; done;
```

1. Main: 601321
2. PR with paranoid_memory_check_key_checksum_on_seek=false: 607885 (+1.1%)
3. PR with paranoid_memory_check_key_checksum_on_seek=true:
3.1 memtable_protection_bytes_per_key=1: 185742 (-69%)
3.2 memtable_protection_bytes_per_key=2: 177167 (-71%)
3.3 memtable_protection_bytes_per_key=4: 185908 (-69%)
3.4 memtable_protection_bytes_per_key=8: 183639 (-69%)

Reviewed By: pdillinger

Differential Revision: D81199245

Pulled By: xingbowang

fbshipit-source-id: e3c29552ab92f2c5f360361366a293fa26934913

Author: Xingbo Wang

db/db_basic_test.cc

@@ -5087,6 +5087,7 @@ TEST_F(DBBasicTest, DisallowMemtableWrite) {
   Options options_disallow = options_allow;
   options_disallow.disallow_memtable_writes = true;
   options_disallow.paranoid_memory_checks = true;
+  options_disallow.memtable_veirfy_per_key_checksum_on_seek = true;
 
   DestroyAndReopen(options_allow);
   // CFs allowing and disallowing memtable write

db/db_memtable_test.cc

@@ -339,6 +339,135 @@ TEST_F(DBMemTableTest, ColumnFamilyId) {
   }
 }
 
+class DBMemTableTestForSeek : public DBMemTableTest,
+                              virtual public ::testing::WithParamInterface<
+                                  std::tuple<bool, bool, bool>> {};
+
+TEST_P(DBMemTableTestForSeek, IntegrityChecks) {
+  // Validate key corruption could be detected during seek.
+  // We insert many keys into skiplist. Then we corrupt the each key one at a
+  // time. With memtable_veirfy_per_key_checksum_on_seek enabled, when the
+  // corrupted key is searched, the checksum of every key visited during the
+  // seek is validated. It will report data corruption. Otherwise seek returns
+  // not found.
+  auto allow_data_in_error = std::get<0>(GetParam());
+  Options options = CurrentOptions();
+  options.allow_data_in_errors = allow_data_in_error;
+  options.paranoid_memory_checks = std::get<1>(GetParam());
+  options.memtable_veirfy_per_key_checksum_on_seek = std::get<2>(GetParam());
+  options.memtable_protection_bytes_per_key = 8;
+  DestroyAndReopen(options);
+
+  // capture the data pointer of all of the keys
+  std::vector<char*> raw_data_pointer;
+
+  // Insert enough keys, so memtable would create multiple levels.
+  auto key_count = 100;
+  for (int i = 0; i < key_count; i++) {
+    // The last digit of the key will be corrupted from value 0 to value 5
+    ASSERT_OK(Put(Key(i * 10), "val0"));
+  }
+
+  ReadOptions rops;
+
+  // Iterate all the keys to get key pointers
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->SetCallBack("InlineSkipList::Iterator::Next::key",
+                                        [&raw_data_pointer](void* key) {
+                                          auto p = static_cast<char*>(key);
+                                          raw_data_pointer.push_back(p);
+                                        });
+  SyncPoint::GetInstance()->EnableProcessing();
+
+  {
+    std::unique_ptr<Iterator> iter{db_->NewIterator(rops)};
+    iter->Seek(Key(0));
+    while (iter->Valid()) {
+      ASSERT_OK(iter->status());
+      iter->Next();
+    }
+    // check status after valid returned false.
+    auto status = iter->status();
+    ASSERT_TRUE(status.ok());
+  }
+
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->ClearAllCallBacks();
+
+  ASSERT_EQ(raw_data_pointer.size(), key_count);
+
+  bool enable_key_validation_on_seek =
+      options.memtable_veirfy_per_key_checksum_on_seek;
+
+  // For each key, corrupt it, validate corruption is detected correctly, then
+  // revert it.
+  for (int i = 0; i < key_count; i++) {
+    std::string key_to_corrupt = Key(i * 10);
+    raw_data_pointer[i][key_to_corrupt.size()] = '5';
+
+    auto corrupted_key = key_to_corrupt;
+    corrupted_key.data()[key_to_corrupt.size() - 1] = '5';
+    auto corrupted_key_slice =
+        Slice(corrupted_key.data(), corrupted_key.length());
+    auto corrupted_key_hex = corrupted_key_slice.ToString(/*hex=*/true);
+
+    {
+      // Test Get API
+      std::string val;
+      auto status = db_->Get(rops, key_to_corrupt, &val);
+      if (enable_key_validation_on_seek) {
+        ASSERT_TRUE(status.IsCorruption()) << key_to_corrupt;
+        ASSERT_EQ(
+            status.ToString().find(corrupted_key_hex) != std::string::npos,
+            allow_data_in_error)
+            << status.ToString() << "\n"
+            << corrupted_key_hex;
+      } else {
+        ASSERT_TRUE(status.IsNotFound());
+      }
+    }
+
+    {
+      // Test MultiGet API
+      std::vector<std::string> vals;
+      std::vector<Status> statuses = db_->MultiGet(
+          rops, {db_->DefaultColumnFamily()}, {key_to_corrupt}, &vals, nullptr);
+      if (enable_key_validation_on_seek) {
+        ASSERT_TRUE(statuses[0].IsCorruption());
+        ASSERT_EQ(
+            statuses[0].ToString().find(corrupted_key_hex) != std::string::npos,
+            allow_data_in_error);
+      } else {
+        ASSERT_TRUE(statuses[0].IsNotFound());
+      }
+    }
+
+    {
+      // Test Iterator Seek API
+      std::unique_ptr<Iterator> iter{db_->NewIterator(rops)};
+      ASSERT_OK(iter->status());
+      iter->Seek(key_to_corrupt);
+      auto status = iter->status();
+      if (enable_key_validation_on_seek) {
+        ASSERT_TRUE(status.IsCorruption());
+        ASSERT_EQ(
+            status.ToString().find(corrupted_key_hex) != std::string::npos,
+            allow_data_in_error);
+      } else {
+        ASSERT_FALSE(iter->Valid());
+        ASSERT_FALSE(status.ok());
+      }
+    }
+
+    // revert the key corruption.
+    raw_data_pointer[i][key_to_corrupt.size()] = '0';
+  }
+}
+
+INSTANTIATE_TEST_CASE_P(DBMemTableTestForSeek, DBMemTableTestForSeek,
+                        ::testing::Combine(::testing::Bool(), ::testing::Bool(),
+                                           ::testing::Bool()));
+
 TEST_F(DBMemTableTest, IntegrityChecks) {
   // We insert keys key000000, key000001 and key000002 into skiplist at fixed
   // height 1 (smallest height). Then we corrupt the second key to aey000001 to

db/memtable.cc

@@ -70,7 +70,9 @@ ImmutableMemTableOptions::ImmutableMemTableOptions(
       protection_bytes_per_key(
           mutable_cf_options.memtable_protection_bytes_per_key),
       allow_data_in_errors(ioptions.allow_data_in_errors),
-      paranoid_memory_checks(mutable_cf_options.paranoid_memory_checks) {}
+      paranoid_memory_checks(mutable_cf_options.paranoid_memory_checks),
+      memtable_veirfy_per_key_checksum_on_seek(
+          mutable_cf_options.memtable_veirfy_per_key_checksum_on_seek) {}
 
 MemTable::MemTable(const InternalKeyComparator& cmp,
                    const ImmutableOptions& ioptions,
@@ -115,7 +117,13 @@ MemTable::MemTable(const InternalKeyComparator& cmp,
       oldest_key_time_(std::numeric_limits<uint64_t>::max()),
       approximate_memory_usage_(0),
       memtable_max_range_deletions_(
-          mutable_cf_options.memtable_max_range_deletions) {
+          mutable_cf_options.memtable_max_range_deletions),
+      key_validation_callback_(
+          (moptions_.protection_bytes_per_key != 0 &&
+           moptions_.memtable_veirfy_per_key_checksum_on_seek)
+              ? std::bind(&MemTable::ValidateKey, this, std::placeholders::_1,
+                          std::placeholders::_2)
+              : std::function<Status(const char*, bool)>(nullptr)) {
   UpdateFlushState();
   // something went wrong if we need to flush before inserting anything
   assert(!ShouldScheduleFlush());
@@ -394,7 +402,11 @@ class MemTableIterator : public InternalIterator {
             !mem.GetImmutableMemTableOptions()->inplace_update_support),
         arena_mode_(arena != nullptr),
         paranoid_memory_checks_(mem.moptions_.paranoid_memory_checks),
-        allow_data_in_error(mem.moptions_.allow_data_in_errors) {
+        validate_on_seek_(
+            mem.moptions_.paranoid_memory_checks ||
+            mem.moptions_.memtable_veirfy_per_key_checksum_on_seek),
+        allow_data_in_error_(mem.moptions_.allow_data_in_errors),
+        key_validation_callback_(mem.key_validation_callback_) {
     if (kind == kRangeDelEntries) {
       iter_ = mem.range_del_table_->GetIterator(arena);
     } else if (prefix_extractor_ != nullptr &&
@@ -463,8 +475,10 @@ class MemTableIterator : public InternalIterator {
         }
       }
     }
-    if (paranoid_memory_checks_) {
-      status_ = iter_->SeekAndValidate(k, nullptr, allow_data_in_error);
+    if (validate_on_seek_) {
+      status_ = iter_->SeekAndValidate(k, nullptr, allow_data_in_error_,
+                                       paranoid_memory_checks_,
+                                       key_validation_callback_);
     } else {
       iter_->Seek(k, nullptr);
     }
@@ -488,8 +502,10 @@ class MemTableIterator : public InternalIterator {
         }
       }
     }
-    if (paranoid_memory_checks_) {
-      status_ = iter_->SeekAndValidate(k, nullptr, allow_data_in_error);
+    if (validate_on_seek_) {
+      status_ = iter_->SeekAndValidate(k, nullptr, allow_data_in_error_,
+                                       paranoid_memory_checks_,
+                                       key_validation_callback_);
     } else {
       iter_->Seek(k, nullptr);
     }
@@ -518,7 +534,7 @@ class MemTableIterator : public InternalIterator {
     PERF_COUNTER_ADD(next_on_memtable_count, 1);
     assert(Valid());
     if (paranoid_memory_checks_) {
-      status_ = iter_->NextAndValidate(allow_data_in_error);
+      status_ = iter_->NextAndValidate(allow_data_in_error_);
     } else {
       iter_->Next();
       TEST_SYNC_POINT_CALLBACK("MemTableIterator::Next:0", iter_);
@@ -540,7 +556,7 @@ class MemTableIterator : public InternalIterator {
     PERF_COUNTER_ADD(prev_on_memtable_count, 1);
     assert(Valid());
     if (paranoid_memory_checks_) {
-      status_ = iter_->PrevAndValidate(allow_data_in_error);
+      status_ = iter_->PrevAndValidate(allow_data_in_error_);
     } else {
       iter_->Prev();
     }
@@ -599,7 +615,9 @@ class MemTableIterator : public InternalIterator {
   bool value_pinned_;
   bool arena_mode_;
   const bool paranoid_memory_checks_;
-  const bool allow_data_in_error;
+  const bool validate_on_seek_;
+  const bool allow_data_in_error_;
+  const std::function<Status(const char*, bool)> key_validation_callback_;
 
   void VerifyEntryChecksum() {
     if (protection_bytes_per_key_ > 0 && Valid()) {
@@ -1493,11 +1511,13 @@ void MemTable::GetFromTable(const LookupKey& key,
   saver.allow_data_in_errors = moptions_.allow_data_in_errors;
   saver.protection_bytes_per_key = moptions_.protection_bytes_per_key;
 
-  if (!moptions_.paranoid_memory_checks) {
+  if (!moptions_.paranoid_memory_checks &&
+      !moptions_.memtable_veirfy_per_key_checksum_on_seek) {
     table_->Get(key, &saver, SaveValue);
   } else {
-    Status check_s = table_->GetAndValidate(key, &saver, SaveValue,
-                                            moptions_.allow_data_in_errors);
+    Status check_s = table_->GetAndValidate(
+        key, &saver, SaveValue, moptions_.allow_data_in_errors,
+        moptions_.paranoid_memory_checks, key_validation_callback_);
     if (check_s.IsCorruption()) {
       *(saver.status) = check_s;
       // Should stop searching the LSM.
@@ -1508,6 +1528,11 @@ void MemTable::GetFromTable(const LookupKey& key,
   *seq = saver.seq;
 }
 
+Status MemTable::ValidateKey(const char* key, bool allow_data_in_errors) {
+  return VerifyEntryChecksum(key, moptions_.protection_bytes_per_key,
+                             allow_data_in_errors);
+}
+
 void MemTable::MultiGet(const ReadOptions& read_options, MultiGetRange* range,
                         ReadCallback* callback, bool immutable_memtable) {
   // The sequence number is updated synchronously in version_set.h

db/memtable.h

@@ -64,6 +64,7 @@ struct ImmutableMemTableOptions {
   uint32_t protection_bytes_per_key;
   bool allow_data_in_errors;
   bool paranoid_memory_checks;
+  bool memtable_veirfy_per_key_checksum_on_seek;
 };
 
 // Batched counters to updated when inserting keys in one write batch.
@@ -826,6 +827,9 @@ class MemTable final : public ReadOnlyMemTable {
                                     uint32_t protection_bytes_per_key,
                                     bool allow_data_in_errors = false);
 
+  // Validate the checksum of the key/value pair.
+  Status ValidateKey(const char* key, bool allow_data_in_errors);
+
  private:
   enum FlushStateEnum { FLUSH_NOT_REQUESTED, FLUSH_REQUESTED, FLUSH_SCHEDULED };
 
@@ -956,6 +960,8 @@ class MemTable final : public ReadOnlyMemTable {
                            SequenceNumber s, char* checksum_ptr);
 
   void MaybeUpdateNewestUDT(const Slice& user_key);
+
+  const std::function<Status(const char*, bool)> key_validation_callback_;
 };
 
 const char* EncodeKey(std::string* scratch, const Slice& target);

db_stress_tool/db_stress_common.h

@@ -275,6 +275,7 @@ DECLARE_string(last_level_temperature);
 DECLARE_string(default_write_temperature);
 DECLARE_string(default_temperature);
 DECLARE_bool(paranoid_memory_checks);
+DECLARE_bool(memtable_veirfy_per_key_checksum_on_seek);
 
 // Options for transaction dbs.
 // Use TransactionDB (a.k.a. Pessimistic Transaction DB)

db_stress_tool/db_stress_gflags.cc

@@ -1484,6 +1484,11 @@ DEFINE_bool(paranoid_memory_checks,
             ROCKSDB_NAMESPACE::Options().paranoid_memory_checks,
             "Sets CF option paranoid_memory_checks.");
 
+DEFINE_bool(
+    memtable_veirfy_per_key_checksum_on_seek,
+    ROCKSDB_NAMESPACE::Options().memtable_veirfy_per_key_checksum_on_seek,
+    "Sets CF option memtable_veirfy_per_key_checksum_on_seek.");
+
 DEFINE_uint32(commit_bypass_memtable_one_in, 0,
               "If greater than zero, transaction option will set "
               "commit_bypass_memtable to per every N transactions on average.");

db_stress_tool/db_stress_test_base.cc

@@ -4442,6 +4442,8 @@ void InitializeOptionsFromFlags(
       FLAGS_memtable_protection_bytes_per_key;
   options.block_protection_bytes_per_key = FLAGS_block_protection_bytes_per_key;
   options.paranoid_memory_checks = FLAGS_paranoid_memory_checks;
+  options.memtable_veirfy_per_key_checksum_on_seek =
+      FLAGS_memtable_veirfy_per_key_checksum_on_seek;
 
   // Integrated BlobDB
   options.enable_blob_files = FLAGS_enable_blob_files;

include/rocksdb/advanced_options.h

@@ -1101,12 +1101,22 @@ struct AdvancedColumnFamilyOptions {
   uint32_t bottommost_file_compaction_delay = 0;
 
   // Enables additional integrity checks during reads/scans.
-  // Specifically, for skiplist-based memtables, we verify that keys visited
-  // are in order. This is helpful to detect corrupted memtable keys during
-  // reads. Enabling this feature incurs a performance overhead due to an
-  // additional key comparison during memtable lookup.
+  // Specifically, for skiplist-based memtables, key ordering validation could
+  // be enabled optionally. This is helpful to detect corrupted memtable keys
+  // during reads. Enabling this feature incurs a performance overhead due to
+  // additional comparison during memtable lookup.
   bool paranoid_memory_checks = false;
 
+  // Enables additional integrity checks during seek.
+  // Specifically, for skiplist-based memtables, key checksum validation could
+  // be enabled during seek optionally. This is helpful to detect corrupted
+  // memtable keys during reads. Enabling this feature incurs a performance
+  // overhead due to additional key checksum validation during memtable seek
+  // operation.
+  // This option depends on memtable_protection_bytes_per_key to be non zero.
+  // If memtable_protection_bytes_per_key is zero, no validation is performed.
+  bool memtable_veirfy_per_key_checksum_on_seek = false;
+
   // When an iterator scans this number of invisible entries (tombstones or
   // hidden puts) from the active memtable during a single iterator operation,
   // we will attempt to flush the memtable. Currently only forward scans are

include/rocksdb/memtablerep.h

@@ -38,6 +38,7 @@
 #include <stdint.h>
 #include <stdlib.h>
 
+#include <functional>
 #include <memory>
 #include <stdexcept>
 #include <unordered_set>
@@ -201,11 +202,12 @@ class MemTableRep {
                    bool (*callback_func)(void* arg, const char* entry));
 
   // Same as Get() but performs data integrity validation.
-  virtual Status GetAndValidate(const LookupKey& /* k */,
-                                void* /* callback_args */,
-                                bool (* /* callback_func */)(void* arg,
-                                                             const char* entry),
-                                bool /*allow_data_in_error*/) {
+  virtual Status GetAndValidate(
+      const LookupKey& /* k */, void* /* callback_args */,
+      bool (* /* callback_func */)(void* arg, const char* entry),
+      bool /* allow_data_in_error */, bool /* detect_key_out_of_order */,
+      const std::function<Status(const char*, bool)>&
+      /* key_validation_callback */) {
     return Status::NotSupported("GetAndValidate() not implemented.");
   }
 
@@ -276,9 +278,11 @@ class MemTableRep {
     // Seek and perform integrity validations on the skip list.
     // Iterator becomes invalid and Corruption is returned if a
     // corruption is found.
-    virtual Status SeekAndValidate(const Slice& /* internal_key */,
-                                   const char* /* memtable_key */,
-                                   bool /* allow_data_in_errors */) {
+    virtual Status SeekAndValidate(
+        const Slice& /* internal_key */, const char* /* memtable_key */,
+        bool /* allow_data_in_errors */, bool /* detect_key_out_of_order */,
+        const std::function<Status(const char*, bool)>&
+        /* key_validation_callback */) {
       return Status::NotSupported("SeekAndValidate() not implemented.");
     }
 

include/rocksdb/options.h

@@ -623,6 +623,7 @@ struct DBOptions {
   // checking for corruption, including
   // * paranoid_file_checks
   // * paranoid_memory_checks
+  // * memtable_veirfy_per_key_checksum_on_seek
   // * DB::VerifyChecksum()
   //
   // Default: true