Feature: graceful shutdown for v2 ( #1382 )

hosting/Windows/Garnet.worker/Program.cs

@@ -1,16 +1,37 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+using System;
 using Garnet;
+using Garnet.server;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 
 class Program
 {
+    // Extra host budget beyond connection drain + data finalization (quiesce, stop listening, etc.).
+    private const int HostShutdownMarginSeconds = 5;
+
     static void Main(string[] args)
     {
+        // Pre-parse shutdown-related arguments so we can configure both the host shutdown budget
+        // and the Worker's connection-drain timeout before the generic host is built.
+        var shutdownTimeoutSeconds = ParseIntOption(args, "--shutdown-timeout", "-shutdown-timeout", GarnetServerOptions.DefaultShutdownTimeoutSeconds);
+        var dataFinalizationTimeoutSeconds = ParseIntOption(args, "--data-finalization-timeout", "-data-finalization-timeout", GarnetServerOptions.DefaultDataFinalizationTimeoutSeconds);
+        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: drain + data finalization + margin.
+        builder.Services.Configure<HostOptions>(opts =>
+        {
+            opts.ShutdownTimeout = shutdownTimeout
+                + TimeSpan.FromSeconds(dataFinalizationTimeoutSeconds)
+                + TimeSpan.FromSeconds(HostShutdownMarginSeconds);
+        });
+
+        builder.Services.AddHostedService(_ => new Worker(args, shutdownTimeout));
 
         builder.Services.AddWindowsService(options =>
         {
@@ -20,4 +41,21 @@ static void Main(string[] args)
         var host = builder.Build();
         host.Run();
     }
+
+    /// <summary>
+    /// Scans <paramref name="args"/> for <paramref name="longName"/> (or <paramref name="shortName"/>) and returns
+    /// the parsed positive integer, or <paramref name="defaultSeconds"/> if the argument is absent or invalid.
+    /// </summary>
+    private static int ParseIntOption(string[] args, string longName, string shortName, int defaultSeconds)
+    {
+        for (var i = 0; i < args.Length - 1; i++)
+        {
+            if (args[i] is var name && (name == longName || name == shortName) &&
+                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 fit within the host <see cref="Microsoft.Extensions.Hosting.HostOptions.ShutdownTimeout"/>,
+        /// which also budgets for data finalization and a small margin (see <c>Program.cs</c>).
+        /// </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

@@ -423,6 +423,15 @@ internal sealed class Options : ICloneable
         [Option("network-connection-limit", Required = false, HelpText = "Maximum number of simultaneously active network connections across all listeners, or -1 for unlimited. Equivalent to the Redis maxclients parameter, and settable at runtime with CONFIG SET maxclients.")]
         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; }
+
+        [IntRangeValidation(1, int.MaxValue)]
+        [Option("data-finalization-timeout", Required = false, HelpText = "Timeout in seconds for AOF commit and checkpoint during graceful shutdown data finalization.")]
+        public int DataFinalizationTimeoutSeconds { 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; }
@@ -1015,6 +1024,8 @@ public GarnetServerOptions GetServerOptions(ILogger logger = null)
                 ThreadPoolMinIOCompletionThreads = ThreadPoolMinIOCompletionThreads,
                 ThreadPoolMaxIOCompletionThreads = ThreadPoolMaxIOCompletionThreads,
                 NetworkConnectionLimit = NetworkConnectionLimit,
+                ShutdownTimeoutSeconds = ShutdownTimeoutSeconds,
+                DataFinalizationTimeoutSeconds = DataFinalizationTimeoutSeconds,
                 DeviceFactoryCreator = deviceType == DeviceType.AzureStorage ? azureFactoryCreator()
                     : new LocalStorageNamedDeviceFactoryCreator(
                         deviceType: deviceType,