feat: add lame duck mode with staggered client shutdown

This commit is contained in:
Joseph Doherty
2026-02-22 23:48:06 -05:00
parent 45de110a84
commit 34067f2b9b
2 changed files with 189 additions and 3 deletions

View File

@@ -40,10 +40,7 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
private int _shutdown;
private int _activeClientCount;
// Used by future lame duck mode implementation
#pragma warning disable CS0649 // Field is never assigned to
private int _lameDuck;
#pragma warning restore CS0649
// Used by future ports file implementation
#pragma warning disable CS0169 // Field is never used
@@ -115,6 +112,81 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
_shutdownComplete.TrySetResult();
}
public async Task LameDuckShutdownAsync()
{
if (IsShuttingDown || Interlocked.CompareExchange(ref _lameDuck, 1, 0) != 0)
return;
_logger.LogInformation("Entering lame duck mode, stop accepting new clients");
// Close listener to stop accepting new connections
_listener?.Close();
// Wait for accept loop to exit
await _acceptLoopExited.Task.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
var gracePeriod = _options.LameDuckGracePeriod;
if (gracePeriod < TimeSpan.Zero) gracePeriod = -gracePeriod;
// If no clients, go straight to shutdown
if (_clients.IsEmpty)
{
await ShutdownAsync();
return;
}
// Wait grace period for clients to drain naturally
_logger.LogInformation("Waiting {GracePeriod}ms grace period", gracePeriod.TotalMilliseconds);
try
{
await Task.Delay(gracePeriod, _quitCts.Token);
}
catch (OperationCanceledException) { return; }
if (_clients.IsEmpty)
{
await ShutdownAsync();
return;
}
// Stagger-close remaining clients
var dur = _options.LameDuckDuration - gracePeriod;
if (dur <= TimeSpan.Zero) dur = TimeSpan.FromSeconds(1);
var clients = _clients.Values.ToList();
var numClients = clients.Count;
if (numClients > 0)
{
_logger.LogInformation("Closing {Count} existing clients over {Duration}ms",
numClients, dur.TotalMilliseconds);
var sleepInterval = dur.Ticks / numClients;
if (sleepInterval < TimeSpan.TicksPerMillisecond)
sleepInterval = TimeSpan.TicksPerMillisecond;
if (sleepInterval > TimeSpan.TicksPerSecond)
sleepInterval = TimeSpan.TicksPerSecond;
for (int i = 0; i < clients.Count; i++)
{
clients[i].MarkClosed(ClosedState.ServerShutdown);
await clients[i].FlushAndCloseAsync(minimalFlush: true);
if (i < clients.Count - 1)
{
var jitter = Random.Shared.NextInt64(sleepInterval / 2, sleepInterval);
try
{
await Task.Delay(TimeSpan.FromTicks(jitter), _quitCts.Token);
}
catch (OperationCanceledException) { break; }
}
}
}
await ShutdownAsync();
}
public NatsServer(NatsOptions options, ILoggerFactory loggerFactory)
{
_options = options;