Implement graceful shutdown for Garnet server ( #1382 )

hosting/Windows/Garnet.worker/Worker.cs

@@ -43,8 +43,24 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
         /// <param name="cancellationToken">Indicates that the shutdown process should no longer be graceful.</param>
         public override async Task StopAsync(CancellationToken cancellationToken)
         {
-            Dispose();
-            await base.StopAsync(cancellationToken).ConfigureAwait(false);
+            try
+            {
+                if (server != null)
+                {
+                    // Perform graceful shutdown with AOF commit and checkpoint
+                    await server.ShutdownAsync(timeout: TimeSpan.FromSeconds(5), token: cancellationToken).ConfigureAwait(false);
+                }
+            }
+            catch (OperationCanceledException)
+            {
+                // Force shutdown requested - proceed to dispose
+            }
+            finally
+            {
+                // Ensure base class cleanup although cancellationToken is cancelled
+                await base.StopAsync(CancellationToken.None).ConfigureAwait(false);
+                Dispose();
+            }
         }
 
         public override void Dispose()
@@ -55,6 +71,8 @@ public override void Dispose()
             }
             server?.Dispose();
             _isDisposed = true;
+            base.Dispose();
+            GC.SuppressFinalize(this);
         }
     }
 }

libs/host/GarnetServer.cs

@@ -10,6 +10,7 @@
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 using Garnet.cluster;
 using Garnet.common;
 using Garnet.networking;
@@ -431,6 +432,179 @@ public void Start()
                 Console.WriteLine("* Ready to accept connections");
         }
 
