fix(bug): server crash info memory while saving (#2637)

* bug: crash info memory while saving

Signed-off-by: adi_holden <[email protected]>

src/server/detail/save_stages_controller.cc

@@ -125,6 +125,7 @@ error_code RdbSnapshot::SaveBody() {
 }
 
 size_t RdbSnapshot::GetSaveBuffersSize() {
+  CHECK(saver_);
   return saver_->GetTotalBuffersSize();
 }
 
@@ -172,17 +173,24 @@ SaveInfo SaveStagesController::Save() {
 
 size_t SaveStagesController::GetSaveBuffersSize() {
   std::atomic<size_t> total_bytes{0};
-  if (use_dfs_format_) {
-    auto cb = [this, &total_bytes](ShardId sid) {
-      total_bytes.fetch_add(snapshots_[sid].first->GetSaveBuffersSize(), memory_order_relaxed);
-    };
-    shard_set->RunBriefInParallel([&](EngineShard* es) { cb(es->shard_id()); });
 
-  } else {
-    // When rdb format save is running, there is only one rdb saver instance, it is running on the
-    // connection thread that runs the save command.
-    total_bytes.store(snapshots_.front().first->GetSaveBuffersSize(), memory_order_relaxed);
+  auto add_snapshot_bytes = [this, &total_bytes](ShardId sid) {
+    if (auto& snapshot = snapshots_[sid].first; snapshot && snapshot->HasStarted()) {
+      total_bytes.fetch_add(snapshot->GetSaveBuffersSize(), memory_order_relaxed);
+    }
+  };
+
+  if (snapshots_.size() > 0) {
+    if (use_dfs_format_) {
+      shard_set->RunBriefInParallel([&](EngineShard* es) { add_snapshot_bytes(es->shard_id()); });
+
+    } else {
+      // When rdb format save is running, there is only one rdb saver instance, it is running on the
+      // connection thread that runs the save command.
+      add_snapshot_bytes(0);
+    }
   }
+
   return total_bytes.load(memory_order_relaxed);
 }
 

tests/dragonfly/snapshot_test.py

@@ -384,6 +384,28 @@ async def test_snapshot(self, df_seeder_factory, df_server):
         await a_client.connection_pool.disconnect()
 
 
+@dfly_args({**BASIC_ARGS, "dbfilename": "info-while-snapshot"})
+async def test_infomemory_while_snapshoting(async_client: aioredis.Redis):
+    await async_client.execute_command("DEBUG POPULATE 10000 key 4048 RAND")
+
+    async def save():
+        await async_client.execute_command("SAVE", "DF")
+
+    save_finished = False
+
+    async def info_in_loop():
+        while not save_finished:
+            await async_client.execute_command("INFO MEMORY")
+            await asyncio.sleep(0.1)
+
+    save_task = asyncio.create_task(save())
+    info_task = asyncio.create_task(info_in_loop())
+
+    await save_task
+    save_finished = True
+    await info_task
+
+
 # If DRAGONFLY_S3_BUCKET is configured, AWS credentials must also be
 # configured.
 @pytest.mark.skipif(