fix server: debug populate consume less memory

Signed-off-by: adi_holden <[email protected]>
dragonflydb · Dec 30, 2024 · 84398b9 · 84398b9
1 parent c4be62e
commit 84398b9
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 21 deletions.
diff --git a/src/server/debugcmd.cc b/src/server/debugcmd.cc
@@ -161,26 +161,31 @@ void DoPopulateBatch(string_view type, string_view prefix, size_t val_size, bool
   absl::InsecureBitGen gen;
   for (unsigned i = 0; i < batch.sz; ++i) {
     string key = absl::StrCat(prefix, ":", batch.index[i]);
+    int32_t elements_left = elements;
+    while (elements_left) {
+      int32_t populate_elements = std::min(5, elements_left);
+      elements_left -= populate_elements;
+      auto [cid, args] =
+          GeneratePopulateCommand(type, key, val_size, random_value, populate_elements,
+                                  *sf->service().mutable_registry(), &gen);
+      if (!cid) {
+        LOG_EVERY_N(WARNING, 10'000) << "Unable to find command, was it renamed?";
+        break;
+      }
 
-    auto [cid, args] = GeneratePopulateCommand(type, std::move(key), val_size, random_value,
-                                               elements, *sf->service().mutable_registry(), &gen);
-    if (!cid) {
-      LOG_EVERY_N(WARNING, 10'000) << "Unable to find command, was it renamed?";
-      break;
-    }
-
-    args_view.clear();
-    for (auto& arg : args) {
-      args_view.push_back(arg);
-    }
-    auto args_span = absl::MakeSpan(args_view);
+      args_view.clear();
+      for (auto& arg : args) {
+        args_view.push_back(arg);
+      }
+      auto args_span = absl::MakeSpan(args_view);
 
-    stub_tx->MultiSwitchCmd(cid);
-    local_cntx.cid = cid;
-    crb.SetReplyMode(ReplyMode::NONE);
-    stub_tx->InitByArgs(cntx->ns, local_cntx.conn_state.db_index, args_span);
+      stub_tx->MultiSwitchCmd(cid);
+      local_cntx.cid = cid;
+      crb.SetReplyMode(ReplyMode::NONE);
+      stub_tx->InitByArgs(cntx->ns, local_cntx.conn_state.db_index, args_span);
 
-    sf->service().InvokeCmd(cid, args_span, &crb, &local_cntx);
+      sf->service().InvokeCmd(cid, args_span, &crb, &local_cntx);
+    }
   }
 
   local_tx->UnlockMulti();

diff --git a/tests/dragonfly/snapshot_test.py b/tests/dragonfly/snapshot_test.py
@@ -591,17 +591,16 @@ async def test_big_value_serialization_memory_limit(df_factory, cont_type):
     await client.execute_command(
         f"debug populate 1 prefix {element_size} TYPE {cont_type} RAND ELEMENTS {elements}"
     )
+    await asyncio.sleep(1)
 
     info = await client.info("ALL")
-    # rss double's because of DEBUG POPULATE
-    assert info["used_memory_peak_rss"] > (one_gb * 2)
+    assert info["used_memory_peak_rss"] < (one_gb * 1.2)
     # if we execute SAVE below without big value serialization we trigger the assertion below.
     # note the peak would reach (one_gb * 3) without it.
     await client.execute_command("SAVE")
     info = await client.info("ALL")
 
-    upper_limit = 2_250_000_000  # 2.25 GB
-    assert info["used_memory_peak_rss"] < upper_limit
+    assert info["used_memory_peak_rss"] < (one_gb * 1.3)
 
     await client.execute_command("FLUSHALL")
     await client.close()