FixedSizeIndex: Handling non-existing key case

byahn0996 · facebook-github-bot · commit d0f47aad65f5 · 2025-09-11T10:21:38.000-07:00
Summary:
With FixedSizeIndex, there's more possibilities for false positive
For the non-existing item, it may try best effort to find it, and may return false positive PackedItemRecord instance.
But when we are sure that it doesn't really exist, it's better to return "it's not there".

To simply indicate it doesn't exist, I changed recordRef() to recordPtr() in ExclusiveLockedBucket and SharedLockedBucket. Those members were already pointer type within the structures, so there's no increased vulnerability by using pointer type directly.

Reviewed By: rlyerly

Differential Revision: D77402423

fbshipit-source-id: ed65de8f94f724cbc86835320c0aac6b6de47d26
diff --git a/cachelib/navy/block_cache/FixedSizeIndex.cpp b/cachelib/navy/block_cache/FixedSizeIndex.cpp
@@ -53,12 +53,12 @@ Index::LookupResult FixedSizeIndex::lookup(uint64_t key) {
   // also not sure if bumping up hit count with couldExist() makes sense.
   ExclusiveLockedBucket elb{key, *this, false};
 
-  if (elb.recordRef().isValid()) {
+  if (elb.isValidRecord()) {
     return LookupResult(true,
-                        ItemRecord(elb.recordRef().address,
-                                   elb.recordRef().getSizeHint(),
+                        ItemRecord(elb.recordPtr()->address,
+                                   elb.recordPtr()->getSizeHint(),
                                    0, /* totalHits */
-                                   elb.recordRef().bumpCurHits()));
+                                   elb.recordPtr()->bumpCurHits()));
   }
 
   return {};
@@ -67,12 +67,12 @@ Index::LookupResult FixedSizeIndex::lookup(uint64_t key) {
 Index::LookupResult FixedSizeIndex::peek(uint64_t key) const {
   SharedLockedBucket slb{key, *this};
 
-  if (slb.recordRef().isValid()) {
+  if (slb.isValidRecord()) {
     return LookupResult(true,
-                        ItemRecord(slb.recordRef().address,
-                                   slb.recordRef().getSizeHint(),
+                        ItemRecord(slb.recordPtr()->address,
+                                   slb.recordPtr()->getSizeHint(),
                                    0, /* totalHits */
-                                   slb.recordRef().info.curHits));
+                                   slb.recordPtr()->info.curHits));
   }
 
   return {};
@@ -84,18 +84,19 @@ Index::LookupResult FixedSizeIndex::insert(uint64_t key,
   LookupResult lr;
   ExclusiveLockedBucket elb{key, *this, true};
 
-  if (elb.recordRef().isValid()) {
+  XDCHECK(elb.bucketExist());
+  if (elb.isValidRecord()) {
     lr = LookupResult(true,
-                      ItemRecord(elb.recordRef().address,
-                                 elb.recordRef().getSizeHint(),
+                      ItemRecord(elb.recordPtr()->address,
+                                 elb.recordPtr()->getSizeHint(),
                                  0, /* totalHits */
-                                 elb.recordRef().info.curHits));
+                                 elb.recordPtr()->info.curHits));
   } else {
     ++elb.validBucketCntRef();
   }
   // TODO: need to combine this two ops into one to make sure updateDistInfo()
   // part is not missed
-  elb.recordRef() = PackedItemRecord{address, sizeHint, /* currentHits */ 0};
+  *elb.recordPtr() = PackedItemRecord{address, sizeHint, /* currentHits */ 0};
   elb.updateDistInfo(key, *this);
 
   return lr;
@@ -106,9 +107,9 @@ bool FixedSizeIndex::replaceIfMatch(uint64_t key,
                                     uint32_t oldAddress) {
   ExclusiveLockedBucket elb{key, *this, false};
 
-  if (elb.recordRef().address == oldAddress) {
-    elb.recordRef().address = newAddress;
-    elb.recordRef().info.curHits = 0;
+  if (elb.isValidRecord() && elb.recordPtr()->address == oldAddress) {
+    elb.recordPtr()->address = newAddress;
+    elb.recordPtr()->info.curHits = 0;
     return true;
   }
   return false;
@@ -117,27 +118,29 @@ bool FixedSizeIndex::replaceIfMatch(uint64_t key,
 Index::LookupResult FixedSizeIndex::remove(uint64_t key) {
   ExclusiveLockedBucket elb{key, *this, false};
 
-  if (elb.recordRef().isValid()) {
-    LookupResult lr{true, ItemRecord(elb.recordRef().address,
-                                     elb.recordRef().getSizeHint(),
+  if (elb.isValidRecord()) {
+    LookupResult lr{true, ItemRecord(elb.recordPtr()->address,
+                                     elb.recordPtr()->getSizeHint(),
                                      0, /* totalHits */
-                                     elb.recordRef().info.curHits)};
+                                     elb.recordPtr()->info.curHits)};
 
     XDCHECK(elb.validBucketCntRef() > 0);
     --elb.validBucketCntRef();
-    elb.recordRef() = PackedItemRecord{};
+    *elb.recordPtr() = PackedItemRecord{};
     return lr;
   }
 
-  elb.recordRef() = PackedItemRecord{};
+  if (elb.bucketExist()) {
+    *elb.recordPtr() = PackedItemRecord{};
+  }
   return {};
 }
 
 bool FixedSizeIndex::removeIfMatch(uint64_t key, uint32_t address) {
   ExclusiveLockedBucket elb{key, *this, false};
 
-  if (elb.recordRef().address == address) {
-    elb.recordRef() = PackedItemRecord{};
+  if (elb.isValidRecord() && elb.recordPtr()->address == address) {
+    *elb.recordPtr() = PackedItemRecord{};
 
     XDCHECK(elb.validBucketCntRef() > 0);
     --elb.validBucketCntRef();
@@ -264,8 +267,8 @@ void FixedSizeIndex::setHitsTestOnly(uint64_t key,
                                      uint8_t totalHits) {
   ExclusiveLockedBucket elb{key, *this, false};
 
-  if (elb.recordRef().isValid()) {
-    elb.recordRef().info.curHits =
+  if (elb.isValidRecord()) {
+    elb.recordPtr()->info.curHits =
         PackedItemRecord::truncateCurHits(currentHits);
     XLOGF(INFO,
           "setHitsTestOnly() for {}. totalHits {} was discarded in "
diff --git a/cachelib/navy/block_cache/FixedSizeIndex.h b/cachelib/navy/block_cache/FixedSizeIndex.h
@@ -288,13 +288,18 @@ class FixedSizeIndex : public Index {
   void initialize();
 
   // Random prime numbers for the distance for next bucket slot to use.
-  static constexpr uint8_t kNextBucketOffset[4] = {0, 23, 61, 97};
+  static constexpr std::array<uint8_t, 4> kNextBucketOffset{0, 23, 61, 97};
+  // This offset will be used to indicate there's no valid bucket slot matching
+  // the given key
+  static constexpr uint8_t kInvalidBucketSlotOffset = 0xff;
+  // This bucket id will be used to indicate there's no valid bucket for the key
+  static constexpr uint64_t kInvalidBucketId = 0xffffffffffffffff;
 
   uint8_t decideBucketOffset(uint64_t bid, uint64_t key) const {
     auto mid = mutexId(bid);
 
     // Check if there's already one matching
-    for (auto i = 0; i < 4; i++) {
+    for (size_t i = 0; i < kNextBucketOffset.size(); i++) {
       auto curBid = calcBucketId(bid, i);
       // Make sure we don't go across the mutex boundary
       XDCHECK(mutexId(curBid) == mid) << bid << " " << i << " " << curBid;
@@ -307,7 +312,7 @@ class FixedSizeIndex : public Index {
     }
 
     // No match. Find the empty one
-    for (auto i = 0; i < 4; i++) {
+    for (size_t i = 0; i < kNextBucketOffset.size(); i++) {
       auto curBid = calcBucketId(bid, i);
       // Make sure we don't go across the mutex boundary
       XDCHECK(mutexId(curBid) == mid) << bid << " " << i << " " << curBid;
@@ -324,7 +329,7 @@ class FixedSizeIndex : public Index {
   uint8_t checkBucketOffset(uint64_t bid, uint64_t key) const {
     auto mid = mutexId(bid);
 
-    for (auto i = 1; i < 4; i++) {
+    for (size_t i = 0; i < kNextBucketOffset.size(); i++) {
       auto curBid = calcBucketId(bid, i);
       // Make sure we don't go across the mutex boundary
       XDCHECK(mutexId(curBid) == mid) << bid << " " << i << " " << curBid;
@@ -335,7 +340,20 @@ class FixedSizeIndex : public Index {
         return i;
       }
     }
-    return 0;
+    return kInvalidBucketSlotOffset;
+  }
+
+  // This helper will get the proper bucket id and record entry
+  // Return value : The pair of <Bucket id, pointer to the record>
+  std::pair<uint64_t, PackedItemRecord*> getBucket(uint64_t orgBid,
+                                                   uint8_t offset) const {
+    if (offset != kInvalidBucketSlotOffset) {
+      auto bid = calcBucketId(orgBid, offset);
+      return std::make_pair(bid, &ht_[bid]);
+    } else {
+      // There's no bucket for the given key
+      return std::make_pair(kInvalidBucketId, nullptr);
+    }
   }
 
   // Updates hits information of a key.
@@ -366,6 +384,7 @@ class FixedSizeIndex : public Index {
     // We don't want to go across the mutex boundary, so if it goes beyond that,
     // it will wrap around and go back to the beginning of current mutex
     // boundary
+    XDCHECK(offset < kNextBucketOffset.size()) << offset;
     return (bid / numBucketsPerMutex_) * numBucketsPerMutex_ +
            ((bid + kNextBucketOffset[offset]) % numBucketsPerMutex_);
   }
@@ -388,8 +407,8 @@ class FixedSizeIndex : public Index {
 
   // A helper class for exclusive locked access to a bucket.
   // It will lock the proper mutex with the given key when it's created.
-  // recordRef() and validBucketCntRef() will return the record and
-  // valid bucket count with exclusively locked bucket reference. Locked
+  // recordPtr() and validBucketCntRef() will return the record and
+  // valid bucket count with exclusively locked bucket. Locked
   // mutex will be released when it's destroyed.
   class ExclusiveLockedBucket {
    public:
@@ -399,59 +418,60 @@ class FixedSizeIndex : public Index {
         : bid_(index.bucketId(key)),
           mid_{index.mutexId(bid_)},
           lg_{index.mutex_[mid_]},
-          record_{&index.ht_[bid_]},
           validBuckets_{index.validBucketsPerMutex_[mid_]} {
       auto offset = alloc ? index.decideBucketOffset(bid_, key)
                           : index.checkBucketOffset(bid_, key);
-      if (offset != 0) {
-        bid_ = index.calcBucketId(bid_, offset);
-        record_ = &index.ht_[bid_];
-        bucketOffset_ = offset;
-      }
+      std::tie(bid_, record_) = index.getBucket(bid_, offset);
+      bucketOffset_ = offset;
     }
 
-    PackedItemRecord& recordRef() { return *record_; }
+    PackedItemRecord* recordPtr() { return record_; }
     size_t& validBucketCntRef() { return validBuckets_; }
     void updateDistInfo(uint64_t key, FixedSizeIndex& index) {
-      index.bucketDistInfo_.updateBucketFillInfo(
-          bid_, bucketOffset_, index.partialKeyBits(key));
+      if (bucketOffset_ != kInvalidBucketSlotOffset) {
+        index.bucketDistInfo_.updateBucketFillInfo(
+            bid_, bucketOffset_, index.partialKeyBits(key));
+      }
+    }
+    bool isValidRecord() const {
+      return (record_ != nullptr && record_->isValid());
     }
+    bool bucketExist() const { return record_ != nullptr; }
 
    private:
     uint64_t bid_;
     uint64_t mid_;
     std::lock_guard<SharedMutex> lg_;
-    PackedItemRecord* record_;
     size_t& validBuckets_;
-    uint8_t bucketOffset_{0};
+    PackedItemRecord* record_{};
+    uint8_t bucketOffset_{kInvalidBucketSlotOffset};
   };
 
   // A helper class for shared locked access to a bucket.
   // It will lock the proper mutex with the given key when it's created.
-  // recordRef() will return the record with shared locked bucket reference.
+  // recordPtr() will return the record with shared locked bucket.
   // Locked mutex will be released when it's destroyed.
   class SharedLockedBucket {
    public:
     explicit SharedLockedBucket(uint64_t key, const FixedSizeIndex& index)
         : bid_(index.bucketId(key)),
           mid_{index.mutexId(bid_)},
-          lg_{index.mutex_[mid_]},
-          record_{&index.ht_[bid_]} {
+          lg_{index.mutex_[mid_]} {
       // check next bucket if it should be used
-      auto offset = (index.checkBucketOffset(bid_, key));
-      if (offset != 0) {
-        bid_ = index.calcBucketId(bid_, offset);
-        record_ = &index.ht_[bid_];
-      }
+      std::tie(bid_, record_) =
+          index.getBucket(bid_, index.checkBucketOffset(bid_, key));
     }
 
-    const PackedItemRecord& recordRef() const { return *record_; }
+    const PackedItemRecord* recordPtr() const { return record_; }
+    bool isValidRecord() const {
+      return (record_ != nullptr && record_->isValid());
+    }
 
    private:
     uint64_t bid_;
     uint64_t mid_;
     std::shared_lock<SharedMutex> lg_;
-    const PackedItemRecord* record_;
+    const PackedItemRecord* record_{};
   };
 
 // For unit tests private member access
