bug(dense set): fix scan function to return only home bucket data (#474)

dense set: fix scan function to return only home bucket data

Signed-off-by: adi_holden <adi@dragonflydb.io>

src/core/dense_set.cc

@@ -26,8 +26,7 @@ constexpr size_t kMinSize = 1 << kMinSizeShift;
 constexpr bool kAllowDisplacements = true;
 
 DenseSet::IteratorBase::IteratorBase(const DenseSet* owner, bool is_end)
-    : owner_(const_cast<DenseSet&>(*owner)),
+    : owner_(const_cast<DenseSet&>(*owner)), curr_entry_(nullptr) {
-      curr_entry_(nullptr) {
   curr_list_ = is_end ? owner_.entries_.end() : owner_.entries_.begin();
   if (curr_list_ != owner->entries_.end()) {
     curr_entry_ = &(*curr_list_);
@@ -173,6 +172,32 @@ bool DenseSet::Equal(DensePtr dptr, const void* ptr, uint32_t cookie) const {
   return ObjEqual(dptr.GetObject(), ptr, cookie);
 }
 
+bool DenseSet::NoItemBelongsBucket(uint32_t bid) const {
+  auto& entries = const_cast<DenseSet*>(this)->entries_;
+  DensePtr* curr = &entries[bid];
+  ExpireIfNeeded(nullptr, curr);
+  if (!curr->IsEmpty() && !curr->IsDisplaced()) {
+    return false;
+  }
+
+  if (bid + 1 < entries_.size()) {
+    DensePtr* right_bucket = &entries[bid + 1];
+    ExpireIfNeeded(nullptr, right_bucket);
+    if (!right_bucket->IsEmpty() && right_bucket->IsDisplaced() &&
+        right_bucket->GetDisplacedDirection() == 1)
+      return false;
+  }
+
+  if (bid > 0) {
+    DensePtr* left_bucket = &entries[bid - 1];
+    ExpireIfNeeded(nullptr, left_bucket);
+    if (!left_bucket->IsEmpty() && left_bucket->IsDisplaced() &&
+        left_bucket->GetDisplacedDirection() == -1)
+      return false;
+  }
+  return true;
+}
+
 auto DenseSet::FindEmptyAround(uint32_t bid) -> ChainVectorIterator {
   ExpireIfNeeded(nullptr, &entries_[bid]);
 
@@ -494,33 +519,45 @@ uint32_t DenseSet::Scan(uint32_t cursor, const ItemCb& cb) const {
 
   auto& entries = const_cast<DenseSet*>(this)->entries_;
 
-  // skip empty entries
+  // First find the bucket to scan, skip empty buckets.
-  do {
+  // A bucket is empty if the current index is empty and the data is not displaced
-    while (entries_idx < entries_.size() && entries_[entries_idx].IsEmpty()) {
+  // to the right or to the left.
-      ++entries_idx;
+  while (entries_idx < entries_.size() && NoItemBelongsBucket(entries_idx)) {
-    }
+    ++entries_idx;
+  }
 
-    if (entries_idx == entries_.size()) {
+  if (entries_idx == entries_.size()) {
-      return 0;
+    return 0;
-    }
+  }
-
-    ExpireIfNeeded(nullptr, &entries[entries_idx]);
-  } while (entries_[entries_idx].IsEmpty());
 
   DensePtr* curr = &entries[entries_idx];
 
-  // when scanning add all entries in a given chain
+  // Check home bucket
-  while (true) {
+  if (!curr->IsEmpty() && !curr->IsDisplaced()) {
-    cb(curr->GetObject());
+    // scanning add all entries in a given chain
-    if (!curr->IsLink())
+    while (true) {
-      break;
+      cb(curr->GetObject());
+      if (!curr->IsLink())
+        break;
 
-    DensePtr* mcurr = const_cast<DensePtr*>(curr);
+      DensePtr* mcurr = const_cast<DensePtr*>(curr);
 
-    if (ExpireIfNeeded(mcurr, &mcurr->AsLink()->next) && !mcurr->IsLink()) {
+      if (ExpireIfNeeded(mcurr, &mcurr->AsLink()->next) && !mcurr->IsLink()) {
-      break;
+        break;
+      }
+      curr = &curr->AsLink()->next;
+    }
+  }
+
+  // Check if the bucket on the left belongs to the home bucket.
+  if (entries_idx > 0) {
+    DensePtr* left_bucket = &entries[entries_idx - 1];
+    ExpireIfNeeded(nullptr, left_bucket);
+
+    if (left_bucket->IsDisplaced() &&
+        left_bucket->GetDisplacedDirection() == -1) {  // left of the home bucket
+      cb(left_bucket->GetObject());
     }
-    curr = &curr->AsLink()->next;
   }
 
   // move to the next index for the next scan and check if we are done
@@ -529,12 +566,13 @@ uint32_t DenseSet::Scan(uint32_t cursor, const ItemCb& cb) const {
     return 0;
   }
 
-  // In case of displacement, we want to fully cover the bucket we traversed, therefore
+  // Check if the bucket on the right belongs to the home bucket.
-  // we check if the bucket on the right belongs to the home bucket.
+  DensePtr* right_bucket = &entries[entries_idx];
-  ExpireIfNeeded(nullptr, &entries[entries_idx]);
+  ExpireIfNeeded(nullptr, right_bucket);
 
-  if (entries[entries_idx].GetDisplacedDirection() == 1) { // right of the home bucket
+  if (right_bucket->IsDisplaced() &&
-    cb(entries[entries_idx].GetObject());
+      right_bucket->GetDisplacedDirection() == 1) {  // right of the home bucket
+    cb(right_bucket->GetObject());
   }
 
   return entries_idx << (32 - capacity_log_);

src/core/dense_set.h

@@ -359,6 +359,9 @@ class DenseSet {
 
   // return a ChainVectorIterator (a.k.a iterator) or end if there is an empty chain found
   ChainVectorIterator FindEmptyAround(uint32_t bid);
+  // return if bucket has no item which is not displaced and right/left bucket has no displaced item
+  // belong to given bid
+  bool NoItemBelongsBucket(uint32_t bid) const;
   void Grow();
 
   // ============ Pseudo Linked List Functions for interacting with Chains ==================

src/server/set_family_test.cc

@@ -192,9 +192,14 @@ TEST_F(SetFamilyTest, SScan) {
   vec = StrArray(resp.GetVec()[1]);
   EXPECT_THAT(vec.size(), 5);
 
+  resp = Run({"sscan", "mystrset", "0", "match", "str-1*"});
+  vec = StrArray(resp.GetVec()[1]);
+  EXPECT_THAT(vec, UnorderedElementsAre("str-1", "str-10", "str-11", "str-12", "str-13", "str-14"));
+
   resp = Run({"sscan", "mystrset", "0", "match", "str-1*", "count", "3"});
   vec = StrArray(resp.GetVec()[1]);
   EXPECT_THAT(vec, IsSubsetOf({"str-1", "str-10", "str-11", "str-12", "str-13", "str-14"}));
+  EXPECT_EQ(vec.size(), 3);
 
   // nothing should match this
   resp = Run({"sscan", "mystrset", "0", "match", "1*"});