This is an automated email from the ASF dual-hosted git repository.

jihuayu pushed a commit to branch unstable
in repository https://gitbox.apache.org/repos/asf/kvrocks.git


The following commit(s) were added to refs/heads/unstable by this push:
     new 7b90ca2bc refactor(cuckoo-filter): move page cache ownership to chain 
(#3609)
7b90ca2bc is described below

commit 7b90ca2bcdd01d31f5efbdf3d348049d22382e26
Author: hurukawa <[email protected]>
AuthorDate: Sat Oct 17 16:07:12 2026 +0800

    refactor(cuckoo-filter): move page cache ownership to chain (#3609)
    
    - Move cuckoo filter page-cache ownership from individual subfilters to
    the chain operation.
    - Share cached pages across subfilters and commit dirty pages with
    metadata at the chain layer.
    - Add slot-level undo for failed kick-out attempts while preserving
    previously staged page mutations.
    - Separate the insert, kick-out, and expansion stages from persistence
    so future multi-item insert operations can reuse the same metadata and
    page cache before committing once.
    
    ## Testing
    
    - `./x.py check format`
    - `cmake --build build --target kvrocks -j2`
    - `cmake --build build --target unittest -j2`
    - `./build/unittest --gtest_filter='*Cuckoo*'`
    - `./build/unittest`
    This PR was written using codex and GPT-5.6 sol
    
    ---------
    
    Co-authored-by: 纪华裕 <[email protected]>
---
 src/types/cuckoo_filter.h                      |  4 +-
 src/types/cuckoo_filter_page.cc                | 51 ++++++++++++++
 src/types/cuckoo_filter_page.h                 | 17 +++++
 src/types/cuckoo_filter_sub_filter.cc          | 57 ++++++++--------
 src/types/cuckoo_filter_sub_filter.h           | 16 ++---
 src/types/redis_cuckoo_chain.cc                | 87 ++++++++++++------------
 src/types/redis_cuckoo_chain.h                 | 21 +++---
 tests/cppunit/types/cuckoo_filter_page_test.cc | 41 ++++++++++++
 tests/cppunit/types/cuckoo_filter_test.cc      | 92 ++++++++++++++++++++++++--
 9 files changed, 287 insertions(+), 99 deletions(-)

diff --git a/src/types/cuckoo_filter.h b/src/types/cuckoo_filter.h
index a885a4c08..5cd02f0a9 100644
--- a/src/types/cuckoo_filter.h
+++ b/src/types/cuckoo_filter.h
@@ -68,8 +68,8 @@ class CuckooFilterHelper {
       return rocksdb::Status::InvalidArgument("bucket_size must be larger than 
0");
     }
 
-    auto max_supported_capacity =
-        static_cast<uint64_t>(kCuckooFilterMaxSupportedBuckets * bucket_size * 
kCuckooFilterLoadFactor);
+    auto max_supported_capacity = static_cast<uint64_t>(
+        static_cast<long double>(kCuckooFilterMaxSupportedBuckets * 
bucket_size) * kCuckooFilterLoadFactor);
     if (capacity > max_supported_capacity) {
       return rocksdb::Status::InvalidArgument("capacity is too large");
     }
diff --git a/src/types/cuckoo_filter_page.cc b/src/types/cuckoo_filter_page.cc
index 1ef4a4c80..4264b00f7 100644
--- a/src/types/cuckoo_filter_page.cc
+++ b/src/types/cuckoo_filter_page.cc
@@ -116,6 +116,47 @@ rocksdb::Status CuckooPageCache::SetBucketSlot(uint16_t 
filter_index, uint32_t n
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status CuckooPageCache::SetBucketSlotWithUndo(uint16_t filter_index, 
uint32_t num_buckets,
+                                                       uint32_t bucket_index, 
uint32_t slot_idx, uint8_t fingerprint,
+                                                       SlotMutation *mutation) 
{
+  if (slot_idx >= bucket_size_) return 
rocksdb::Status::InvalidArgument("invalid cuckoo filter bucket slot");
+
+  BucketRef bucket;
+  auto s = ensureBucketLoaded(filter_index, num_buckets, bucket_index, 
&bucket);
+  if (!s.ok()) return s;
+
+  mutation->filter_index = filter_index;
+  mutation->num_buckets = num_buckets;
+  mutation->bucket_index = bucket_index;
+  mutation->slot_idx = slot_idx;
+  mutation->old_fingerprint = getBucketRefSlot(bucket, slot_idx);
+  mutation->page_was_dirty = bucket.page->is_dirty;
+  setBucketRefSlot(bucket, slot_idx, fingerprint);
+  return rocksdb::Status::OK();
+}
+
+rocksdb::Status CuckooPageCache::RestoreBucketSlot(const SlotMutation 
&mutation) {
+  if (mutation.slot_idx >= bucket_size_) {
+    return rocksdb::Status::Corruption("invalid cuckoo filter slot mutation");
+  }
+
+  BucketLocation location;
+  auto s = resolveBucketLocation(mutation.filter_index, mutation.num_buckets, 
mutation.bucket_index, &location);
+  if (!s.ok()) return s;
+
+  auto it = pages_.find(location.page_key);
+  if (it == pages_.end()) return rocksdb::Status::Corruption("cuckoo filter 
slot mutation page is not cached");
+
+  auto offset = location.offset + mutation.slot_idx;
+  if (offset >= it->second.data.size()) {
+    return rocksdb::Status::Corruption("invalid cuckoo filter slot mutation");
+  }
+
+  it->second.data[offset] = static_cast<char>(mutation.old_fingerprint);
+  it->second.is_dirty = mutation.page_was_dirty;
+  return rocksdb::Status::OK();
+}
+
 rocksdb::Status CuckooPageCache::WriteBackDirtyPages(rocksdb::WriteBatchBase 
*batch) {
   for (const auto &entry : pages_) {
     if (!entry.second.is_dirty) continue;
@@ -125,6 +166,16 @@ rocksdb::Status 
CuckooPageCache::WriteBackDirtyPages(rocksdb::WriteBatchBase *ba
   return rocksdb::Status::OK();
 }
 
+void CuckooPageCache::DiscardCleanPages() {
+  for (auto it = pages_.begin(); it != pages_.end();) {
+    if (it->second.is_dirty) {
+      ++it;
+    } else {
+      it = pages_.erase(it);
+    }
+  }
+}
+
 void CuckooPageCache::DiscardCachedPages() { pages_.clear(); }
 
 rocksdb::Status CuckooPageCache::resolveBucketLocation(uint16_t filter_index, 
uint32_t num_buckets,
diff --git a/src/types/cuckoo_filter_page.h b/src/types/cuckoo_filter_page.h
index ff2efeb50..0012cb82c 100644
--- a/src/types/cuckoo_filter_page.h
+++ b/src/types/cuckoo_filter_page.h
@@ -34,9 +34,22 @@ namespace redis {
 
 class CuckooPageCache {
  public:
+  struct SlotMutation {
+    uint8_t OldFingerprint() const { return old_fingerprint; }
+
+    uint16_t filter_index = 0;
+    uint32_t num_buckets = 0;
+    uint32_t bucket_index = 0;
+    uint32_t slot_idx = 0;
+    uint8_t old_fingerprint = 0;
+    bool page_was_dirty = false;
+  };
+
   CuckooPageCache(engine::Storage *storage, engine::Context &ctx, const Slice 
&ns_key, bool slot_id_encoded,
                   uint64_t version, uint8_t bucket_size, uint32_t page_size);
 
+  uint8_t BucketSize() const { return bucket_size_; }
+
   rocksdb::Status PrefetchBuckets(uint16_t filter_index, uint32_t num_buckets, 
uint32_t bucket1_index,
                                   uint32_t bucket2_index);
   rocksdb::Status TryInsertInBucket(uint16_t filter_index, uint32_t 
num_buckets, uint32_t bucket_index,
@@ -45,8 +58,12 @@ class CuckooPageCache {
                                 uint8_t *fingerprint);
   rocksdb::Status SetBucketSlot(uint16_t filter_index, uint32_t num_buckets, 
uint32_t bucket_index, uint32_t slot_idx,
                                 uint8_t fingerprint);
+  rocksdb::Status SetBucketSlotWithUndo(uint16_t filter_index, uint32_t 
num_buckets, uint32_t bucket_index,
+                                        uint32_t slot_idx, uint8_t 
fingerprint, SlotMutation *mutation);
+  rocksdb::Status RestoreBucketSlot(const SlotMutation &mutation);
   rocksdb::Status WriteBackDirtyPages(rocksdb::WriteBatchBase *batch);
 
+  void DiscardCleanPages();
   void DiscardCachedPages();
 
  private:
diff --git a/src/types/cuckoo_filter_sub_filter.cc 
b/src/types/cuckoo_filter_sub_filter.cc
index 0e631bf2c..e4824b872 100644
--- a/src/types/cuckoo_filter_sub_filter.cc
+++ b/src/types/cuckoo_filter_sub_filter.cc
@@ -20,17 +20,15 @@
 
 #include "cuckoo_filter_sub_filter.h"
 
+#include <vector>
+
 #include "cuckoo_filter.h"
+#include "cuckoo_filter_page.h"
 
 namespace redis {
 
-CuckooSubFilter::CuckooSubFilter(engine::Storage *storage, engine::Context 
&ctx, const Slice &ns_key,
-                                 bool slot_id_encoded, uint64_t version, 
uint8_t bucket_size, uint32_t page_size,
-                                 uint16_t filter_index, uint32_t num_buckets)
-    : bucket_size_(bucket_size),
-      filter_index_(filter_index),
-      num_buckets_(num_buckets),
-      pages_(storage, ctx, ns_key, slot_id_encoded, version, bucket_size, 
page_size) {}
+CuckooSubFilter::CuckooSubFilter(CuckooPageCache &pages, uint16_t 
filter_index, uint32_t num_buckets)
+    : bucket_size_(pages.BucketSize()), filter_index_(filter_index), 
num_buckets_(num_buckets), pages_(pages) {}
 
 rocksdb::Status CuckooSubFilter::TryInsert(uint64_t hash, uint8_t fingerprint, 
bool *inserted) {
   *inserted = false;
@@ -83,20 +81,29 @@ rocksdb::Status CuckooSubFilter::TryKickOutInsert(uint64_t 
hash, uint8_t fingerp
   uint32_t current_bucket_idx = getPrimaryBucketIndex(hash);
   uint8_t current_fp = fingerprint;
   uint32_t victim_slot = 0;
+  std::vector<CuckooPageCache::SlotMutation> mutations;
+  mutations.reserve(max_iterations);
 
-  for (uint16_t iteration = 0; iteration < max_iterations; ++iteration) {
-    uint8_t old_fp = 0;
-    auto s = pages_.GetBucketSlot(filter_index_, num_buckets_, 
current_bucket_idx, victim_slot, &old_fp);
-    if (!s.ok()) {
-      pages_.DiscardCachedPages();
-      return s;
-    }
-    s = pages_.SetBucketSlot(filter_index_, num_buckets_, current_bucket_idx, 
victim_slot, current_fp);
-    if (!s.ok()) {
-      pages_.DiscardCachedPages();
-      return s;
+  auto rollback = [&]() {
+    for (auto it = mutations.rbegin(); it != mutations.rend(); ++it) {
+      auto s = pages_.RestoreBucketSlot(*it);
+      if (!s.ok()) return s;
     }
-    current_fp = old_fp;
+    return rocksdb::Status::OK();
+  };
+
+  auto rollback_and_return = [&](const rocksdb::Status &status) {
+    auto rollback_status = rollback();
+    return rollback_status.ok() ? status : rollback_status;
+  };
+
+  for (uint16_t iteration = 0; iteration < max_iterations; ++iteration) {
+    CuckooPageCache::SlotMutation mutation;
+    auto s = pages_.SetBucketSlotWithUndo(filter_index_, num_buckets_, 
current_bucket_idx, victim_slot, current_fp,
+                                          &mutation);
+    if (!s.ok()) return rollback_and_return(s);
+    current_fp = mutation.OldFingerprint();
+    mutations.push_back(mutation);
 
     if (current_fp == 0) {
       *inserted = true;
@@ -107,10 +114,7 @@ rocksdb::Status CuckooSubFilter::TryKickOutInsert(uint64_t 
hash, uint8_t fingerp
 
     bool inserted_in_alt_bucket = false;
     s = pages_.TryInsertInBucket(filter_index_, num_buckets_, alt_bucket_idx, 
current_fp, &inserted_in_alt_bucket);
-    if (!s.ok()) {
-      pages_.DiscardCachedPages();
-      return s;
-    }
+    if (!s.ok()) return rollback_and_return(s);
     if (inserted_in_alt_bucket) {
       *inserted = true;
       return rocksdb::Status::OK();
@@ -120,12 +124,7 @@ rocksdb::Status CuckooSubFilter::TryKickOutInsert(uint64_t 
hash, uint8_t fingerp
     victim_slot = (victim_slot + 1) % bucket_size_;
   }
 
-  pages_.DiscardCachedPages();
-  return rocksdb::Status::OK();
-}
-
-rocksdb::Status CuckooSubFilter::WriteToBatch(rocksdb::WriteBatchBase *batch) {
-  return pages_.WriteBackDirtyPages(batch);
+  return rollback();
 }
 
 uint32_t CuckooSubFilter::getPrimaryBucketIndex(uint64_t hash) const { return 
hash % num_buckets_; }
diff --git a/src/types/cuckoo_filter_sub_filter.h 
b/src/types/cuckoo_filter_sub_filter.h
index d3ca1f170..27d968e3e 100644
--- a/src/types/cuckoo_filter_sub_filter.h
+++ b/src/types/cuckoo_filter_sub_filter.h
@@ -21,29 +21,25 @@
 #pragma once
 
 #include <rocksdb/status.h>
-#include <rocksdb/write_batch.h>
 
 #include <cstdint>
 
-#include "cuckoo_filter_page.h"
-
 namespace redis {
 
+class CuckooPageCache;
+
 class CuckooSubFilter {
  public:
-  CuckooSubFilter(engine::Storage *storage, engine::Context &ctx, const Slice 
&ns_key, bool slot_id_encoded,
-                  uint64_t version, uint8_t bucket_size, uint32_t page_size, 
uint16_t filter_index,
-                  uint32_t num_buckets);
+  CuckooSubFilter(CuckooPageCache &pages, uint16_t filter_index, uint32_t 
num_buckets);
 
   uint16_t Index() const { return filter_index_; }
   uint32_t NumBuckets() const { return num_buckets_; }
 
   rocksdb::Status TryInsert(uint64_t hash, uint8_t fingerprint, bool 
*inserted);
   rocksdb::Status Contains(uint64_t hash, uint8_t fingerprint, bool *exists);
-  // Performs speculative kick-out mutations in the page cache. On success, 
dirty pages remain staged for
-  // WriteToBatch(); on inserted=false or non-OK status, cached pages are 
discarded before returning.
+  // Performs speculative kick-out mutations in the page cache. Failed 
attempts restore only the slots changed by
+  // this call, preserving mutations staged by the owning chain operation.
   rocksdb::Status TryKickOutInsert(uint64_t hash, uint8_t fingerprint, 
uint16_t max_iterations, bool *inserted);
-  rocksdb::Status WriteToBatch(rocksdb::WriteBatchBase *batch);
 
  private:
   uint32_t getPrimaryBucketIndex(uint64_t hash) const;
@@ -52,7 +48,7 @@ class CuckooSubFilter {
   uint8_t bucket_size_ = 0;
   uint16_t filter_index_ = 0;
   uint32_t num_buckets_ = 0;
-  CuckooPageCache pages_;
+  CuckooPageCache &pages_;
 };
 
 }  // namespace redis
diff --git a/src/types/redis_cuckoo_chain.cc b/src/types/redis_cuckoo_chain.cc
index ef74d17b0..cbb724e3f 100644
--- a/src/types/redis_cuckoo_chain.cc
+++ b/src/types/redis_cuckoo_chain.cc
@@ -21,6 +21,7 @@
 #include "redis_cuckoo_chain.h"
 
 #include "cuckoo_filter.h"
+#include "cuckoo_filter_page.h"
 #include "cuckoo_filter_sub_filter.h"
 #include "logging.h"
 
@@ -151,37 +152,45 @@ rocksdb::Status CuckooChain::Add(engine::Context &ctx, 
const Slice &user_key, co
   // Calculate hash and fingerprint for the item
   uint64_t hash = CuckooFilterHelper::Hash(item.data(), item.size());
   uint8_t fingerprint = CuckooFilterHelper::GenerateFingerprint(hash);
+  CuckooPageCache pages(storage_, ctx, ns_key, storage_->IsSlotIdEncoded(), 
metadata.version, metadata.bucket_size,
+                        metadata.page_size);
 
   bool inserted = false;
-  s = tryCuckooInsert(ctx, user_key, ns_key, &metadata, hash, fingerprint, 
&inserted);
+  s = tryInsertFingerprint(&metadata, pages, hash, fingerprint, &inserted);
   if (!s.ok()) return s;
-  if (inserted) {
-    *added = true;
-    return rocksdb::Status::OK();
+  if (!inserted) {
+    *added = false;
+    return rocksdb::Status::Aborted("filter is full");
   }
 
-  s = tryCuckooKickOut(ctx, user_key, ns_key, &metadata, hash, fingerprint, 
&inserted);
+  s = commitPagesAndMetadata(ctx, user_key, ns_key, &metadata, pages);
   if (!s.ok()) return s;
-  if (inserted) {
-    *added = true;
-    return rocksdb::Status::OK();
-  }
 
-  s = expandAndInsertCuckooChain(ctx, user_key, ns_key, &metadata, hash, 
fingerprint, &inserted);
+  *added = true;
+  return rocksdb::Status::OK();
+}
+
+rocksdb::Status CuckooChain::tryInsertFingerprint(CuckooChainMetadata 
*metadata, CuckooPageCache &pages, uint64_t hash,
+                                                  uint8_t fingerprint, bool 
*inserted) {
+  auto s = tryCuckooInsert(metadata, pages, hash, fingerprint, inserted);
   if (!s.ok()) return s;
-  if (inserted) {
-    *added = true;
-    return rocksdb::Status::OK();
+
+  if (!*inserted) {
+    s = tryCuckooKickOut(metadata, pages, hash, fingerprint, inserted);
+    if (!s.ok()) return s;
   }
 
-  // No expansion allowed and filter is full
-  *added = false;
-  return rocksdb::Status::Aborted("filter is full");
+  if (!*inserted) {
+    s = expandAndInsertCuckooChain(metadata, pages, hash, fingerprint, 
inserted);
+    if (!s.ok()) return s;
+  }
+
+  if (*inserted) metadata->size++;
+  return rocksdb::Status::OK();
 }
 
-rocksdb::Status CuckooChain::tryCuckooInsert(engine::Context &ctx, const Slice 
&user_key, const std::string &ns_key,
-                                             CuckooChainMetadata *metadata, 
uint64_t hash, uint8_t fingerprint,
-                                             bool *inserted) {
+rocksdb::Status CuckooChain::tryCuckooInsert(const CuckooChainMetadata 
*metadata, CuckooPageCache &pages, uint64_t hash,
+                                             uint8_t fingerprint, bool 
*inserted) {
   *inserted = false;
 
   // RedisBloom prioritizes the newest sub-filter to avoid repeatedly probing 
older, fuller filters.
@@ -192,26 +201,23 @@ rocksdb::Status 
CuckooChain::tryCuckooInsert(engine::Context &ctx, const Slice &
                                                      metadata->bucket_size, 
current_filter_idx, &num_buckets);
     if (!s.ok()) return s;
 
-    CuckooSubFilter sub_filter(storage_, ctx, ns_key, 
storage_->IsSlotIdEncoded(), metadata->version,
-                               metadata->bucket_size, metadata->page_size, 
current_filter_idx, num_buckets);
+    CuckooSubFilter sub_filter(pages, current_filter_idx, num_buckets);
     bool current_inserted = false;
     s = sub_filter.TryInsert(hash, fingerprint, &current_inserted);
     if (!s.ok()) return s;
 
     if (current_inserted) {
-      s = commitSubFilterAndMetadata(ctx, user_key, ns_key, metadata, 
&sub_filter);
-      if (!s.ok()) return s;
       *inserted = true;
       return rocksdb::Status::OK();
     }
+    pages.DiscardCleanPages();
   }
 
   return rocksdb::Status::OK();
 }
 
-rocksdb::Status CuckooChain::tryCuckooKickOut(engine::Context &ctx, const 
Slice &user_key, const std::string &ns_key,
-                                              CuckooChainMetadata *metadata, 
uint64_t hash, uint8_t fingerprint,
-                                              bool *inserted) {
+rocksdb::Status CuckooChain::tryCuckooKickOut(const CuckooChainMetadata 
*metadata, CuckooPageCache &pages,
+                                              uint64_t hash, uint8_t 
fingerprint, bool *inserted) {
   *inserted = false;
 
   // No space found in any filter, try kick-out on the last filter
@@ -221,14 +227,11 @@ rocksdb::Status 
CuckooChain::tryCuckooKickOut(engine::Context &ctx, const Slice
                                                    last_filter_idx, 
&num_buckets);
   if (!s.ok()) return s;
 
-  CuckooSubFilter last_filter(storage_, ctx, ns_key, 
storage_->IsSlotIdEncoded(), metadata->version,
-                              metadata->bucket_size, metadata->page_size, 
last_filter_idx, num_buckets);
+  CuckooSubFilter last_filter(pages, last_filter_idx, num_buckets);
   bool kickout_inserted = false;
   s = last_filter.TryKickOutInsert(hash, fingerprint, 
metadata->max_iterations, &kickout_inserted);
   if (!s.ok()) return s;
   if (kickout_inserted) {
-    s = commitSubFilterAndMetadata(ctx, user_key, ns_key, metadata, 
&last_filter);
-    if (!s.ok()) return s;
     *inserted = true;
     return rocksdb::Status::OK();
   }
@@ -236,8 +239,7 @@ rocksdb::Status 
CuckooChain::tryCuckooKickOut(engine::Context &ctx, const Slice
   return rocksdb::Status::OK();
 }
 
-rocksdb::Status CuckooChain::expandAndInsertCuckooChain(engine::Context &ctx, 
const Slice &user_key,
-                                                        const std::string 
&ns_key, CuckooChainMetadata *metadata,
+rocksdb::Status CuckooChain::expandAndInsertCuckooChain(CuckooChainMetadata 
*metadata, CuckooPageCache &pages,
                                                         uint64_t hash, uint8_t 
fingerprint, bool *inserted) {
   *inserted = false;
 
@@ -256,33 +258,28 @@ rocksdb::Status 
CuckooChain::expandAndInsertCuckooChain(engine::Context &ctx, co
   }
   if (!s.ok()) return s;
 
-  CuckooSubFilter new_filter(storage_, ctx, ns_key, 
storage_->IsSlotIdEncoded(), metadata->version,
-                             metadata->bucket_size, metadata->page_size, 
new_filter_idx, new_num_buckets);
+  CuckooSubFilter new_filter(pages, new_filter_idx, new_num_buckets);
   bool new_filter_inserted = false;
   s = new_filter.TryInsert(hash, fingerprint, &new_filter_inserted);
   if (!s.ok()) return s;
   if (!new_filter_inserted) return rocksdb::Status::Corruption("failed to 
insert into new cuckoo filter");
 
   metadata->n_filters++;
-  s = commitSubFilterAndMetadata(ctx, user_key, ns_key, metadata, &new_filter);
-  if (!s.ok()) return s;
-
   *inserted = true;
   return rocksdb::Status::OK();
 }
 
-rocksdb::Status CuckooChain::commitSubFilterAndMetadata(engine::Context &ctx, 
const Slice &user_key,
-                                                        const std::string 
&ns_key, CuckooChainMetadata *metadata,
-                                                        CuckooSubFilter 
*sub_filter) {
+rocksdb::Status CuckooChain::commitPagesAndMetadata(engine::Context &ctx, 
const Slice &user_key,
+                                                    const std::string &ns_key, 
const CuckooChainMetadata *metadata,
+                                                    CuckooPageCache &pages) {
   auto batch = storage_->GetWriteBatchBase();
   WriteBatchLogData log_data(kRedisCuckooFilter, 
std::vector<std::string>{"add", user_key.ToString()});
   auto s = batch->PutLogData(log_data.Encode());
   if (!s.ok()) return s;
 
-  s = sub_filter->WriteToBatch(batch.Get());
+  s = pages.WriteBackDirtyPages(batch.Get());
   if (!s.ok()) return s;
 
-  metadata->size++;
   std::string metadata_bytes;
   metadata->Encode(&metadata_bytes);
   s = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes);
@@ -321,6 +318,8 @@ rocksdb::Status CuckooChain::MExists(engine::Context &ctx, 
const Slice &user_key
   }
 
   size_t found_count = 0;
+  CuckooPageCache pages(storage_, ctx, ns_key, storage_->IsSlotIdEncoded(), 
metadata.version, metadata.bucket_size,
+                        metadata.page_size);
   for (int filter_idx = static_cast<int>(metadata.n_filters) - 1; filter_idx 
>= 0; --filter_idx) {
     auto current_filter_idx = static_cast<uint16_t>(filter_idx);
     uint32_t num_buckets = 0;
@@ -328,8 +327,7 @@ rocksdb::Status CuckooChain::MExists(engine::Context &ctx, 
const Slice &user_key
                                                 current_filter_idx, 
&num_buckets);
     if (!s.ok()) return s;
 
-    CuckooSubFilter sub_filter(storage_, ctx, ns_key, 
storage_->IsSlotIdEncoded(), metadata.version,
-                               metadata.bucket_size, metadata.page_size, 
current_filter_idx, num_buckets);
+    CuckooSubFilter sub_filter(pages, current_filter_idx, num_buckets);
     for (size_t i = 0; i < items.size(); ++i) {
       if ((*exists)[i]) continue;
 
@@ -345,6 +343,7 @@ rocksdb::Status CuckooChain::MExists(engine::Context &ctx, 
const Slice &user_key
     if (found_count == items.size()) {
       break;
     }
+    pages.DiscardCleanPages();
   }
 
   return rocksdb::Status::OK();
diff --git a/src/types/redis_cuckoo_chain.h b/src/types/redis_cuckoo_chain.h
index d3a4381e5..18eae088e 100644
--- a/src/types/redis_cuckoo_chain.h
+++ b/src/types/redis_cuckoo_chain.h
@@ -28,7 +28,7 @@
 
 namespace redis {
 
-class CuckooSubFilter;
+class CuckooPageCache;
 
 // Default values for a newly created cuckoo filter.
 const uint32_t kCFDefaultCapacity = 1024;
@@ -62,15 +62,16 @@ class CuckooChain : public Database {
 
   static rocksdb::Status validateMetadata(const CuckooChainMetadata &metadata);
 
-  rocksdb::Status tryCuckooInsert(engine::Context &ctx, const Slice &user_key, 
const std::string &ns_key,
-                                  CuckooChainMetadata *metadata, uint64_t 
hash, uint8_t fingerprint, bool *inserted);
-  rocksdb::Status tryCuckooKickOut(engine::Context &ctx, const Slice 
&user_key, const std::string &ns_key,
-                                   CuckooChainMetadata *metadata, uint64_t 
hash, uint8_t fingerprint, bool *inserted);
-  rocksdb::Status expandAndInsertCuckooChain(engine::Context &ctx, const Slice 
&user_key, const std::string &ns_key,
-                                             CuckooChainMetadata *metadata, 
uint64_t hash, uint8_t fingerprint,
-                                             bool *inserted);
-  rocksdb::Status commitSubFilterAndMetadata(engine::Context &ctx, const Slice 
&user_key, const std::string &ns_key,
-                                             CuckooChainMetadata *metadata, 
CuckooSubFilter *sub_filter);
+  static rocksdb::Status tryInsertFingerprint(CuckooChainMetadata *metadata, 
CuckooPageCache &pages, uint64_t hash,
+                                              uint8_t fingerprint, bool 
*inserted);
+  static rocksdb::Status tryCuckooInsert(const CuckooChainMetadata *metadata, 
CuckooPageCache &pages, uint64_t hash,
+                                         uint8_t fingerprint, bool *inserted);
+  static rocksdb::Status tryCuckooKickOut(const CuckooChainMetadata *metadata, 
CuckooPageCache &pages, uint64_t hash,
+                                          uint8_t fingerprint, bool *inserted);
+  static rocksdb::Status expandAndInsertCuckooChain(CuckooChainMetadata 
*metadata, CuckooPageCache &pages,
+                                                    uint64_t hash, uint8_t 
fingerprint, bool *inserted);
+  rocksdb::Status commitPagesAndMetadata(engine::Context &ctx, const Slice 
&user_key, const std::string &ns_key,
+                                         const CuckooChainMetadata *metadata, 
CuckooPageCache &pages);
 };
 
 }  // namespace redis
diff --git a/tests/cppunit/types/cuckoo_filter_page_test.cc 
b/tests/cppunit/types/cuckoo_filter_page_test.cc
index 8be87aefa..099159220 100644
--- a/tests/cppunit/types/cuckoo_filter_page_test.cc
+++ b/tests/cppunit/types/cuckoo_filter_page_test.cc
@@ -343,6 +343,47 @@ TEST_F(RedisCuckooPageCacheTest, 
DiscardCachedPagesDropsDirtyPages) {
   EXPECT_TRUE(s.IsNotFound()) << s.ToString();
 }
 
+TEST_F(RedisCuckooPageCacheTest, DiscardCleanPagesPreservesDirtyPages) {
+  auto metadata = makeMetadata(4);
+  auto dirty_page_key = makePageKey(metadata, 0, 0);
+  auto clean_page_key = makePageKey(metadata, 1, 0);
+  writePage(dirty_page_key, std::string(8, static_cast<char>(1)));
+  writePage(clean_page_key, std::string(8, static_cast<char>(2)));
+  redis::CuckooPageCache pages(storage_.get(), *ctx_, ns_key_, 
storage_->IsSlotIdEncoded(), metadata.version,
+                               metadata.bucket_size, metadata.page_size);
+
+  auto s = pages.SetBucketSlot(0, 2, 0, 0, 11);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  uint8_t fingerprint = 0;
+  s = pages.GetBucketSlot(1, 2, 0, 0, &fingerprint);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  ASSERT_EQ(fingerprint, 2);
+
+  pages.DiscardCleanPages();
+  writePage(dirty_page_key, std::string(8, static_cast<char>(3)));
+  writePage(clean_page_key, std::string(8, static_cast<char>(4)));
+
+  s = pages.GetBucketSlot(0, 2, 0, 0, &fingerprint);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(fingerprint, 11);
+  s = pages.GetBucketSlot(1, 2, 0, 0, &fingerprint);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(fingerprint, 4);
+
+  auto batch = storage_->GetWriteBatchBase();
+  s = pages.WriteBackDirtyPages(batch.Get());
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  commitBatch(batch.Get());
+
+  std::string page;
+  s = readPage(dirty_page_key, &page);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(static_cast<uint8_t>(page[0]), 11);
+  s = readPage(clean_page_key, &page);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(static_cast<uint8_t>(page[0]), 4);
+}
+
 TEST_F(RedisCuckooPageCacheTest, PageKeyUsesMetadataVersion) {
   auto old_metadata = makeMetadata(4, 100);
   auto new_metadata = makeMetadata(4, 101);
diff --git a/tests/cppunit/types/cuckoo_filter_test.cc 
b/tests/cppunit/types/cuckoo_filter_test.cc
index 8f92263b3..9fb5fafcd 100644
--- a/tests/cppunit/types/cuckoo_filter_test.cc
+++ b/tests/cppunit/types/cuckoo_filter_test.cc
@@ -665,7 +665,90 @@ TEST_F(RedisCuckooFilterTest, 
KickOutSuccessWritesDirtyPages) {
   EXPECT_EQ(static_cast<uint8_t>(page[evicted_bucket]), first.fingerprint);
 }
 
-TEST_F(RedisCuckooFilterTest, KickOutErrorDiscardsDirtyPages) {
+TEST_F(RedisCuckooFilterTest, 
FailedMultiStepKickOutPreservesEarlierCachedMutation) {
+  constexpr uint8_t bucket_size = 1;
+  constexpr uint32_t page_size = 2;
+  constexpr uint32_t num_buckets = 2;
+  constexpr uint64_t version = 1;
+
+  CuckooChainMetadata metadata(false);
+  metadata.version = version;
+  std::string original_page{1, 1};
+  writePage(makePageKey(key_, metadata, 0, 0), original_page);
+
+  redis::CuckooPageCache page_cache(storage_.get(), *ctx_, 
db_->AppendNamespacePrefix(key_),
+                                    storage_->IsSlotIdEncoded(), version, 
bucket_size, page_size);
+  redis::CuckooSubFilter earlier_sub_filter(page_cache, 1, num_buckets);
+
+  bool inserted = false;
+  auto s = earlier_sub_filter.TryInsert(0, 9, &inserted);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  ASSERT_TRUE(inserted);
+
+  ASSERT_EQ(redis::CuckooFilterHelper::GetAltBucketIndex(0, 1, num_buckets), 
1);
+  ASSERT_EQ(redis::CuckooFilterHelper::GetAltBucketIndex(1, 1, num_buckets), 
0);
+  ASSERT_EQ(redis::CuckooFilterHelper::GetAltBucketIndex(0, 2, num_buckets), 
0);
+  redis::CuckooSubFilter full_sub_filter(page_cache, 0, num_buckets);
+  s = full_sub_filter.TryKickOutInsert(0, 2, 3, &inserted);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  ASSERT_FALSE(inserted);
+
+  auto batch = storage_->GetWriteBatchBase();
+  s = page_cache.WriteBackDirtyPages(batch.Get());
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(batch->GetWriteBatch()->Count(), 1);
+  s = storage_->Write(*ctx_, storage_->DefaultWriteOptions(), 
batch->GetWriteBatch());
+  ASSERT_TRUE(s.ok()) << s.ToString();
+
+  std::string restored_page;
+  s = readPage(makePageKey(key_, metadata, 0, 0), &restored_page);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(restored_page, original_page);
+
+  std::string earlier_page;
+  s = readPage(makePageKey(key_, metadata, 1, 0), &earlier_page);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(earlier_page, std::string({9, 0}));
+}
+
+TEST_F(RedisCuckooFilterTest, 
FailedMultiStepKickOutPreservesEarlierMutationOnSamePage) {
+  constexpr uint8_t bucket_size = 1;
+  constexpr uint32_t page_size = 2;
+  constexpr uint32_t num_buckets = 2;
+  constexpr uint64_t version = 1;
+
+  CuckooChainMetadata metadata(false);
+  metadata.version = version;
+  writePage(makePageKey(key_, metadata, 0, 0), std::string({1, 1}));
+
+  redis::CuckooPageCache page_cache(storage_.get(), *ctx_, 
db_->AppendNamespacePrefix(key_),
+                                    storage_->IsSlotIdEncoded(), version, 
bucket_size, page_size);
+  auto s = page_cache.SetBucketSlot(0, num_buckets, 0, 0, 3);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+
+  ASSERT_EQ(redis::CuckooFilterHelper::GetAltBucketIndex(0, 3, num_buckets), 
1);
+  ASSERT_EQ(redis::CuckooFilterHelper::GetAltBucketIndex(1, 1, num_buckets), 
0);
+  ASSERT_EQ(redis::CuckooFilterHelper::GetAltBucketIndex(0, 2, num_buckets), 
0);
+  redis::CuckooSubFilter full_sub_filter(page_cache, 0, num_buckets);
+  bool inserted = true;
+  s = full_sub_filter.TryKickOutInsert(0, 2, 3, &inserted);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  ASSERT_FALSE(inserted);
+
+  auto batch = storage_->GetWriteBatchBase();
+  s = page_cache.WriteBackDirtyPages(batch.Get());
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(batch->GetWriteBatch()->Count(), 1);
+  s = storage_->Write(*ctx_, storage_->DefaultWriteOptions(), 
batch->GetWriteBatch());
+  ASSERT_TRUE(s.ok()) << s.ToString();
+
+  std::string page;
+  s = readPage(makePageKey(key_, metadata, 0, 0), &page);
+  ASSERT_TRUE(s.ok()) << s.ToString();
+  EXPECT_EQ(page, std::string({3, 1}));
+}
+
+TEST_F(RedisCuckooFilterTest, KickOutErrorRestoresDirtyPages) {
   constexpr uint8_t bucket_size = 1;
   constexpr uint32_t page_size = 1;
   constexpr uint32_t num_buckets = 2;
@@ -689,16 +772,17 @@ TEST_F(RedisCuckooFilterTest, 
KickOutErrorDiscardsDirtyPages) {
   writePage(makePageKey(key_, metadata, 0, 0), original_page);
   writePage(makePageKey(key_, metadata, 0, 1), std::string(2, 
static_cast<char>(9)));
 
-  redis::CuckooSubFilter sub_filter(storage_.get(), *ctx_, 
db_->AppendNamespacePrefix(key_),
+  redis::CuckooPageCache page_cache(storage_.get(), *ctx_, 
db_->AppendNamespacePrefix(key_),
                                     storage_->IsSlotIdEncoded(), 
metadata.version, metadata.bucket_size,
-                                    metadata.page_size, 0, num_buckets);
+                                    metadata.page_size);
+  redis::CuckooSubFilter sub_filter(page_cache, 0, num_buckets);
   bool inserted = true;
   auto s = sub_filter.TryKickOutInsert(hash, fingerprint, 
metadata.max_iterations, &inserted);
   ASSERT_TRUE(s.IsCorruption()) << s.ToString();
   ASSERT_FALSE(inserted);
 
   auto batch = storage_->GetWriteBatchBase();
-  s = sub_filter.WriteToBatch(batch.Get());
+  s = page_cache.WriteBackDirtyPages(batch.Get());
   ASSERT_TRUE(s.ok()) << s.ToString();
   s = storage_->Write(*ctx_, storage_->DefaultWriteOptions(), 
batch->GetWriteBatch());
   ASSERT_TRUE(s.ok()) << s.ToString();

Reply via email to