+        /// <summary>
+        /// Performs graceful shutdown of the server.
+        /// Stops accepting new connections, waits for active connections to complete, commits AOF, and takes checkpoint if needed.
+        /// </summary>
+        /// <param name="timeout">Timeout for waiting on active connections (default: 30 seconds)</param>
+        /// <param name="token">Cancellation token</param>
+        /// <returns>Task representing the async shutdown operation</returns>
+        public async Task ShutdownAsync(TimeSpan? timeout = null, CancellationToken token = default)
+        {
+            var shutdownTimeout = timeout ?? TimeSpan.FromSeconds(30);
+
+            try
+            {
+                // Stop accepting new connections first
+                StopListening();
+
+                // Wait for existing connections to complete
+                await WaitForActiveConnectionsAsync(shutdownTimeout, token).ConfigureAwait(false);
+
+                // Commit AOF and take checkpoint if needed
+                await FinalizeDataAsync(token).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                // Force shutdown requested
+            }
+            catch (Exception ex)
+            {
+                logger?.LogError(ex, "Error during graceful shutdown");
+            }
+        }
+
+        /// <summary>
+        /// Stop all servers from accepting new connections.
+        /// </summary>
+        private void StopListening()
+        {
+            if (servers == null) return;
+
+            logger?.LogDebug("Stopping listeners to prevent new connections...");
+            foreach (var server in servers)
+            {
+                try
+                {
+                    server?.StopListening();
+                }
+                catch (Exception ex)
+                {
+                    logger?.LogWarning(ex, "Error stopping listener");
+                }
+            }
+        }
+
+        /// <summary>
+        /// Waits for active connections to complete within the specified timeout.
+        /// </summary>
+        private async Task WaitForActiveConnectionsAsync(TimeSpan timeout, CancellationToken token)
+        {
+            if (servers == null) return;
+
+            // Linked Token : between external token and timeout
+            using var cts = CancellationTokenSource.CreateLinkedTokenSource(token);
+            cts.CancelAfter(timeout);
+
+            var delays = new[] { 50, 300, 1000 };
+            var delayIndex = 0;
+
+            try
+            {
+                while (!cts.Token.IsCancellationRequested)
+                {
+                    var activeConnections = GetActiveConnectionCount();
+                    if (activeConnections == 0)
+                    {
+                        logger?.LogInformation("All connections have been closed gracefully.");
+                        return;
+                    }
+
+                    logger?.LogInformation("Waiting for {ActiveConnections} active connections to complete...", activeConnections);
+
+                    var currentDelay = delays[delayIndex];
+                    if (delayIndex < delays.Length - 1) delayIndex++;
+
+                    await Task.Delay(currentDelay, cts.Token).ConfigureAwait(false);
+                }
+            }
+            catch (OperationCanceledException) when (token.IsCancellationRequested)
+            {
+                throw;
+            }
+            catch (OperationCanceledException)
+            {
+                // timeout reached error logging
+                logger?.LogWarning("Timeout reached after {TimeoutSeconds} seconds. Some connections may still be active.",
+                    timeout.TotalSeconds);
+            }
+            catch (Exception ex)
+            {
+                logger?.LogWarning(ex, "Error checking active connections");
+                delayIndex = 0;
+                await Task.Delay(500, token).ConfigureAwait(false);
+            }
+        }
+
+        /// <summary>
+        /// Gets the current number of active connections directly from server instances.
+        /// </summary>
+        private long GetActiveConnectionCount()
+        {
+            long count = 0;
+            if (servers != null)
+            {
+                foreach (var garnetServer in servers)
+                {
+                    if (garnetServer is GarnetServerBase garnetServerBase)
+                    {
+                        count += garnetServerBase.get_conn_active();
+                    }
+                }
+            }
+            return count;
+        }
+
+        /// <summary>
+        /// Commits AOF and takes checkpoint for data durability during shutdown.
+        /// </summary>
+        private async Task FinalizeDataAsync(CancellationToken token)
+        {
+            // Commit AOF before checkpoint/shutdown
+            if (opts.EnableAOF)
+            {
+                logger?.LogDebug("Committing AOF before shutdown...");
+                try
+                {
+                    var commitSuccess = await Store.CommitAOFAsync(token).ConfigureAwait(false);
+                    if (commitSuccess)
+                    {
+                        logger?.LogDebug("AOF committed successfully.");
+                    }
+                    else
+                    {
+                        logger?.LogInformation("AOF commit skipped (another commit in progress or replica mode).");
+                    }
+                }
+                catch (Exception ex)
+                {
+                    logger?.LogError(ex, "Error committing AOF during shutdown");
+                }
+            }
+
+            // Take checkpoint for tiered storage
+            if (opts.EnableStorageTier)
+            {
+                logger?.LogDebug("Taking checkpoint for tiered storage...");
+                try
+                {
+                    var checkpointSuccess = Store.TakeCheckpoint(background: false, token);
+                    if (checkpointSuccess)
+                    {
+                        logger?.LogDebug("Checkpoint completed successfully.");
+                    }
+                    else
+                    {
+                        logger?.LogInformation("Checkpoint skipped (another checkpoint in progress or replica mode).");
+                    }
+                }
+                catch (Exception ex)
+                {
+                    logger?.LogError(ex, "Error taking checkpoint during shutdown");
+                }
+            }
+        }
+
         /// <summary>
         /// Dispose store (including log and checkpoint directory)
         /// </summary>

libs/server/Servers/GarnetServerBase.cs

@@ -154,6 +154,12 @@ public bool AddSession(WireFormat protocol, ref ISessionProvider provider, INetw
         /// <inheritdoc />
         public abstract void Start();
 
+        /// <inheritdoc />
+        public virtual void StopListening()
+        {
+            // Base implementation does nothing; derived classes should override
+        }
+
         /// <inheritdoc />
         public virtual void Dispose()
         {

libs/server/Servers/GarnetServerTcp.cs

@@ -28,6 +28,7 @@ public class GarnetServerTcp : GarnetServerBase, IServerHook
         readonly int networkConnectionLimit;
         readonly string unixSocketPath;
         readonly UnixFileMode unixSocketPermission;
+        volatile bool isListening;
 
         /// <inheritdoc/>
         public override IEnumerable<IMessageConsumer> ActiveConsumers()
@@ -117,19 +118,43 @@ public override void Start()
             }
 
             listenSocket.Listen(512);
+            isListening = true;
             if (!listenSocket.AcceptAsync(acceptEventArg))
                 AcceptEventArg_Completed(null, acceptEventArg);
         }
 
