From b493e96bc181e03d1ab5dfe556a866d6ffb16cec Mon Sep 17 00:00:00 2001
From: Mark Nelson <mnelson@redhat.com>
Date: Wed, 6 Nov 2019 02:04:14 +0000
Subject: [PATCH] kv/RocksDBStore: Add minimum key limit before invoking
 DeleteRange.

Co-authored-by: Zengran Zhang zhangzengran@sangfor.com.cn
Signed-off-by: Mark Nelson <mnelson@redhat.com>
---
 src/common/options.cc  |  4 +++
 src/kv/RocksDBStore.cc | 69 ++++++++++++++++++++++++++++++++----------
 src/kv/RocksDBStore.h  |  4 ++-
 3 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/src/common/options.cc b/src/common/options.cc
index 1f76931f805..e849683a8da 100644
--- a/src/common/options.cc
+++ b/src/common/options.cc
@@ -3597,6 +3597,10 @@ std::vector<Option> get_global_options() {
     .set_default(false)
     .set_description(""),
 
+    Option("rocksdb_delete_range_threshold", Option::TYPE_UINT, Option::LEVEL_ADVANCED)
+    .set_default(1024)
+    .set_description("The number of keys required to invoke DeleteRange when deleting muliple keys."),
+
     Option("rocksdb_bloom_bits_per_key", Option::TYPE_UINT, Option::LEVEL_ADVANCED)
     .set_default(20)
     .set_description("Number of bits per key to use for RocksDB's bloom filters.")
diff --git a/src/kv/RocksDBStore.cc b/src/kv/RocksDBStore.cc
index 03a265bc34f..19df97111f6 100644
--- a/src/kv/RocksDBStore.cc
+++ b/src/kv/RocksDBStore.cc
@@ -953,16 +953,32 @@ void RocksDBStore::RocksDBTransactionImpl::rm_single_key(const string &prefix,
 void RocksDBStore::RocksDBTransactionImpl::rmkeys_by_prefix(const string &prefix)
 {
   auto cf = db->get_cf_handle(prefix);
-  if (cf) {
-    string endprefix("\xff\xff\xff\xff");  // FIXME: this is cheating...
-    bat.DeleteRange(cf, string(), endprefix);
-  } else {
-    string endprefix = prefix;
-    endprefix.push_back('\x01');
-    bat.DeleteRange(db->default_cf,
-      combine_strings(prefix, string()),
-      combine_strings(endprefix, string()));
+  uint64_t cnt = db->delete_range_threshold;
+  bat.SetSavePoint();
+  auto it = db->get_iterator(prefix);
+  for (it->seek_to_first(); it->valid(); it->next()) {
+    if (!cnt) {
+      bat.RollbackToSavePoint();
+      if (cf) {
+        string endprefix = "\xff\xff\xff\xff";  // FIXME: this is cheating...
+        bat.DeleteRange(cf, string(), endprefix);
+      } else {
+        string endprefix = prefix;
+        endprefix.push_back('\x01');
+	bat.DeleteRange(db->default_cf,
+                        combine_strings(prefix, string()),
+                        combine_strings(endprefix, string()));
+      }
+      return;
+    }
+    if (cf) {
+      bat.Delete(cf, rocksdb::Slice(it->key()));
+    } else {
+      bat.Delete(db->default_cf, combine_strings(prefix, it->key()));
+    }
+    --cnt;
   }
+  bat.PopSavePoint();
 }
 
 void RocksDBStore::RocksDBTransactionImpl::rm_range_keys(const string &prefix,
@@ -970,14 +986,35 @@ void RocksDBStore::RocksDBTransactionImpl::rm_range_keys(const string &prefix,
                                                          const string &end)
 {
   auto cf = db->get_cf_handle(prefix);
-  if (cf) {
-    bat.DeleteRange(cf, rocksdb::Slice(start), rocksdb::Slice(end));
-  } else {
-    bat.DeleteRange(
-        db->default_cf,
-        rocksdb::Slice(combine_strings(prefix, start)),
-        rocksdb::Slice(combine_strings(prefix, end)));
+
+  uint64_t cnt = db->delete_range_threshold;
+  auto it = db->get_iterator(prefix);
+  bat.SetSavePoint();
+  it->lower_bound(start);
+  while (it->valid()) {
+    if (it->key() >= end) {
+      break;
+    }
+    if (!cnt) {
+      bat.RollbackToSavePoint();
+      if (cf) {
+        bat.DeleteRange(cf, rocksdb::Slice(start), rocksdb::Slice(end));
+      } else {
+        bat.DeleteRange(db->default_cf,
+                        rocksdb::Slice(combine_strings(prefix, start)),
+                        rocksdb::Slice(combine_strings(prefix, end)));
+      }
+      return;
+    }
+    if (cf) {
+      bat.Delete(cf, rocksdb::Slice(it->key()));
+    } else {
+      bat.Delete(db->default_cf, combine_strings(prefix, it->key()));
+    }
+    it->next();
+    --cnt;
   }
+  bat.PopSavePoint();
 }
 
 void RocksDBStore::RocksDBTransactionImpl::merge(
diff --git a/src/kv/RocksDBStore.h b/src/kv/RocksDBStore.h
index 254f6f3afda..0674ad6dc8e 100644
--- a/src/kv/RocksDBStore.h
+++ b/src/kv/RocksDBStore.h
@@ -120,6 +120,7 @@ public:
   /// compact the underlying rocksdb store
   bool compact_on_mount;
   bool disableWAL;
+  const uint64_t delete_range_threshold;
   void compact() override;
 
   void compact_async() override {
@@ -157,7 +158,8 @@ public:
     compact_queue_stop(false),
     compact_thread(this),
     compact_on_mount(false),
-    disableWAL(false)
+    disableWAL(false),
+    delete_range_threshold(cct->_conf.get_val<uint64_t>("rocksdb_delete_range_threshold"))
   {}
 
   ~RocksDBStore() override;
-- 
2.39.5