feat: Add eviction based on rss memory usage (#5791)

This PR introduced eviction based on RSS memory usage. It uses max_memory and eviction_memory_budget_threshold flags to evict items. The eviction starts when rss_mem_current  is above max_memory_limit * eviction_memory_budget_threshold.

Another important step was introducing the eviction_state_ in EngineShard. This state is used to track how much memory has already been evicted but not yet freed, meaning it hasn’t reduced rss_mem_current counter yet. Once we observe a decrease in rss_mem_current, we also reduce the corresponding amount of evicted bytes in eviction_state_.

This is necessary because rss_mem_current updates with a delay, which can cause the eviction process to remove more items than actually needed. 

Flag `enable_heartbeat_rss_eviction` was added which can be used to enable or disable RSS eviction.

Closes #4011

Signed-off-by: Stepan Bagritsevich <[email protected]>
Co-authored-by: mkaruza <[email protected]>

Author: BagritsevichStepan

src/server/engine_shard.cc

@@ -54,13 +54,13 @@ ABSL_FLAG(string, tiered_prefix, "",
 
 ABSL_FLAG(bool, enable_heartbeat_eviction, true,
           "Enable eviction during heartbeat when memory is under pressure.");
-
+ABSL_FLAG(bool, enable_heartbeat_rss_eviction, true,
+          "Enable eviction during heartbeat when rss memory is under pressure. Evicition based "
+          "on used_memory will still be enabled.");
 ABSL_FLAG(double, eviction_memory_budget_threshold, 0.1,
           "Eviction starts when the free memory (including RSS memory) drops below "
           "eviction_memory_budget_threshold * max_memory_limit.");
-
 ABSL_FLAG(bool, background_heartbeat, false, "Whether to run heartbeat as a background fiber");
-
 ABSL_DECLARE_FLAG(uint32_t, max_eviction_per_heartbeat);
 
 namespace dfly {
@@ -116,46 +116,6 @@ size_t CalculateHowManyBytesToEvictOnShard(size_t global_memory_limit, size_t gl
   return shard_budget < shard_memory_threshold ? (shard_memory_threshold - shard_budget) : 0;
 }
 
-/* Calculates the number of bytes to evict based on memory and rss memory usage. */
-size_t CalculateEvictionBytes() {
-  const size_t shards_count = shard_set->size();
-  const double eviction_memory_budget_threshold = GetFlag(FLAGS_eviction_memory_budget_threshold);
-
-  size_t limit = max_memory_limit.load(memory_order_relaxed);
-  const size_t shard_memory_budget_threshold =
-      size_t(limit * eviction_memory_budget_threshold) / shards_count;
-
-  const size_t global_used_memory = used_mem_current.load(memory_order_relaxed);
-
-  // Calculate how many bytes we need to evict on this shard
-  size_t goal_bytes =
-      CalculateHowManyBytesToEvictOnShard(limit, global_used_memory, shard_memory_budget_threshold);
-
-  // TODO: Eviction due to rss usage is not working well as it causes eviction
-  // of to many keys untill we finally see decrease in rss. We need to improve
-  // this logic before we enable it.
-  /*
-  const double rss_oom_deny_ratio = ServerState::tlocal()->rss_oom_deny_ratio;
-  // If rss_oom_deny_ratio is set, we should evict depending on rss memory too
-  if (rss_oom_deny_ratio > 0.0) {
-    const size_t max_rss_memory = size_t(rss_oom_deny_ratio * max_memory_limit);
-    // We start eviction when we have less than eviction_memory_budget_threshold * 100% of free rss
-    memory const size_t shard_rss_memory_budget_threshold =
-        size_t(max_rss_memory * eviction_memory_budget_threshold) / shards_count;
-
-    // Calculate how much rss memory is used by all shards
-    const size_t global_used_rss_memory = rss_mem_current.load(memory_order_relaxed);
-
-    // Try to evict more bytes if we are close to the rss memory limit
-    goal_bytes = std::max(
-        goal_bytes, CalculateHowManyBytesToEvictOnShard(max_rss_memory, global_used_rss_memory,
-                                                        shard_rss_memory_budget_threshold));
-  }
-  */
-
-  return goal_bytes;
-}
-
 }  // namespace
 
 __thread EngineShard* EngineShard::shard_ = nullptr;
@@ -265,10 +225,7 @@ bool EngineShard::DefragTaskState::CheckRequired() {
   if (res == 0) {
     // finished checking.
     last_check_time = time(nullptr);
-    if (finfo.committed != finfo.committed_golden) {
-      LOG_FIRST_N(ERROR, 100) << "committed memory computed incorrectly: " << finfo.committed
-                              << " vs " << finfo.committed_golden;
-    }
+
     const double waste_threshold = GetFlag(FLAGS_mem_defrag_waste_threshold);
     if (finfo.wasted > size_t(finfo.committed * waste_threshold)) {
       VLOG(1) << "memory fragmentation issue found: " << finfo.wasted << " " << finfo.committed;
@@ -425,6 +382,8 @@ void EngineShard::StartPeriodicHeartbeatFiber(util::ProactorBase* pb) {
     return;
   }
   auto heartbeat = [this]() { Heartbeat(); };
+
+  eviction_state_.rss_eviction_enabled_ = GetFlag(FLAGS_enable_heartbeat_rss_eviction);
   std::chrono::milliseconds period_ms(*cycle_ms);
 
   fb2::Fiber::Opts fb_opts{.priority = absl::GetFlag(FLAGS_background_heartbeat)
@@ -702,6 +661,7 @@ void EngineShard::RetireExpiredAndEvict() {
   DbContext db_cntx;
   db_cntx.time_now_ms = GetCurrentTimeMs();
 
+  size_t deleted_bytes = 0;
   size_t eviction_goal = GetFlag(FLAGS_enable_heartbeat_eviction) ? CalculateEvictionBytes() : 0;
 
   for (unsigned i = 0; i < db_slice.db_array_size(); ++i) {
@@ -713,6 +673,7 @@ void EngineShard::RetireExpiredAndEvict() {
     if (!expt->Empty()) {
       DbSlice::DeleteExpiredStats stats = db_slice.DeleteExpiredStep(db_cntx, ttl_delete_target);
 
+      deleted_bytes += stats.deleted_bytes;
       eviction_goal -= std::min(eviction_goal, size_t(stats.deleted_bytes));
       counter_[TTL_TRAVERSE].IncBy(stats.traversed);
       counter_[TTL_DELETE].IncBy(stats.deleted);
@@ -734,9 +695,71 @@ void EngineShard::RetireExpiredAndEvict() {
                << " bytes. Max eviction per heartbeat: "
                << GetFlag(FLAGS_max_eviction_per_heartbeat);
 
+      deleted_bytes += evicted_bytes;
       eviction_goal -= std::min(eviction_goal, evicted_bytes);
     }
   }
+
+  eviction_state_.deleted_bytes_before_rss_update += deleted_bytes;
+}
+
+size_t EngineShard::CalculateEvictionBytes() {
+  const size_t shards_count = shard_set->size();
+  const double eviction_memory_budget_threshold = GetFlag(FLAGS_eviction_memory_budget_threshold);
+
+  // Calculate threshold for both used_memory and rss_memory
+  size_t limit = max_memory_limit.load(memory_order_relaxed);
+  const size_t shard_memory_budget_threshold =
+      size_t(limit * eviction_memory_budget_threshold) / shards_count;
+
+  const size_t global_used_memory = used_mem_current.load(memory_order_relaxed);
+
+  // Calculate how many bytes we need to evict on this shard
+  size_t goal_bytes =
+      CalculateHowManyBytesToEvictOnShard(limit, global_used_memory, shard_memory_budget_threshold);
+
+  VLOG_IF_EVERY_N(1, goal_bytes > 0, 50)
+      << "Used memory goal bytes: " << goal_bytes << ", used memory: " << global_used_memory
+      << ", memory limit: " << max_memory_limit;
+
+  // Check for `enable_heartbeat_rss_eviction` flag since it dynamic. And reset
+  // state if flag has changed.
+  bool rss_eviction_enabled_flag = GetFlag(FLAGS_enable_heartbeat_rss_eviction);
+  if (eviction_state_.rss_eviction_enabled_ != rss_eviction_enabled_flag) {
+    eviction_state_.global_rss_memory_at_prev_eviction =
+        eviction_state_.deleted_bytes_before_rss_update = 0;
+    eviction_state_.rss_eviction_enabled_ = rss_eviction_enabled_flag;
+  }
+  if (eviction_state_.rss_eviction_enabled_) {
+    // Calculate how much rss memory is used by all shards
+    const size_t global_used_rss_memory = rss_mem_current.load(memory_order_relaxed);
+    auto& global_rss_memory_at_prev_eviction = eviction_state_.global_rss_memory_at_prev_eviction;
+    auto& deleted_bytes_before_rss_update = eviction_state_.deleted_bytes_before_rss_update;
+    if (global_used_rss_memory < eviction_state_.global_rss_memory_at_prev_eviction) {
+      auto decrease_delete_bytes_before_rss_update =
+          std::min(deleted_bytes_before_rss_update,
+                   (global_rss_memory_at_prev_eviction - global_used_rss_memory) / shards_count);
+      VLOG_EVERY_N(1, 50) << "deleted_bytes_before_rss_update: " << deleted_bytes_before_rss_update
+                          << " decrease_delete_bytes_before_rss_update: "
+                          << decrease_delete_bytes_before_rss_update;
+      deleted_bytes_before_rss_update -= decrease_delete_bytes_before_rss_update;
+    }
+
+    global_rss_memory_at_prev_eviction = global_used_rss_memory;
+
+    // Try to evict more bytes if we are close to the rss memory limit
+    const size_t rss_goal_bytes = CalculateHowManyBytesToEvictOnShard(
+        limit, global_used_rss_memory - deleted_bytes_before_rss_update * shards_count,
+        shard_memory_budget_threshold);
+
+    VLOG_IF_EVERY_N(1, rss_goal_bytes > 0, 50)
+        << "Rss memory goal bytes: " << rss_goal_bytes
+        << ", rss used memory: " << global_used_rss_memory << ", rss memory limit: " << limit
+        << ", deleted_bytes_before_rss_update: " << deleted_bytes_before_rss_update;
+
+    goal_bytes = std::max(goal_bytes, rss_goal_bytes);
+  }
+  return goal_bytes;
 }
 
 void EngineShard::CacheStats() {

src/server/engine_shard.h

@@ -225,6 +225,12 @@ class EngineShard {
     void ResetScanState();
   };
 
+  struct EvictionTaskState {
+    bool rss_eviction_enabled_ = true;
+    size_t deleted_bytes_before_rss_update = 0;
+    size_t global_rss_memory_at_prev_eviction = 0;
+  };
+
   EngineShard(util::ProactorBase* pb, mi_heap_t* heap);
 
   // blocks the calling fiber.
@@ -236,6 +242,9 @@ class EngineShard {
   void Heartbeat();
   void RetireExpiredAndEvict();
 
+  /* Calculates the number of bytes to evict based on memory and rss memory usage. */
+  size_t CalculateEvictionBytes();
+
   void CacheStats();
 
   // We are running a task that checks whether we need to
@@ -275,6 +284,7 @@ class EngineShard {
   IntentLock shard_lock_;
 
   uint32_t defrag_task_ = 0;
+  EvictionTaskState eviction_state_;  // Used on eviction fiber
   util::fb2::Fiber fiber_heartbeat_periodic_;
   util::fb2::Done fiber_heartbeat_periodic_done_;
 

src/server/main_service.cc

@@ -916,6 +916,7 @@ void Service::Init(util::AcceptServer* acceptor, std::vector<facade::Listener*>
   config_registry.Register("dbnum");  // equivalent to databases in redis.
   config_registry.Register("dir");
   config_registry.RegisterMutable("enable_heartbeat_eviction");
+  config_registry.RegisterMutable("enable_heartbeat_rss_eviction");
   config_registry.RegisterMutable("masterauth");
   config_registry.RegisterMutable("masteruser");
   config_registry.RegisterMutable("max_eviction_per_heartbeat");

src/server/test_utils.cc

@@ -32,6 +32,8 @@ ABSL_DECLARE_FLAG(double, rss_oom_deny_ratio);
 ABSL_DECLARE_FLAG(uint32_t, num_shards);
 ABSL_FLAG(bool, force_epoll, false, "If true, uses epoll api instead iouring to run tests");
 ABSL_DECLARE_FLAG(uint32_t, acllog_max_len);
+ABSL_DECLARE_FLAG(bool, enable_heartbeat_rss_eviction);
+
 namespace dfly {
 
 namespace {
@@ -185,6 +187,8 @@ void BaseFamilyTest::SetUpTestSuite() {
 
   absl::SetFlag(&FLAGS_rss_oom_deny_ratio, -1);
   absl::SetFlag(&FLAGS_dbfilename, "");
+  // We don't want rss eviction
+  absl::SetFlag(&FLAGS_enable_heartbeat_rss_eviction, false);
 
   static bool init = true;
   if (exchange(init, false)) {

tests/dragonfly/memory_test.py

@@ -186,52 +186,61 @@ async def test_eval_with_oom(df_factory: DflyInstanceFactory):
     assert rss_before_eval * 1.01 > info["used_memory_rss"]
 
 
-@pytest.mark.skip("rss eviction disabled")
-@pytest.mark.asyncio
-@dfly_args(
-    {
-        "proactor_threads": 1,
-        "cache_mode": "true",
-        "maxmemory": "5gb",
-        "rss_oom_deny_ratio": 0.8,
-        "max_eviction_per_heartbeat": 100,
-    }
-)
-async def test_cache_eviction_with_rss_deny_oom(
-    async_client: aioredis.Redis,
-):
-    """
-    Test to verify that cache eviction is triggered even if used memory is small but rss memory is above limit
-    """
-
-    max_memory = 5 * 1024 * 1024 * 1024  # 5G
-    rss_max_memory = int(max_memory * 0.8)
+@pytest.mark.parametrize("heartbeat_rss_eviction", [True, False])
+async def test_eviction_on_rss_treshold(df_factory: DflyInstanceFactory, heartbeat_rss_eviction):
+    max_memory = 1024 * 1024**2  # 10242mb
+
+    df_server = df_factory.create(
+        proactor_threads=3,
+        cache_mode="yes",
+        maxmemory=max_memory,
+        enable_heartbeat_eviction="false",
+        enable_heartbeat_rss_eviction=heartbeat_rss_eviction,
+    )
+    df_server.start()
+    client = df_server.client()
 
-    data_fill_size = int(0.9 * rss_max_memory)  # 95% of rss_max_memory
+    data_fill_size = int(0.70 * max_memory)  # 70% of max_memory
 
     val_size = 1024 * 5  # 5 kb
     num_keys = data_fill_size // val_size
 
-    await async_client.execute_command("DEBUG", "POPULATE", num_keys, "key", val_size)
-    # Test that used memory is less than 90% of max memory
-    memory_info = await async_client.info("memory")
-    assert (
-        memory_info["used_memory"] < max_memory * 0.9
-    ), "Used memory should be less than 90% of max memory."
-    assert (
-        memory_info["used_memory_rss"] > rss_max_memory * 0.9
-    ), "RSS memory should be less than 90% of rss max memory (max_memory * rss_oom_deny_ratio)."
-
-    # Get RSS memory after creating new connections
-    memory_info = await async_client.info("memory")
-    while memory_info["used_memory_rss"] > rss_max_memory * 0.9:
-        await asyncio.sleep(1)
-        memory_info = await async_client.info("memory")
-        logging.info(
-            f'Current rss: {memory_info["used_memory_rss"]}. rss eviction threshold: {rss_max_memory * 0.9}.'
-        )
-        stats_info = await async_client.info("stats")
-        logging.info(f'Current evicted: {stats_info["evicted_keys"]}. Total keys: {num_keys}.')
+    await client.execute_command("DEBUG", "POPULATE", num_keys, "key", val_size)
+
+    # Create huge list which can be used with LRANGE to increase RSS memory only
+    for name in ["list_1", "list_2"]:
+        for i in range(1, 1000):
+            rand_str = "".join(random.choices(string.ascii_letters, k=val_size))
+            await client.execute_command(f"LPUSH {name} {rand_str}")
+
+    # Make them STICK so we don't evict them
+    await client.execute_command(f"STICK list_1")
+    await client.execute_command(f"STICK list_2")
+
+    await client.execute_command("CONFIG SET enable_heartbeat_eviction true")
+
+    memory_info_before = await client.info("memory")
+
+    # This will increase only RSS memory above treshold
+    p = client.pipeline()
+    for _ in range(50):
+        p.execute_command("LRANGE list_1 0 -1")
+        p.execute_command("LRANGE list_2 0 -1")
+    await p.execute()
+
+    # Wait for some time
+    await asyncio.sleep(3)
+    memory_info_after = await client.info("memory")
+    stats_info_after = await client.info("stats")
+
+    if heartbeat_rss_eviction:
+        # We should see used memory deacrease and number of some number of evicted keys
+        assert memory_info_after["used_memory"] < memory_info_before["used_memory"]
+        assert stats_info_after["evicted_keys"]
+    else:
+        # If heartbeat rss eviction is disabled there should be no chage
+        assert memory_info_after["used_memory"] == memory_info_before["used_memory"]
+        assert stats_info_after["evicted_keys"] == 0
 
 
 @pytest.mark.asyncio