+        /// <inheritdoc />
+        public override void StopListening()
+        {
+            if (!isListening)
+                return;
+
+            isListening = false;
+            try
+            {
+                // Close the listen socket to stop accepting new connections
+                // This will cause any pending AcceptAsync to complete with an error
+                listenSocket.Close();
+                logger?.LogDebug("Stopped accepting new connections on {endpoint}", EndPoint);
+            }
+            catch (Exception ex)
+            {
+                logger?.LogDebug(ex, "Error closing listen socket on {endpoint}", EndPoint);
+            }
+        }
+
         private void AcceptEventArg_Completed(object sender, SocketAsyncEventArgs e)
         {
             try
             {
                 do
                 {
+                    // Check isListening flag before processing and before calling AcceptAsync again
+                    if (!isListening) break;
+
                     if (!HandleNewConnection(e)) break;
                     e.AcceptSocket = null;
-                } while (!listenSocket.AcceptAsync(e));
+                } while (isListening && !listenSocket.AcceptAsync(e));
             }
             // socket disposed
             catch (ObjectDisposedException) { }

libs/server/Servers/IGarnetServer.cs

@@ -46,5 +46,11 @@ public interface IGarnetServer : IDisposable
         /// Start server
         /// </summary>
         public void Start();
+
+        /// <summary>
+        /// Stop accepting new connections (for graceful shutdown).
+        /// Existing connections remain active until they complete or are disposed.
+        /// </summary>
+        public void StopListening();
     }
 }

libs/server/Servers/StoreApi.cs

@@ -130,6 +130,35 @@ public bool FlushDB(int dbId = 0, bool unsafeTruncateLog = false)
             }
         }
 
+        /// <summary>
+        /// Take checkpoint for all active databases
+        /// </summary>
+        /// <param name="background">True if method can return before checkpoint is taken</param>
+        /// <param name="token">Cancellation token</param>
+        /// <returns>false if checkpoint was skipped due to node state or another checkpoint in progress</returns>
+        public bool TakeCheckpoint(bool background = false, CancellationToken token = default)
+        {
+            using (PreventRoleChange(out var acquired))
+            {
+                if (!acquired || IsReplica)
+                {
+                    return false;
+                }
+
+                return storeWrapper.TakeCheckpoint(background, logger: null, token: token);
+            }
+        }
+
+        /// <summary>
+        /// Check if storage tier is enabled
+        /// </summary>
+        public bool IsStorageTierEnabled => storeWrapper.serverOptions.EnableStorageTier;
+
+        /// <summary>
+        /// Check if AOF is enabled
+        /// </summary>
+        public bool IsAOFEnabled => storeWrapper.serverOptions.EnableAOF;
+
         /// <summary>
         /// Helper to disable role changes during a using block.
         /// 

main/GarnetServer/Program.cs

@@ -22,7 +22,18 @@ static void Main(string[] args)
                 // Start the server
                 server.Start();
 
-                Thread.Sleep(Timeout.Infinite);
+                var shutdownEvent = new ManualResetEventSlim(false);
+
+                Console.CancelKeyPress += (sender, e) =>
+                {
+                    e.Cancel = true;
+                    // Graceful shutdown: drain connections, commit AOF, take checkpoint
+                    server.ShutdownAsync(TimeSpan.FromSeconds(5))
+                        .GetAwaiter().GetResult();
+                    shutdownEvent.Set();
+                };
+
+                shutdownEvent.Wait();
             }
             catch (Exception ex)
             {