Implement graceful shutdown for Garnet server ( #1382 )

hosting/Windows/Garnet.worker/Program.cs

@@ -1,16 +1,35 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+using System;
 using Garnet;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 
 class Program
 {
+    // Data finalization (AOF commit / checkpoint) uses up to 15 seconds internally (see GarnetServer.FinalizeDataAsync).
+    // Add this buffer on top of the connection-drain timeout so the host shutdown budget covers the full shutdown sequence.
+    private const int DataFinalizationBufferSeconds = 20;
+
     static void Main(string[] args)
     {
+        // Pre-parse only the shutdown-timeout argument so we can configure both
+        // the host shutdown budget and the Worker's connection-drain timeout from a single value.
+        var shutdownTimeoutSeconds = ParseShutdownTimeoutSeconds(args, defaultSeconds: 5);
+        var shutdownTimeout = TimeSpan.FromSeconds(shutdownTimeoutSeconds);
+
         var builder = Host.CreateApplicationBuilder(args);
-        builder.Services.AddHostedService(_ => new Worker(args));
+
+        // Tell the .NET host (and the Windows SCM via WindowsServiceLifetime) how long to wait
+        // before forcibly killing the process. We add DataFinalizationBufferSeconds so that AOF
+        // commit / checkpoint can complete after connection draining finishes.
+        builder.Services.Configure<HostOptions>(opts =>
+        {
+            opts.ShutdownTimeout = shutdownTimeout + TimeSpan.FromSeconds(DataFinalizationBufferSeconds);
+        });
+
+        builder.Services.AddHostedService(_ => new Worker(args, shutdownTimeout));
 
         builder.Services.AddWindowsService(options =>
         {
@@ -20,4 +39,22 @@ static void Main(string[] args)
         var host = builder.Build();
         host.Run();
     }
+
+    /// <summary>
+    /// Scans <paramref name="args"/> for <c>--shutdown-timeout &lt;value&gt;</c> and returns
+    /// the parsed integer, or <paramref name="defaultSeconds"/> if the argument is absent or invalid.
+    /// This lightweight pre-parse avoids a full CommandLineParser pass before the host is built.
+    /// </summary>
+    private static int ParseShutdownTimeoutSeconds(string[] args, int defaultSeconds)
+    {
+        for (var i = 0; i < args.Length - 1; i++)
+        {
+            if (args[i] is "--shutdown-timeout" or "-shutdown-timeout" &&
+                int.TryParse(args[i + 1], out var value) && value > 0)
+            {
+                return value;
+            }
+        }
+        return defaultSeconds;
+    }
 }

hosting/Windows/Garnet.worker/Worker.cs

@@ -12,12 +12,20 @@ public class Worker : BackgroundService
     {
         private bool _isDisposed = false;
         private readonly string[] args;
+        private readonly TimeSpan _shutdownTimeout;
 
         private GarnetServer server;
 
-        public Worker(string[] args)
+        /// <param name="args">Command line arguments forwarded to <see cref="GarnetServer"/>.</param>
+        /// <param name="shutdownTimeout">
+        /// How long to wait for active connections to drain during graceful shutdown.
+        /// Must be less than the host <see cref="Microsoft.Extensions.Hosting.HostOptions.ShutdownTimeout"/>
+        /// so that data finalization (AOF commit / checkpoint) can also complete within the host budget.
+        /// </param>
+        public Worker(string[] args, TimeSpan shutdownTimeout)
         {
             this.args = args;
+            _shutdownTimeout = shutdownTimeout;
         }
 
         protected override async Task ExecuteAsync(CancellationToken stoppingToken)
@@ -43,8 +51,26 @@ 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)
+                {
+                    // If cancellation is requested, we will skip the graceful shutdown and proceed to dispose immediately
+                    bool isForceShutdown = cancellationToken.IsCancellationRequested;
+                    // Perform graceful shutdown with AOF commit and checkpoint when not forced Shutdown From OS.
+                    await server.ShutdownAsync(timeout: _shutdownTimeout, noSave: isForceShutdown, 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 +81,8 @@ public override void Dispose()
             }
             server?.Dispose();
             _isDisposed = true;
+            base.Dispose();
+            GC.SuppressFinalize(this);
         }
     }
 }

libs/host/Configuration/Options.cs

@@ -406,6 +406,11 @@ internal sealed class Options : ICloneable
         [Option("network-connection-limit", Required = false, HelpText = "Maximum number of simultaneously active network connections.")]
         public int NetworkConnectionLimit { get; set; }
 
+        [IntRangeValidation(1, int.MaxValue)]
+        [Option("shutdown-timeout", Required = false, HelpText = "Timeout in seconds to wait for active connections to drain during graceful shutdown. " +
+            "The Windows SCM default pre-kill wait is 5 seconds, so values below 5 are not recommended when running as a Windows service.")]
+        public int ShutdownTimeoutSeconds { get; set; }
+
         [OptionValidation]
         [Option("use-azure-storage", Required = false, HelpText = "Use Azure Page Blobs for storage instead of local storage.")]
         public bool? UseAzureStorage { get; set; }
@@ -938,6 +943,7 @@ endpoint is IPEndPoint listenEp && clusterAnnounceEndpoint[0] is IPEndPoint anno
                 ThreadPoolMinIOCompletionThreads = ThreadPoolMinIOCompletionThreads,
                 ThreadPoolMaxIOCompletionThreads = ThreadPoolMaxIOCompletionThreads,
                 NetworkConnectionLimit = NetworkConnectionLimit,
+                ShutdownTimeoutSeconds = ShutdownTimeoutSeconds,
                 DeviceFactoryCreator = deviceType == DeviceType.AzureStorage ? azureFactoryCreator()
                     : new LocalStorageNamedDeviceFactoryCreator(deviceType: deviceType, logger: logger),
                 CheckpointThrottleFlushDelayMs = CheckpointThrottleFlushDelayMs,

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;
@@ -450,6 +451,212 @@ 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="noSave">If true, skip data persistence (AOF commit and checkpoint) during shutdown</param>
+        /// <param name="token">Cancellation token</param>
+        /// <returns>Task representing the async shutdown operation</returns>
+        public async Task ShutdownAsync(TimeSpan? timeout = null, bool noSave = false, CancellationToken token = default)
+        {
+            var shutdownTimeout = timeout ?? TimeSpan.FromSeconds(30);
+
+            try
+            {
+                // Quiesce existing sessions first: they will reject the next incoming message
+                // and close themselves, so FinalizeDataAsync runs with no concurrent writers.
+                if (servers != null)
+                {
+                    foreach (var server in servers)
+                        server.BeginQuiesce();
+                }
+
+                // Quiesce pub/sub fan-out so no new messages are delivered after this point.
+                subscribeBroker?.BeginQuiesce();
+
+                // Stop accepting new connections.
+                StopListening();
+
+                // Wait for existing connections to complete (cancellable)
+                try
+                {
+                    await WaitForActiveConnectionsAsync(shutdownTimeout, token).ConfigureAwait(false);
+                }
+                catch (OperationCanceledException)
+                {
+                    logger?.LogWarning("Connection draining was cancelled. Proceeding with data finalization...");
+                }
+            }
+            catch (Exception ex)
+            {
+                logger?.LogError(ex, "Error during graceful shutdown");
+            }
+            finally
+            {
+                if (!noSave)
+                {
+                    // Attempt AOF commit or checkpoint as best-effort,
+                    // even if connection draining was cancelled or failed.
+                    // Use a bounded timeout instead of the caller's token to ensure completion.
+                    using var finalizeCts = new CancellationTokenSource(TimeSpan.FromSeconds(15));
+                    try
+                    {
+                        await FinalizeDataAsync(finalizeCts.Token).ConfigureAwait(false);
+                    }
+                    catch (Exception ex)
+                    {
+                        logger?.LogError(ex, "Error during data finalization");
+                    }
+                }
+                else
+                {
+                    logger?.LogInformation("Shutdown with noSave flag - skipping data persistence.");
+                }
+            }
+        }
+
+        /// <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");
+                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>
+        /// Persists data during shutdown using AOF or checkpoint based on configuration.
+        /// </summary>
+        private async Task FinalizeDataAsync(CancellationToken token)
+        {
+            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");
+                }
+
+                return;
+            }
+
+            if (!opts.EnableStorageTier)
+                return;
+
+            logger?.LogDebug("Taking checkpoint for tiered storage...");
+            try
+            {
+                var checkpointSuccess = await Store.TakeCheckpointAsync(background: false, token: token).ConfigureAwait(false);
+                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/host/defaults.conf

@@ -296,6 +296,9 @@
 	/* Maximum number of simultaneously active network connections. */
 	"NetworkConnectionLimit" : -1,
 
+	/* Timeout in seconds to wait for active connections to drain during graceful shutdown. */
+	"ShutdownTimeoutSeconds" : 5,
+
 	/* Use Azure Page Blobs for storage instead of local storage. */
 	"UseAzureStorage" : false,