core: Add count min sketch to track sizes

Signed-off-by: Abhijat Malviya <abhijat@dragonflydb.io>

Author: abhijat

src/core/CMakeLists.txt

@@ -27,7 +27,7 @@ add_library(dfly_core allocation_tracker.cc bloom.cc compact_object.cc dense_set
     interpreter.cc glob_matcher.cc mi_memory_resource.cc qlist.cc sds_utils.cc
     segment_allocator.cc score_map.cc small_string.cc sorted_map.cc task_queue.cc
     tx_queue.cc string_set.cc string_map.cc top_keys.cc detail/bitpacking.cc
-        page_usage_stats.cc)
+        page_usage_stats.cc count_min_sketch.cc)
 
 cxx_link(dfly_core base dfly_search_core fibers2 jsonpath
     absl::flat_hash_map absl::str_format absl::random_random redis_lib
@@ -56,6 +56,7 @@ cxx_test(qlist_test dfly_core DATA testdata/list.txt.zst LABELS DFLY)
 cxx_test(zstd_test dfly_core TRDP::zstd LABELS DFLY)
 cxx_test(top_keys_test dfly_core LABELS DFLY)
 cxx_test(page_usage_stats_test dfly_core LABELS DFLY)
+cxx_test(count_min_sketch_test dfly_core LABELS DFLY)
 
 if(LIB_PCRE2)
   target_compile_definitions(dfly_core_test PRIVATE USE_PCRE2=1)

src/core/count_min_sketch.cc

@@ -0,0 +1,129 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "core/count_min_sketch.h"
+
+#include <absl/time/clock.h>
+#include <xxhash.h>
+
+#include <cmath>
+#include <iostream>
+
+namespace {
+
+constexpr auto MAX = std::numeric_limits<dfly::CountMinSketch::SizeT>::max();
+
+uint64_t GetCurrentMS() {
+  return absl::GetCurrentTimeNanos() / 1000 / 1000;
+}
+
+// Value halves every 5000 ms: ln(2) / 5000
+constexpr double EXP_DECAY_CONST = 0.000138629;
+
+uint64_t ExponentialDecay(uint64_t value, int64_t time_delta) {
+  return value * std::exp(-time_delta * EXP_DECAY_CONST);
+}
+
+// Value decrements by one every 1000 ms
+constexpr double LIN_DECAY_CONST = 0.001;
+
+uint64_t LinearDecay(uint64_t value, int64_t time_delta) {
+  const double decay = time_delta * LIN_DECAY_CONST;
+  return value - std::min(static_cast<double>(value), decay);
+}
+
+uint64_t ApplyDecay(uint64_t value, int64_t time_delta, dfly::MultiSketch::Decay decay) {
+  switch (decay) {
+    case dfly::MultiSketch::Decay::Exponential:
+      return ExponentialDecay(value, time_delta);
+    case dfly::MultiSketch::Decay::Linear:
+      return LinearDecay(value, time_delta);
+    case dfly::MultiSketch::Decay::SlidingWindow:
+      return value;
+  }
+  ABSL_UNREACHABLE();
+}
+
+}  // namespace
+
+namespace dfly {
+
+CountMinSketch::CountMinSketch(double epsilon, double delta) {
+  width_ = std::exp(1) / epsilon;
+  depth_ = std::log(1.0 / delta);
+  counters_.reserve(depth_);
+  for (uint64_t i = 0; i < depth_; ++i) {
+    counters_.emplace_back(width_, 0);
+  }
+}
+
+void CountMinSketch::Update(uint64_t key, CountMinSketch::SizeT incr) {
+  uint64_t i = 0;
+  std::for_each(counters_.begin(), counters_.end(), [&](auto& ctr) {
+    const uint64_t index = Hash(key, i++);
+    const SizeT curr = ctr[index];
+    const SizeT updated = curr + incr;
+    ctr[index] = updated < curr ? MAX : updated;
+  });
+}
+
+CountMinSketch::SizeT CountMinSketch::EstimateFrequency(uint64_t key) const {
+  uint64_t i = 0;
+  auto it = counters_.begin();
+  SizeT estimate = (*it)[Hash(key, i++)];
+  std::for_each(++it, counters_.end(),
+                [&](const auto& ctr) { estimate = std::min(estimate, ctr[Hash(key, i++)]); });
+  return estimate;
+}
+
+void CountMinSketch::Reset() {
+  for (auto& ctr : counters_) {
+    std::fill(ctr.begin(), ctr.end(), 0);
+  }
+}
+
+uint64_t CountMinSketch::Hash(uint64_t key, uint64_t i) const {
+  return XXH3_64bits_withSeed(&key, sizeof(key), i) % width_;
+}
+
+MultiSketch::MultiSketch(uint64_t rollover_ms, double epsilon, double delta, Decay decay)
+    : rollover_ms_(rollover_ms), current_sketch_(sketches_.size() - 1), decay_t_(decay) {
+  uint64_t now = GetCurrentMS();
+  for (uint64_t i = 0; i < sketches_.size(); ++i) {
+    sketches_[i] = SketchWithTimestamp{CountMinSketch{epsilon, delta}, now, now};
+  }
+}
+
+void MultiSketch::Update(uint64_t key, CountMinSketch::SizeT incr) {
+  if (++rollover_check_ >= rollover_check_every_) {
+    MaybeRolloverCurrentSketch();
+    rollover_check_ = 0;
+  }
+  sketches_[current_sketch_].sketch_.Update(key, incr);
+}
+
+CountMinSketch::SizeT MultiSketch::EstimateFrequency(uint64_t key) const {
+  CountMinSketch::SizeT estimate = 0;
+  const uint64_t now = GetCurrentMS();
+
+  for (const auto& sketch : sketches_) {
+    const auto e = sketch.sketch_.EstimateFrequency(key);
+    // TODO use average time of sketch to compute delta
+    estimate += ApplyDecay(e, now - sketch.start_time_, decay_t_);
+  }
+  return estimate;
+}
+
+void MultiSketch::MaybeRolloverCurrentSketch() {
+  const uint64_t now = GetCurrentMS();
+  const uint64_t oldest = (current_sketch_ + 1) % sketches_.size();
+  if (const uint64_t oldest_ts = sketches_[oldest].start_time_; now - oldest_ts > rollover_ms_) {
+    sketches_[oldest].sketch_.Reset();
+    sketches_[oldest].start_time_ = now;
+    sketches_[current_sketch_].end_time_ = now;
+    current_sketch_ = oldest;
+  }
+}
+
+}  // namespace dfly

src/core/count_min_sketch.h

@@ -0,0 +1,95 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <vector>
+
+namespace dfly {
+
+class CountMinSketch {
+ public:
+  using SizeT = uint16_t;
+
+  // epsilon is the maximum deviation from actual frequency allowed per element times the sum of all
+  // frequencies:
+  //    f_actual <= f_estimated <= f_actual + epsilon * N
+  //      where N is the sum of all frequencies
+  // delta is the probability that f_estimated overshoots the epsilon threshold for a single
+  // estimate, aka failure probability.
+
+  // With default values, the dimension of the counter table is 27182 x 9, and the size is
+  // around 490 KiBs.
+  explicit CountMinSketch(double epsilon = 0.0001, double delta = 0.0001);
+
+  void Update(uint64_t key, SizeT incr = 1);
+
+  SizeT EstimateFrequency(uint64_t key) const;
+
+  void Reset();
+
+  CountMinSketch(const CountMinSketch& other) = delete;
+  CountMinSketch& operator=(const CountMinSketch& other) = delete;
+
+  CountMinSketch(CountMinSketch&& other) noexcept = default;
+  CountMinSketch& operator=(CountMinSketch&& other) noexcept = default;
+
+ private:
+  uint64_t Hash(uint64_t key, uint64_t i) const;
+
+  std::vector<std::vector<SizeT>> counters_;
+  uint64_t width_;
+  uint64_t depth_;
+};
+
+// Maintains a list of three sketches with timestamps. Updates are made to the current sketch.
+// Once the oldest sketch is older than a fixed limit, it is discarded and becomes the current
+// sketch. Estimates are the sum across all sketches.
+class MultiSketch {
+  struct SketchWithTimestamp {
+    CountMinSketch sketch_;
+    uint64_t start_time_{0};
+    uint64_t end_time_{0};
+  };
+
+ public:
+  enum class Decay : uint8_t {
+    Exponential,
+    Linear,
+    SlidingWindow,
+  };
+
+  explicit MultiSketch(uint64_t rollover_ms = 1000, double epsilon = 0.0001, double delta = 0.0001,
+                       Decay decay = Decay::Exponential);
+
+  MultiSketch(const MultiSketch& other) = delete;
+  MultiSketch& operator=(const MultiSketch& other) = delete;
+
+  MultiSketch(MultiSketch&& other) noexcept = default;
+  MultiSketch& operator=(MultiSketch&& other) noexcept = default;
+
+  void Update(uint64_t key, CountMinSketch::SizeT incr = 1);
+
+  CountMinSketch::SizeT EstimateFrequency(uint64_t key) const;
+
+  // For unit tests, allow setting a smaller limit
+  void SetRolloverCheckLimit(uint64_t rollover_check_limit) {
+    rollover_check_every_ = rollover_check_limit;
+  }
+
+ private:
+  void MaybeRolloverCurrentSketch();
+
+  std::array<SketchWithTimestamp, 3> sketches_;
+  uint64_t rollover_ms_;
+  uint64_t current_sketch_;
+
+  // Do a rollover check every N calls to avoid expensive GetTime calls
+  uint64_t rollover_check_every_{512};
+  uint64_t rollover_check_{0};
+  Decay decay_t_;
+};
+
+}  // namespace dfly