diff --git a/src/server/detail/save_stages_controller.cc b/src/server/detail/save_stages_controller.cc index a6fa0525dff6..b6d5a5728b87 100644 --- a/src/server/detail/save_stages_controller.cc +++ b/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 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); } diff --git a/tests/dragonfly/snapshot_test.py b/tests/dragonfly/snapshot_test.py index 9241e9eccc54..0ecd51702da5 100644 --- a/tests/dragonfly/snapshot_test.py +++ b/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(