feat (server): add on move hook to dash insert new and bumpup (#5278)

* server: add on move hook to dash insert new and bumpup

---------

Signed-off-by: adi_holden <[email protected]>
Signed-off-by: Roman Gershman <[email protected]>
Co-authored-by: Roman Gershman <[email protected]>

Author: adiholden

src/core/dash.h

@@ -85,6 +85,9 @@ class DashTable : public detail::DashTableBase {
       return true;
     }
 
+    void OnMove(Cursor source, Cursor dest) {
+    }
+
     void RecordSplit(SegmentType* segment) {
     }
     /*
@@ -286,9 +289,13 @@ class DashTable : public detail::DashTableBase {
   // Returns true if an element was deleted i.e the rightmost slot was busy.
   bool ShiftRight(bucket_iterator it);
 
-  template <typename BumpPolicy> iterator BumpUp(iterator it, const BumpPolicy& bp) {
-    SegmentIterator seg_it =
-        segment_[it.seg_id_]->BumpUp(it.bucket_id_, it.slot_id_, DoHash(it->first), bp);
+  template <typename BumpPolicy> iterator BumpUp(iterator it, BumpPolicy& bp) {
+    SegmentIterator seg_it = segment_[it.seg_id_]->BumpUp(
+        it.bucket_id_, it.slot_id_, DoHash(it->first), bp,
+        [&](uint32_t segment_id, detail::PhysicalBid from, detail::PhysicalBid to) {
+          // OnMove is used to notify policy about the items moves across buckets.
+          bp.OnMove(Cursor{global_depth_, segment_id, from}, Cursor{global_depth_, segment_id, to});
+        });
 
     return iterator{this, it.seg_id_, seg_it.index, seg_it.slot};
   }
@@ -314,7 +321,7 @@ class DashTable : public detail::DashTableBase {
                                            InsertMode mode);
 
   void IncreaseDepth(unsigned new_depth);
-  void Split(uint32_t seg_id);
+  template <typename EvictionPolicy> void Split(uint32_t seg_id, EvictionPolicy& ev);
 
   // Segment directory contains multiple segment pointers, some of them pointing to
   // the same object. IterateDistinct goes over all distinct segments in the table.
@@ -786,12 +793,19 @@ auto DashTable<_Key, _Value, Policy>::InsertInternal(U&& key, V&& value, Evictio
     typename SegmentType::Iterator it;
     bool res = true;
     unsigned num_buckets = target->num_buckets();
+
+    auto move_cb = [&](uint32_t segment_id, detail::PhysicalBid from, detail::PhysicalBid to) {
+      // OnMove is used to notify policy about the move of items across buckets.
+      ev.OnMove(Cursor{global_depth_, segment_id, from}, Cursor{global_depth_, segment_id, to});
+    };
+
     if (mode == InsertMode::kForceInsert) {
-      it = target->InsertUniq(std::forward<U>(key), std::forward<V>(value), key_hash, true);
+      it =
+          target->InsertUniq(std::forward<U>(key), std::forward<V>(value), key_hash, true, move_cb);
       res = it.found();
     } else {
-      std::tie(it, res) =
-          target->Insert(std::forward<U>(key), std::forward<V>(value), key_hash, EqPred(key));
+      std::tie(it, res) = target->Insert(std::forward<U>(key), std::forward<V>(value), key_hash,
+                                         EqPred(key), move_cb);
     }
 
     if (res) {  // success
@@ -846,7 +860,7 @@ auto DashTable<_Key, _Value, Policy>::InsertInternal(U&& key, V&& value, Evictio
       }
 
       auto hash_fn = [this](const auto& k) { return policy_.HashFn(k); };
-      unsigned moved = target->UnloadStash(hash_fn);
+      unsigned moved = target->UnloadStash(hash_fn, move_cb);
       if (moved > 0) {
         stash_unloaded_ += moved;
         continue;
@@ -876,7 +890,7 @@ auto DashTable<_Key, _Value, Policy>::InsertInternal(U&& key, V&& value, Evictio
     }
 
     ev.RecordSplit(target);
-    Split(target_seg_id);
+    Split(target_seg_id, ev);
   }
 
   return std::make_pair(iterator{}, false);
@@ -893,12 +907,14 @@ void DashTable<_Key, _Value, Policy>::IncreaseDepth(unsigned new_depth) {
   for (int i = prev_sz - 1; i >= 0; --i) {
     size_t offs = i * repl_cnt;
     std::fill(segment_.begin() + offs, segment_.begin() + offs + repl_cnt, segment_[i]);
+    segment_[i]->set_segment_id(offs);  // update segment id.
   }
   global_depth_ = new_depth;
 }
 
 template <typename _Key, typename _Value, typename Policy>
-void DashTable<_Key, _Value, Policy>::Split(uint32_t seg_id) {
+template <typename EvictionPolicy>
+void DashTable<_Key, _Value, Policy>::Split(uint32_t seg_id, EvictionPolicy& ev) {
   SegmentType* source = segment_[seg_id];
 
   uint32_t chunk_size = 1u << (global_depth_ - source->local_depth());
@@ -911,7 +927,15 @@ void DashTable<_Key, _Value, Policy>::Split(uint32_t seg_id) {
 
   // remove current segment bucket count.
   bucket_count_ -= (source->num_buckets() + target->num_buckets());
-  source->Split(std::move(hash_fn), target);  // increases the depth.
+
+  source->Split(
+      std::move(hash_fn), target,
+      [&](uint32_t segment_from, detail::PhysicalBid from, uint32_t segment_to,
+          detail::PhysicalBid to) {
+        // OnMove is used to notify eviction policy about the moves across buckets/segments
+        // during the split.
+        ev.OnMove(Cursor{global_depth_, segment_from, from}, Cursor{global_depth_, segment_to, to});
+      });
 
   // add back the updated bucket count.
   bucket_count_ += (target->num_buckets() + source->num_buckets());

src/core/dash_internal.h

@@ -413,10 +413,12 @@ class Segment {
 
   // Returns (iterator, true) if insert succeeds,
   // (iterator, false) for duplicate and (invalid-iterator, false) if it's full
-  template <typename K, typename V, typename Pred>
-  std::pair<Iterator, bool> Insert(K&& key, V&& value, Hash_t key_hash, Pred&& pred);
+  template <typename K, typename V, typename Pred, typename OnMoveCb>
+  std::pair<Iterator, bool> Insert(K&& key, V&& value, Hash_t key_hash, Pred&& pred,
+                                   OnMoveCb&& on_move_cb);
 
-  template <typename HashFn> void Split(HashFn&& hfunc, Segment* dest);
+  template <typename HashFn, typename OnMoveCb>
+  void Split(HashFn&& hfunc, Segment* dest, OnMoveCb&& on_move_cb);
 
   void Delete(const Iterator& it, Hash_t key_hash);
 
@@ -523,8 +525,8 @@ class Segment {
   // otherwise chooses home bucket first.
   // TODO: I am actually not sure if spread optimization is helpful. Worth checking
   // whether we get higher occupancy rates when using it.
-  template <typename U, typename V>
-  Iterator InsertUniq(U&& key, V&& value, Hash_t key_hash, bool spread);
+  template <typename U, typename V, typename OnMoveCb>
+  Iterator InsertUniq(U&& key, V&& value, Hash_t key_hash, bool spread, OnMoveCb&& on_move_cb);
 
   // capture version change in case of insert.
   // Returns ids of buckets whose version would cross ver_threshold upon insertion of key_hash
@@ -558,19 +560,25 @@ class Segment {
   }
 
   // Bumps up this entry making it more "important" for the eviction policy.
-  template <typename BumpPolicy>
-  Iterator BumpUp(PhysicalBid bid, SlotId slot, Hash_t key_hash, const BumpPolicy& ev);
+  template <typename BumpPolicy, typename OnMoveCb>
+  Iterator BumpUp(PhysicalBid bid, SlotId slot, Hash_t key_hash, const BumpPolicy& ev,
+                  OnMoveCb&& cb);
 
   // Tries to move stash entries back to their normal buckets (exact or neighbour).
   // Returns number of entries that succeeded to unload.
   // Important! Affects versions of the moved items and the items in the destination
   // buckets.
-  template <typename HFunc> unsigned UnloadStash(HFunc&& hfunc);
+  template <typename HFunc, typename OnMoveCb> unsigned UnloadStash(HFunc&& hfunc, OnMoveCb&& cb);
 
   unsigned num_buckets() const {
     return kBucketNum + kStashBucketNum;
   }
 
+  // needed only when DashTable grows its segment table.
+  void set_segment_id(uint32_t new_id) {
+    segment_id_ = new_id;
+  }
+
  private:
   static_assert(sizeof(Iterator) == 2);
 
@@ -1084,15 +1092,16 @@ auto Segment<Key, Value, Policy>::TryMoveFromStash(unsigned stash_id, unsigned s
 }
 
 template <typename Key, typename Value, typename Policy>
-template <typename U, typename V, typename Pred>
-auto Segment<Key, Value, Policy>::Insert(U&& key, V&& value, Hash_t key_hash, Pred&& pred)
-    -> std::pair<Iterator, bool> {
+template <typename U, typename V, typename Pred, typename OnMoveCb>
+auto Segment<Key, Value, Policy>::Insert(U&& key, V&& value, Hash_t key_hash, Pred&& pred,
+                                         OnMoveCb&& on_move_cb) -> std::pair<Iterator, bool> {
   Iterator it = FindIt(key_hash, pred);
   if (it.found()) {
     return std::make_pair(it, false); /* duplicate insert*/
   }
 
