1 files changed, 96 insertions, 0 deletions
diff --git a/src/rocksdb/db_stress_tool/db_stress_compaction_filter.h b/src/rocksdb/db_stress_tool/db_stress_compaction_filter.h
new file mode 100644
index 000000000..408bb48f3
--- /dev/null
+++ b/src/rocksdb/db_stress_tool/db_stress_compaction_filter.h
@@ -0,0 +1,96 @@
+//  Copyright (c) 2011-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).
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+
+#pragma once
+
+#include "db_stress_tool/db_stress_common.h"
+#include "db_stress_tool/db_stress_shared_state.h"
+#include "rocksdb/compaction_filter.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+// DbStressCompactionFilter is safe to use with db_stress as it does not perform
+// any mutation. It only makes `kRemove` decisions for keys that are already
+// non-existent according to the `SharedState`.
+class DbStressCompactionFilter : public CompactionFilter {
+ public:
+  DbStressCompactionFilter(SharedState* state, int cf_id)
+      : state_(state), cf_id_(cf_id) {}
+
+  Decision FilterV2(int /*level*/, const Slice& key, ValueType /*value_type*/,
+                    const Slice& /*existing_value*/, std::string* /*new_value*/,
+                    std::string* /*skip_until*/) const override {
+    if (state_ == nullptr) {
+      return Decision::kKeep;
+    }
+    if (key.empty() || ('0' <= key[0] && key[0] <= '9')) {
+      // It is likely leftover from a test_batches_snapshots run. Below this
+      // conditional, the test_batches_snapshots key format is not handled
+      // properly. Just keep it to be safe.
+      return Decision::kKeep;
+    }
+    uint64_t key_num = 0;
+    {
+      Slice ukey_without_ts = key;
+      assert(ukey_without_ts.size() >= FLAGS_user_timestamp_size);
+      ukey_without_ts.remove_suffix(FLAGS_user_timestamp_size);
+      [[maybe_unused]] bool ok =
+          GetIntVal(ukey_without_ts.ToString(), &key_num);
+      assert(ok);
+    }
+    port::Mutex* key_mutex = state_->GetMutexForKey(cf_id_, key_num);
+    if (!key_mutex->TryLock()) {
+      return Decision::kKeep;
+    }
+    // Reaching here means we acquired the lock.
+
+    bool key_exists = state_->Exists(cf_id_, key_num);
+    const bool allow_overwrite = state_->AllowsOverwrite(key_num);
+
+    key_mutex->Unlock();
+
+    if (!key_exists) {
+      return allow_overwrite ? Decision::kRemove : Decision::kPurge;
+    }
+    return Decision::kKeep;
+  }
+
+  const char* Name() const override { return "DbStressCompactionFilter"; }
+
+ private:
+  SharedState* const state_;
+  const int cf_id_;
+};
+
+class DbStressCompactionFilterFactory : public CompactionFilterFactory {
+ public:
+  DbStressCompactionFilterFactory() : state_(nullptr) {}
+
+  void SetSharedState(SharedState* state) {
+    MutexLock state_mutex_guard(&state_mutex_);
+    state_ = state;
+  }
+
+  std::unique_ptr<CompactionFilter> CreateCompactionFilter(
+      const CompactionFilter::Context& context) override {
+    MutexLock state_mutex_guard(&state_mutex_);
+    return std::unique_ptr<CompactionFilter>(
+        new DbStressCompactionFilter(state_, context.column_family_id));
+  }
+
+  const char* Name() const override {
+    return "DbStressCompactionFilterFactory";
+  }
+
+ private:
+  port::Mutex state_mutex_;
+  SharedState* state_;
+};
+
+}  // namespace ROCKSDB_NAMESPACE