From f37c57c704188a3319140f6fe614faf595daf1da Mon Sep 17 00:00:00 2001
From: adiholden <adi@dragonflydb.io>
Date: Sun, 7 Jan 2024 12:21:09 +0200
Subject: [PATCH] fix(server): crash on rename save command on background save
 (#2375)

* fix(server): crash on rename save command on baground save

Signed-off-by: adi_holden <adi@dragonflydb.io>
---
 src/server/main_service.cc       |  2 +-
 tests/dragonfly/snapshot_test.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index a522dc059..2e5f60634 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -1463,7 +1463,7 @@ facade::ConnectionContext* Service::CreateContext(util::FiberSocketBase* peer,
 }
 
 const CommandId* Service::FindCmd(std::string_view cmd) const {
-  return registry_.Find(cmd);
+  return registry_.Find(registry_.RenamedOrOriginal(cmd));
 }
 
 bool Service::IsLocked(DbIndex db_index, std::string_view key) const {
diff --git a/tests/dragonfly/snapshot_test.py b/tests/dragonfly/snapshot_test.py
index 9a3867fa8..9241e9ecc 100644
--- a/tests/dragonfly/snapshot_test.py
+++ b/tests/dragonfly/snapshot_test.py
@@ -196,6 +196,35 @@ class TestSetsnapshot_cron(SnapshotTestBase):
         assert super().get_main_file("test-set-snapshot_cron-summary.dfs")
 
 
+@dfly_args(
+    {**BASIC_ARGS, "dbfilename": "test-save-rename-command", "rename_command": "save=save-foo"}
+)
+class TestSnapshotShutdownWithRenameSave(SnapshotTestBase):
+    """Test set snapshot_cron flag"""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, tmp_dir: Path):
+        super().setup(tmp_dir)
+
+    @pytest.mark.asyncio
+    async def test_snapshot(self, df_server, df_seeder_factory):
+        """Checks that on shutdown we save snapshot"""
+        seeder = df_seeder_factory.create(port=df_server.port)
+        await seeder.run(target_deviation=0.1)
+
+        start_capture = await seeder.capture()
+        a_client = aioredis.Redis(port=df_server.port)
+
+        df_server.stop()
+        df_server.start()
+
+        a_client = aioredis.Redis(port=df_server.port)
+        await wait_available_async(a_client)
+        await a_client.connection_pool.disconnect()
+
+        assert await seeder.compare(start_capture, port=df_server.port)
+
+
 @dfly_args({**BASIC_ARGS})
 class TestOnlyOneSaveAtATime(SnapshotTestBase):
     """Dragonfly does not allow simultaneous save operations, send 2 save operations and make sure one is rejected"""