Implement graceful shutdown for Garnet server ( #1382 ) #1551

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

yuseok-kim-edushare wants to merge 59 commits into microsoft:main from yuseok-kim-edushare:yuseok-kim/graceful_shutdown

hosting/Windows/Garnet.worker/Worker.cs

-Original file line number
+Diff line change
@@ Expand Up @@
             /// <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: TimeSpan.FromSeconds(5), 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()
@@ Expand All / @@ -55,6 +73,8 @@ public override void Dispose() @@
                 }
                 server?.Dispose();
                 _isDisposed = true;
+                base.Dispose();
+                GC.SuppressFinalize(this);
             }
         }
     }

libs/host/GarnetServer.cs

-Original file line number
+Diff line change
@@ Expand Up / @@ -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;
@@ Expand Down Expand Up / @@ -450,6 +451,201 @@ 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
+                {
+                    // Stop accepting new connections first
+                    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 = Store.TakeCheckpoint(background: false, token: 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>
@@ Expand Down @@

libs/server/Servers/GarnetServerBase.cs

-Original file line number
+Diff line change
@@ Expand Up @@
             /// <inheritdoc />
             public abstract void Start();
+            /// <inheritdoc />
+            public virtual void StopListening()
+            {
+                // Base implementation does nothing; derived classes should override
+            }
             /// <inheritdoc />
             public virtual void Dispose()
             {
@@ Expand Down @@

libs/server/Servers/GarnetServerTcp.cs

-Original file line number
+Diff line change
@@ Expand Up / @@ -121,6 +121,22 @@ public override void Start() @@
                     AcceptEventArg_Completed(null, acceptEventArg);
             }
+            /// <inheritdoc />
+            public override void StopListening()
+            {
+                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
@@ Expand Down @@

libs/server/Servers/IGarnetServer.cs

-Original file line number
+Diff line change
@@ Expand Up / @@ -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

-Original file line number
+Diff line change
@@ Expand Up @@
                 }
             }
+            /// <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.
             ///
@@ Expand Down @@

main/GarnetServer/Program.cs

-Original file line number
+Diff line change
@@ Expand Up / @@ -22,7 +22,18 @@ static void Main(string[] args) @@
                     // Start the server
                     server.Start();
-                    Thread.Sleep(Timeout.Infinite);
+                    using 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)
                 {
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement graceful shutdown for Garnet server ( #1382 ) #1551

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Implement graceful shutdown for Garnet server ( #1382 ) #1551

Are you sure you want to change the base?

Uh oh!

Implement graceful shutdown for Garnet server ( #1382 ) #1551

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!