feat: allow extracting the expiry time from string_set (#2020)

We introduce `iterator Find(member)` function
as well as iterator members to actually get the expiry time.

src/core/dense_set.cc

@@ -410,38 +410,49 @@ void DenseSet::AddUnique(void* obj, bool has_ttl, uint64_t hashcode) {
   ++size_;
 }
 
-auto DenseSet::Find(const void* ptr, uint32_t bid, uint32_t cookie) -> pair<DensePtr*, DensePtr*> {
+auto DenseSet::Find2(const void* ptr, uint32_t bid, uint32_t cookie)
-  // could do it with zigzag decoding but this is clearer.
+    -> tuple<size_t, DensePtr*, DensePtr*> {
-  int offset[] = {0, -1, 1};
+  DensePtr* curr = &entries_[bid];
+  ExpireIfNeeded(nullptr, curr);
+
+  if (Equal(*curr, ptr, cookie)) {
+    return {bid, nullptr, curr};
+  }
 
   // first look for displaced nodes since this is quicker than iterating a potential long chain
-  for (int j = 0; j < 3; ++j) {
+  if (bid > 0) {
-    if ((bid == 0 && j == 1) || (bid + 1 == entries_.size() && j == 2))
+    curr = &entries_[bid - 1];
-      continue;
-
-    DensePtr* curr = &entries_[bid + offset[j]];
-
     ExpireIfNeeded(nullptr, curr);
+
     if (Equal(*curr, ptr, cookie)) {
-      return make_pair(nullptr, curr);
+      return {bid - 1, nullptr, curr};
+    }
+  }
+
+  if (bid + 1 < entries_.size()) {
+    curr = &entries_[bid + 1];
+    ExpireIfNeeded(nullptr, curr);
+
+    if (Equal(*curr, ptr, cookie)) {
+      return {bid + 1, nullptr, curr};
     }
   }
 
   // if the node is not displaced, search the correct chain
   DensePtr* prev = &entries_[bid];
-  DensePtr* curr = prev->Next();
+  curr = prev->Next();
   while (curr != nullptr) {
     ExpireIfNeeded(prev, curr);
 
     if (Equal(*curr, ptr, cookie)) {
-      return make_pair(prev, curr);
+      return {bid, prev, curr};
     }
     prev = curr;
     curr = curr->Next();
   }
 
   // not in the Set
-  return make_pair(nullptr, nullptr);
+  return {0, nullptr, nullptr};
 }
 
 void DenseSet::Delete(DensePtr* prev, DensePtr* ptr) {

src/core/dense_set.h

@@ -174,6 +174,22 @@ class DenseSet {
   using ChainVectorConstIterator = std::vector<DensePtr, DensePtrAllocator>::const_iterator;
 
   class IteratorBase {
+    friend class DenseSet;
+
+   public:
+    IteratorBase(DenseSet* owner, ChainVectorIterator list_it, DensePtr* e)
+        : owner_(owner), curr_list_(list_it), curr_entry_(e) {
+    }
+
+    // returns the expiry time of the current entry or UINT32_MAX if no ttl is set.
+    uint32_t ExpiryTime() const {
+      return curr_entry_->HasTtl() ? owner_->ObjExpireTime(curr_entry_->GetObject()) : UINT32_MAX;
+    }
+
+    bool HasExpiry() const {
+      return curr_entry_->HasTtl();
+    }
+
    protected:
     IteratorBase() : owner_(nullptr), curr_entry_(nullptr) {
     }
@@ -254,6 +270,15 @@ class DenseSet {
   }
 
   void* FindInternal(const void* obj, uint64_t hashcode, uint32_t cookie) const;
+
+  IteratorBase FindIt(const void* ptr, uint32_t cookie) {
+    auto [bid, _, curr] = Find2(ptr, BucketId(ptr, cookie), cookie);
+    if (curr) {
+      return IteratorBase(this, entries_.begin() + bid, curr);
+    }
+    return IteratorBase{};
+  }
+
   void* PopInternal();
 
   // Note this does not free any dynamic allocations done by derived classes, that a DensePtr
@@ -322,7 +347,13 @@ class DenseSet {
   // ============ Pseudo Linked List in DenseSet end ==================
 
   // returns (prev, item) pair. If item is root, then prev is null.
-  std::pair<DensePtr*, DensePtr*> Find(const void* ptr, uint32_t bid, uint32_t cookie);
+  std::pair<DensePtr*, DensePtr*> Find(const void* ptr, uint32_t bid, uint32_t cookie) {
+    auto [_, p, c] = Find2(ptr, bid, cookie);
+    return {p, c};
+  }
+
+  // returns bid and (prev, item) pair. If item is root, then prev is null.
+  std::tuple<size_t, DensePtr*, DensePtr*> Find2(const void* ptr, uint32_t bid, uint32_t cookie);
 
   DenseLinkKey* NewLink(void* data, DensePtr next);
 

src/core/string_set.h

@@ -52,6 +52,9 @@ class StringSet : public DenseSet {
     using pointer = sds*;
     using reference = sds&;
 
+    explicit iterator(const IteratorBase& o) : IteratorBase(o) {
+    }
+
     iterator() : IteratorBase() {
     }
 
@@ -78,6 +81,9 @@ class StringSet : public DenseSet {
     value_type operator->() {
       return (value_type)curr_entry_->GetObject();
     }
+
+    using IteratorBase::ExpiryTime;
+    using IteratorBase::HasExpiry;
   };
 
   class const_iterator : private IteratorBase {
@@ -122,16 +128,19 @@ class StringSet : public DenseSet {
   iterator end() {
     return iterator{this, true};
   }
+  /*
+    const_iterator cbegin() const {
+      return const_iterator{this, false};
+    }
 
-  const_iterator cbegin() const {
+    const_iterator cend() const {
-    return const_iterator{this, false};
+      return const_iterator{this, true};
-  }
+    }
-
+  */
-  const_iterator cend() const {
-    return const_iterator{this, true};
-  }
-
   uint32_t Scan(uint32_t, const std::function<void(sds)>&) const;
+  iterator Find(std::string_view member) {
+    return iterator{FindIt(&member, 1)};
+  }
 
  protected:
   uint64_t Hash(const void* ptr, uint32_t cookie) const override;

src/core/string_set_test.cc

@@ -389,6 +389,9 @@ TEST_F(StringSetTest, Iteration) {
 TEST_F(StringSetTest, Ttl) {
   EXPECT_TRUE(ss_->Add("bla"sv, 1));
   EXPECT_FALSE(ss_->Add("bla"sv, 1));
+  auto it = ss_->Find("bla"sv);
+  EXPECT_EQ(1u, it.ExpiryTime());
+
   ss_->set_time(1);
   EXPECT_TRUE(ss_->Add("bla"sv, 1));
   EXPECT_EQ(1u, ss_->Size());
@@ -397,11 +400,16 @@ TEST_F(StringSetTest, Ttl) {
     EXPECT_TRUE(ss_->Add(StrCat("foo", i), 1));
   }
   EXPECT_EQ(101u, ss_->Size());
+  it = ss_->Find("foo50");
+  EXPECT_STREQ("foo50", *it);
+  EXPECT_EQ(2u, it.ExpiryTime());
 
   ss_->set_time(2);
   for (unsigned i = 0; i < 100; ++i) {
     EXPECT_TRUE(ss_->Add(StrCat("bar", i)));
   }
+  it = ss_->Find("bar50");
+  EXPECT_FALSE(it.HasExpiry());
 
   for (auto it = ss_->begin(); it != ss_->end(); ++it) {
     ASSERT_TRUE(absl::StartsWith(*it, "bar")) << *it;