feat: SHUTDOWN: fast path for NOW/FORCE; unify SAVE/SAFE; support NOSAVE (#5783)

Fixed: #5444

Author: vyavdoshenko

src/facade/dragonfly_listener.cc

@@ -66,6 +66,9 @@ CONFIG_enum(tls_auth_clients, "yes", "", tls_auth_clients_enum, tls_auth_clients
 
 namespace facade {
 
+// See dragonfly_listener.h
+std::atomic<bool> g_shutdown_fast{false};
+
 using namespace util;
 using util::detail::SafeErrorMessage;
 
@@ -252,7 +255,12 @@ bool Listener::IsMainInterface() const {
 }
 
 void Listener::PreShutdown() {
-  // Iterate on all connections and allow them to finish their commands for
+  // If NOW/FORCE requested, expedite shutdown without waiting.
+  if (g_shutdown_fast.load(std::memory_order_acquire)) {
+    return;
+  }
+
+  // Otherwise: Iterate on all connections and allow them to finish their commands for
   // a short period.
   // Executed commands can be visible in snapshots or replicas, but if we close the client
   // connections too fast we might not send the acknowledgment for those commands.

src/facade/dragonfly_listener.h

@@ -7,6 +7,7 @@
 #include <absl/base/internal/spinlock.h>
 #include <absl/time/time.h>
 
+#include <atomic>
 #include <memory>
 #include <system_error>
 #include <utility>
@@ -107,4 +108,9 @@ class DispatchTracker {
   bool ignore_blocked_;
 };
 
+// Global shutdown tuning flag, controlled by SHUTDOWN options.
+// When true, listeners perform expedited shutdown without waiting for
+// in-flight dispatches (used by NOW/FORCE).
+extern std::atomic<bool> g_shutdown_fast;
+
 }  // namespace facade

src/server/server_family.cc

@@ -3781,24 +3781,60 @@ void ServerFamily::Latency(CmdArgList args, const CommandContext& cmd_cntx) {
 }
 
 void ServerFamily::ShutdownCmd(CmdArgList args, const CommandContext& cmd_cntx) {
-  if (args.size() > 1) {
+  // Supported options (case-insensitive):
+  // SAVE | NOSAVE, NOW, FORCE, ABORT, SAFE (Valkey-specific, the same as SAVE in Dragonfly)
+  enum ShutBits : uint32_t {
+    SB_SAVE = 1u << 0,
+    SB_NOSAVE = 1u << 1,
+    SB_NOW = 1u << 2,
+    SB_FORCE = 1u << 3,
+    SB_ABORT = 1u << 4,
+  };
+
+  uint32_t sb = 0;
+
+  CmdArgParser parser(args);
+  while (parser.HasNext()) {
+    // Map SAFE to SAVE directly (fallthrough behavior)
+    ShutBits opt = parser.MapNext("SAVE", SB_SAVE, "NOSAVE", SB_NOSAVE, "NOW", SB_NOW, "FORCE",
+                                  SB_FORCE, "ABORT", SB_ABORT, "SAFE", SB_SAVE);
+    sb |= static_cast<uint32_t>(opt);
+  }
+
+  if (auto perr = parser.TakeError(); perr) {
+    return cmd_cntx.rb->SendError(perr.MakeReply());
+  }
+
+  // Conflicting toggles
+  if ((sb & SB_SAVE) && (sb & SB_NOSAVE)) {
     cmd_cntx.rb->SendError(kSyntaxErr);
     return;
   }
 
-  if (args.size() == 1) {
-    auto sub_cmd = ArgS(args, 0);
-    if (absl::EqualsIgnoreCase(sub_cmd, "SAVE")) {
-    } else if (absl::EqualsIgnoreCase(sub_cmd, "NOSAVE")) {
-      save_on_shutdown_ = false;
-    } else {
-      cmd_cntx.rb->SendError(kSyntaxErr);
-      return;
-    }
+  if (sb & SB_ABORT) {
+    // We currently do not support aborting an in-progress shutdown sequence.
+    cmd_cntx.rb->SendError("SHUTDOWN ABORT is not supported");
+    return;
   }
 
+  // Configure save behavior on shutdown according to options.
+  if (sb & SB_FORCE) {
+    // FORCE implies no snapshot on shutdown regardless of SAVE/SAFE
+    save_on_shutdown_ = false;
+  } else if (sb & SB_NOSAVE) {
+    save_on_shutdown_ = false;
+  } else if (sb & SB_SAVE) {
+    save_on_shutdown_ = true;
+  }
+
+  // Wire NOW/FORCE to a single fast-shutdown flag for listeners.
+  facade::g_shutdown_fast.store((sb & (SB_NOW | SB_FORCE)) != 0, std::memory_order_seq_cst);
+
   CHECK_NOTNULL(acceptor_)->Stop();
   cmd_cntx.rb->SendOk();
+
+  // Reset flag for any subsequent restarts (mainly for tests).
+  facade::g_shutdown_fast.store(false, std::memory_order_seq_cst);
 }
 
 void ServerFamily::Dfly(CmdArgList args, const CommandContext& cmd_cntx) {

tests/dragonfly/shutdown_test.py

@@ -54,3 +54,83 @@ async def delayed_takeover():
             assert acknowleged_value == int(value_from_snapshot)
 
         await client.connection_pool.disconnect()
+
+
+@dfly_args({"proactor_threads": "2"})
+class TestShutdownOptions:
+    @pytest.mark.asyncio
+    async def test_shutdown_abort_and_invalid_option(self, df_factory):
+        df_args = {"dbfilename": "dump", **BASIC_ARGS, "port": 1121}
+        df_server = df_factory.create(**df_args)
+        df_server.start()
+
+        client = aioredis.Redis(port=df_server.port)
+
+        # ABORT should be rejected and server should remain responsive
+        with pytest.raises(redis.exceptions.ResponseError):
+            await client.execute_command("SHUTDOWN ABORT")
+
+        pong = await client.ping()
+        assert pong is True
+
+        # Invalid option -> syntax error
+        with pytest.raises(redis.exceptions.ResponseError):
+            await client.execute_command("SHUTDOWN FOO")
+
+        await client.connection_pool.disconnect()
+        df_server.stop()
+
+    @pytest.mark.asyncio
+    async def test_shutdown_safe_persists_snapshot(self, df_factory, tmp_path):
+        # Ensure snapshot dir exists and is used
+        snap_dir = tmp_path
+        df_args = {"dbfilename": "dump", "dir": str(snap_dir) + "/", "port": 1122}
+
+        df_server = df_factory.create(**df_args)
+        df_server.start()
+
+        client = aioredis.Redis(port=df_server.port)
+        await client.set("safe_key", "safe_value")
+
+        # SHUTDOWN SAFE should save synchronously and then stop
+        try:
+            await client.execute_command("SHUTDOWN SAFE")
+        except Exception:
+            # Connection may be dropped as part of shutdown; this is acceptable
+            pass
+
+        await client.connection_pool.disconnect()
+
+        # Restart and verify data persisted
+        df_server.start()
+        client = aioredis.Redis(port=df_server.port)
+        val = await client.get("safe_key")
+        assert val == b"safe_value"
+        await client.connection_pool.disconnect()
+        df_server.stop()
+
+    @pytest.mark.asyncio
+    async def test_shutdown_save_persists_snapshot(self, df_factory, tmp_path):
+        # SAVE should follow the same synchronous path as SAFE
+        snap_dir = tmp_path
+        df_args = {"dbfilename": "dump", "dir": str(snap_dir) + "/", "port": 1123}
+
+        df_server = df_factory.create(**df_args)
+        df_server.start()
+
+        client = aioredis.Redis(port=df_server.port)
+        await client.set("save_key", "save_value")
+
+        try:
+            await client.execute_command("SHUTDOWN SAVE")
+        except Exception:
+            pass
+
+        await client.connection_pool.disconnect()
+
+        df_server.start()
+        client = aioredis.Redis(port=df_server.port)
+        val = await client.get("save_key")
+        assert val == b"save_value"
+        await client.connection_pool.disconnect()
+        df_server.stop()