From 2543d00314b70b1390481476e553acd161ea783d Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Sun, 12 Feb 2023 03:29:19 +0300
Subject: [PATCH] Add erase_batch to key-values.

---
 tddb/td/db/BinlogKeyValue.h        | 17 +++++++++++++++++
 tddb/td/db/KeyValueSyncInterface.h |  2 ++
 tddb/td/db/SeqKeyValue.h           | 17 +++++++++++++++++
 tddb/td/db/SqliteKeyValue.cpp      |  6 ++++++
 tddb/td/db/SqliteKeyValue.h        |  2 ++
 tddb/td/db/TsSeqKeyValue.h         |  5 +++++
 6 files changed, 49 insertions(+)

diff --git a/tddb/td/db/BinlogKeyValue.h b/tddb/td/db/BinlogKeyValue.h
index ba3aa9d04..6b6e68422 100644
--- a/tddb/td/db/BinlogKeyValue.h
+++ b/tddb/td/db/BinlogKeyValue.h
@@ -166,6 +166,23 @@ class BinlogKeyValue final : public KeyValueSyncInterface {
     return seq_no;
   }
 
+  SeqNo erase_batch(vector<string> keys) final {
+    auto lock = rw_mutex_.lock_write().move_as_ok();
+    vector<uint64> log_event_ids;
+    for (auto &key : keys) {
+      auto it = map_.find(key);
+      if (it != map_.end()) {
+        log_event_ids.push_back(it->second.second);
+        map_.erase(it);
+      }
+    }
+    if (log_event_ids.empty()) {
+      return 0;
+    }
+    VLOG(binlog) << "Remove value of keys " << keys;
+    return binlog_->erase_batch(std::move(log_event_ids));
+  }
+
   void add_event(uint64 seq_no, BufferSlice &&event) {
     binlog_->add_raw_event(BinlogDebugInfo{__FILE__, __LINE__}, seq_no, std::move(event));
   }
diff --git a/tddb/td/db/KeyValueSyncInterface.h b/tddb/td/db/KeyValueSyncInterface.h
index c3f5f5b99..111882fb2 100644
--- a/tddb/td/db/KeyValueSyncInterface.h
+++ b/tddb/td/db/KeyValueSyncInterface.h
@@ -40,6 +40,8 @@ class KeyValueSyncInterface {
 
   virtual SeqNo erase(const string &key) = 0;
 
+  virtual SeqNo erase_batch(vector<string> keys) = 0;
+
   virtual void erase_by_prefix(Slice prefix) = 0;
 
   virtual void force_sync(Promise<> &&promise) = 0;
diff --git a/tddb/td/db/SeqKeyValue.h b/tddb/td/db/SeqKeyValue.h
index 7e5ca7a9d..9870ac939 100644
--- a/tddb/td/db/SeqKeyValue.h
+++ b/tddb/td/db/SeqKeyValue.h
@@ -43,6 +43,23 @@ class SeqKeyValue {
     return next_seq_no();
   }
 
+  SeqNo erase_batch(vector<string> keys) {
+    size_t count = 0;
+    for (auto &key : keys) {
+      auto it = map_.find(key);
+      if (it != map_.end()) {
+        map_.erase(it);
+        count++;
+      }
+    }
+    if (count == 0) {
+      return 0;
+    }
+    SeqNo result = current_id_ + 1;
+    current_id_ += count;
+    return result;
+  }
+
   SeqNo seq_no() const {
     return current_id_ + 1;
   }
diff --git a/tddb/td/db/SqliteKeyValue.cpp b/tddb/td/db/SqliteKeyValue.cpp
index 0e933dec5..42b6d8bce 100644
--- a/tddb/td/db/SqliteKeyValue.cpp
+++ b/tddb/td/db/SqliteKeyValue.cpp
@@ -88,6 +88,12 @@ void SqliteKeyValue::erase(Slice key) {
   erase_stmt_.reset();
 }
 
+void SqliteKeyValue::erase_batch(vector<string> keys) {
+  for (auto &key : keys) {
+    erase(key);
+  }
+}
+
 void SqliteKeyValue::erase_by_prefix(Slice prefix) {
   auto next = next_prefix(prefix);
   if (next.empty()) {
diff --git a/tddb/td/db/SqliteKeyValue.h b/tddb/td/db/SqliteKeyValue.h
index 9d0d7cd76..914825da8 100644
--- a/tddb/td/db/SqliteKeyValue.h
+++ b/tddb/td/db/SqliteKeyValue.h
@@ -47,6 +47,8 @@ class SqliteKeyValue {
 
   void erase(Slice key);
 
+  void erase_batch(vector<string> keys);
+
   Status begin_read_transaction() TD_WARN_UNUSED_RESULT {
     return db_.begin_read_transaction();
   }
diff --git a/tddb/td/db/TsSeqKeyValue.h b/tddb/td/db/TsSeqKeyValue.h
index 6e7019c2d..ebb413289 100644
--- a/tddb/td/db/TsSeqKeyValue.h
+++ b/tddb/td/db/TsSeqKeyValue.h
@@ -45,6 +45,11 @@ class TsSeqKeyValue {
     return kv_.erase(key);
   }
 
+  SeqNo erase_batch(vector<string> keys) {
+    auto lock = rw_mutex_.lock_write().move_as_ok();
+    return kv_.erase_batch(std::move(keys));
+  }
+
   std::pair<SeqNo, RwMutex::WriteLock> erase_and_lock(const string &key) {
     auto lock = rw_mutex_.lock_write().move_as_ok();
     return std::make_pair(kv_.erase(key), std::move(lock));