-  it = InsertUniq(std::forward<U>(key), std::forward<V>(value), key_hash, true);
+  it = InsertUniq(std::forward<U>(key), std::forward<V>(value), key_hash, true,
+                  std::forward<OnMoveCb>(on_move_cb));
 
   return std::make_pair(it, it.found());
 }
@@ -1210,8 +1219,8 @@ void Segment<Key, Value, Policy>::Delete(const Iterator& it, Hash_t key_hash) {
 // Split items from the left segment to the right during the growth phase.
 // right segment will have all the items with lsb at local_depth ==1 .
 template <typename Key, typename Value, typename Policy>
-template <typename HFunc>
-void Segment<Key, Value, Policy>::Split(HFunc&& hfn, Segment* dest_right) {
+template <typename HFunc, typename MoveCb>
+void Segment<Key, Value, Policy>::Split(HFunc&& hfn, Segment* dest_right, MoveCb&& on_move_cb) {
   ++local_depth_;
   dest_right->local_depth_ = local_depth_;
 
@@ -1243,8 +1252,11 @@ void Segment<Key, Value, Policy>::Split(HFunc&& hfn, Segment* dest_right) {
 
       invalid_mask |= (1u << slot);
 
+      // We pass dummy callback because we are not interested to track movements in the newly
+      // created segment.
       Iterator it = dest_right->InsertUniq(std::forward<Key_t>(bucket->key[slot]),
-                                           std::forward<Value_t>(bucket->value[slot]), hash, false);
+                                           std::forward<Value_t>(bucket->value[slot]), hash, false,
+                                           [](auto&&...) {});
 
       // we move items residing in a regular bucket to a new segment.
       // Note 1: in case we are somehow attacked with items that after the split
@@ -1266,6 +1278,7 @@ void Segment<Key, Value, Policy>::Split(HFunc&& hfn, Segment* dest_right) {
       // selective bias will be able to hit our dashtable with items with the same bucket id.
       assert(it.found());
       update_version(*bucket, it.index);
+      on_move_cb(segment_id_, i, dest_right->segment_id_, it.index);
     };
 
     bucket_[i].ForEachSlot(std::move(cb));
@@ -1286,17 +1299,20 @@ void Segment<Key, Value, Policy>::Split(HFunc&& hfn, Segment* dest_right) {
         Iterator it = TryMoveFromStash(i, slot, hash);
         if (it.found()) {
           invalid_mask |= (1u << slot);
+          on_move_cb(segment_id_, i, segment_id_, it.index);
         }
 
         return;
       }
 
       invalid_mask |= (1u << slot);
       auto it = dest_right->InsertUniq(std::forward<Key_t>(bucket->key[slot]),
-                                       std::forward<Value_t>(bucket->value[slot]), hash, false);
+                                       std::forward<Value_t>(bucket->value[slot]), hash, false,
+                                       /* not interested in these movements */ [](auto&&...) {});
       (void)it;
       assert(it.index != kNanBid);
       update_version(*bucket, it.index);
+      on_move_cb(segment_id_, i, dest_right->segment_id_, it.index);
 
       // Remove stash reference pointing to stash bucket i.
       RemoveStashReference(i, hash);
@@ -1348,9 +1364,9 @@ bool Segment<Key, Value, Policy>::CheckIfMovesToOther(bool own_items, unsigned f
 }
 
 template <typename Key, typename Value, typename Policy>
-template <typename U, typename V>
-auto Segment<Key, Value, Policy>::InsertUniq(U&& key, V&& value, Hash_t key_hash, bool spread)
-    -> Iterator {
+template <typename U, typename V, typename OnMoveCb>
+auto Segment<Key, Value, Policy>::InsertUniq(U&& key, V&& value, Hash_t key_hash, bool spread,
+                                             OnMoveCb&& on_move_cb) -> Iterator {
   const uint8_t bid = HomeIndex(key_hash);
   const uint8_t nid = NextBid(bid);
 
@@ -1385,13 +1401,15 @@ auto Segment<Key, Value, Policy>::InsertUniq(U&& key, V&& value, Hash_t key_hash
   int displace_index = MoveToOther(true, nid, NextBid(nid));
   if (displace_index >= 0) {
     neighbor.Insert(displace_index, std::forward<U>(key), std::forward<V>(value), meta_hash, true);
+    on_move_cb(segment_id_, nid, NextBid(nid));
     return Iterator{nid, uint8_t(displace_index)};
   }
 
   unsigned prev_idx = PrevBid(bid);
   displace_index = MoveToOther(false, bid, prev_idx);
   if (displace_index >= 0) {
     target.Insert(displace_index, std::forward<U>(key), std::forward<V>(value), meta_hash, false);
+    on_move_cb(segment_id_, bid, prev_idx);
     return Iterator{bid, uint8_t(displace_index)};
   }
 
@@ -1597,9 +1615,9 @@ auto Segment<Key, Value, Policy>::FindValidStartingFrom(PhysicalBid bid, unsigne
 }
 
 template <typename Key, typename Value, typename Policy>
-template <typename BumpPolicy>
+template <typename BumpPolicy, typename OnMoveCb>
 auto Segment<Key, Value, Policy>::BumpUp(uint8_t bid, SlotId slot, Hash_t key_hash,
-                                         const BumpPolicy& bp) -> Iterator {
+                                         const BumpPolicy& bp, OnMoveCb&& on_move_cb) -> Iterator {
   auto& from = GetBucket(bid);
 
   if (!bp.CanBump(from.key[slot])) {
@@ -1624,6 +1642,7 @@ auto Segment<Key, Value, Policy>::BumpUp(uint8_t bid, SlotId slot, Hash_t key_ha
   if (Iterator it = TryMoveFromStash(stash_pos, slot, key_hash); it.found()) {
     // TryMoveFromStash handles versions internally.
     from.Delete(slot);
+    on_move_cb(segment_id_, bid, it.index);
     return it;
   }
 
@@ -1690,12 +1709,14 @@ auto Segment<Key, Value, Policy>::BumpUp(uint8_t bid, SlotId slot, Hash_t key_ha
     swapb.SetStashPtr(stash_pos, swap_fp, bucket_ + next_bid);
   }
 
+  on_move_cb(segment_id_, bid, swap_bid);
+  on_move_cb(segment_id_, swap_bid, bid);
   return Iterator{swap_bid, kLastSlot};
 }
 
 template <typename Key, typename Value, typename Policy>
-template <typename HFunc>
-unsigned Segment<Key, Value, Policy>::UnloadStash(HFunc&& hfunc) {
+template <typename HFunc, typename OnMoveCb>
+unsigned Segment<Key, Value, Policy>::UnloadStash(HFunc&& hfunc, OnMoveCb&& on_move_cb) {
   unsigned moved = 0;
 
   for (unsigned i = 0; i < kStashBucketNum; ++i) {
@@ -1710,6 +1731,7 @@ unsigned Segment<Key, Value, Policy>::UnloadStash(HFunc&& hfunc) {
       if (res.found()) {
         ++moved;
         invalid_mask |= (1u << slot);
+        on_move_cb(segment_id_, i, res.index);
       }
     };