From c74bc6690607df9804ff71407103f9f9a9008888 Mon Sep 17 00:00:00 2001
From: Borys <borys@dragonflydb.io>
Date: Mon, 14 Oct 2024 20:42:10 +0300
Subject: [PATCH] chore: optimize DenseSet link allocation

---
 src/core/dense_set.cc | 18 ++++++++++++++----
 src/core/dense_set.h  | 14 +++++++++++---
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/core/dense_set.cc b/src/core/dense_set.cc
index 67bf4d6c255f..02b90fe4a9a6 100644
--- a/src/core/dense_set.cc
+++ b/src/core/dense_set.cc
@@ -90,6 +90,11 @@ DenseSet::DenseSet(MemoryResource* mr) : entries_(mr) {
 }
 
 DenseSet::~DenseSet() {
+  while (cached_link_num) {
+    mr()->deallocate(cached_links[--cached_link_num], sizeof(DenseLinkKey), alignof(DenseLinkKey));
+    --num_links_;
+  }
+
   // We can not call Clear from the base class because it internally calls ObjDelete which is
   // a virtual function. Therefore, destructor of the derived classes must clean up the table.
   CHECK(entries_.empty());
@@ -798,13 +803,18 @@ uint32_t DenseSet::Scan(uint32_t cursor, const ItemCb& cb) const {
 }
 
 auto DenseSet::NewLink(void* data, DensePtr next) -> DenseLinkKey* {
-  LinkAllocator la(mr());
-  DenseLinkKey* lk = la.allocate(1);
-  la.construct(lk);
+  DenseLinkKey* lk = nullptr;
+  if (cached_link_num) {
+    lk = cached_links[--cached_link_num];
+  } else {
+    LinkAllocator la(mr());
+    lk = la.allocate(1);
+    la.construct(lk);
+    ++num_links_;
+  }
 
   lk->next = next;
   lk->SetObject(data);
-  ++num_links_;
 
   return lk;
 }
diff --git a/src/core/dense_set.h b/src/core/dense_set.h
index 503044f943c7..176bfa7bb73b 100644
--- a/src/core/dense_set.h
+++ b/src/core/dense_set.h
@@ -390,9 +390,13 @@ class DenseSet {
   DenseLinkKey* NewLink(void* data, DensePtr next);
 
   inline void FreeLink(DenseLinkKey* plink) {
-    // deallocate the link if it is no longer a link as it is now in an empty list
-    mr()->deallocate(plink, sizeof(DenseLinkKey), alignof(DenseLinkKey));
-    --num_links_;
+    if (cached_link_num < max_cached_links) {
+      cached_links[cached_link_num++] = plink;
+    } else {
+      // deallocate the link if it is no longer a link as it is now in an empty list
+      mr()->deallocate(plink, sizeof(DenseLinkKey), alignof(DenseLinkKey));
+      --num_links_;
+    }
   }
 
   // Returns true if *node was deleted.
@@ -420,6 +424,10 @@ class DenseSet {
   uint32_t time_now_ = 0;
 
   mutable bool expiration_used_ = false;
+
+  static constexpr uint32_t max_cached_links = 8;
+  DenseLinkKey* cached_links[max_cached_links];
+  uint32_t cached_link_num = 0;
 };
 
 inline void* DenseSet::FindInternal(const void* obj, uint64_t hashcode, uint32_t cookie) const {