fix(dragonfly_test): Fix DflyEngineTest.Bug207

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

src/server/dragonfly_test.cc

@@ -27,6 +27,7 @@ ABSL_DECLARE_FLAG(float, mem_defrag_waste_threshold);
 ABSL_DECLARE_FLAG(uint32_t, mem_defrag_check_sec_interval);
 ABSL_DECLARE_FLAG(std::vector<std::string>, rename_command);
 ABSL_DECLARE_FLAG(bool, lua_resp2_legacy_float);
+ABSL_DECLARE_FLAG(double, eviction_memory_budget_threshold);
 
 namespace dfly {
 
@@ -455,14 +456,19 @@ TEST_F(DflyEngineTest, OOM) {
 /// Reproduces the case where items with expiry data were evicted,
 /// and then written with the same key.
 TEST_F(DflyEngineTest, Bug207) {
-  max_memory_limit = 1200000;  // 1.2mb = 300000 * 4
+  /* Sometimes filling the cache is much faster than the first heartbeat, leading to OOM.
+     To prevent OOM during cache filling, we need to increase the max_memory_limit and eviction
+     threshold. */
+  max_memory_limit = 1350000;  // 1.35mb
+  absl::FlagSaver fs;
+  absl::SetFlag(&FLAGS_eviction_memory_budget_threshold, 0.6);
+
   shard_set->TEST_EnableCacheMode();
 
   ssize_t i = 0;
   RespExpr resp;
   for (; i < 10000; ++i) {
     resp = Run({"setex", StrCat("key", i), "30", "bar"});
-    // we evict some items because 5000 is too much when max_memory_limit is 300000.
     ASSERT_EQ(resp, "OK");
   }
 

src/server/engine_shard.cc

@@ -60,6 +60,10 @@ ABSL_FLAG(float, tiered_offload_threshold, 0.5,
 ABSL_FLAG(bool, enable_heartbeat_eviction, true,
           "Enable eviction during heartbeat when memory is under pressure.");
 
+ABSL_FLAG(double, eviction_memory_budget_threshold, 0.1,
+          "Eviction starts when free memory (including rss memory) drops below "
+          "eviction_memory_budget_threshold * max_memory_limit");
+
 namespace dfly {
 
 using absl::GetFlag;
@@ -198,11 +202,6 @@ optional<uint32_t> GetPeriodicCycleMs() {
   return clock_cycle_ms;
 }
 
-/* Eviction begins when the shard's free memory falls below
-  (max_memory_limit * kEvictionMemoryThresholdFactor) / <number of shards>.
-  Same logic for the RSS memory. */
-constexpr double kEvictionMemoryThresholdFactor = 0.1;
-
 ssize_t CalculateMemoryBudget(size_t max_memory, size_t used_memory) {
   if (max_memory >= used_memory) {
     return max_memory - used_memory;
@@ -749,15 +748,17 @@ void EngineShard::RetireExpiredAndEvict() {
 
   // We start eviction when we have less than 10% of free memory
   const ssize_t shard_memory_budget_threshold =
-      ssize_t(max_memory_limit * kEvictionMemoryThresholdFactor) / shards_count;
+      ssize_t(max_memory_limit * absl::GetFlag(FLAGS_eviction_memory_budget_threshold)) /
+      shards_count;
 
   size_t max_rss_memory;
   ssize_t shard_rss_memory_budget_threshold;  // Threshold for free rss memory
   if (rss_oom_deny_ratio > 0.0) {
     max_rss_memory = size_t(rss_oom_deny_ratio * max_memory_limit);
     // We start eviction when we have less than 10% of free rss memory
     shard_rss_memory_budget_threshold =
-        ssize_t(max_rss_memory * kEvictionMemoryThresholdFactor) / shards_count;
+        ssize_t(max_rss_memory * absl::GetFlag(FLAGS_eviction_memory_budget_threshold)) /
+        shards_count;
   } else {
     max_rss_memory = std::numeric_limits<size_t>::max();
     // We should never evict based on rss memory