Prune idle sessions before starting new ones (#701)

This commit is contained in:
Stephen Halter 2025-08-18 11:04:27 -07:00 • committed by GitHub
commit b067261ca7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,18 +1,16 @@
using System.Runtime.InteropServices;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using ModelContextProtocol.Server;
namespace ModelContextProtocol.AspNetCore;
internal sealed partial class IdleTrackingBackgroundService(
StreamableHttpHandler handler,
StatefulSessionManager sessions,
IOptions<HttpServerTransportOptions> options,
IHostApplicationLifetime appLifetime,
ILogger<IdleTrackingBackgroundService> logger) : BackgroundService
{
// The compiler will complain about the parameter being unused otherwise despite the source generator.
// Workaround for https://github.com/dotnet/runtime/issues/91121. This is fixed in .NET 9 and later.
private readonly ILogger _logger = logger;
protected override async Task ExecuteAsync(CancellationToken stoppingToken)
@ -30,65 +28,9 @@ internal sealed partial class IdleTrackingBackgroundService(
var timeProvider = options.Value.TimeProvider;
using var timer = new PeriodicTimer(TimeSpan.FromSeconds(5), timeProvider);
var idleTimeoutTicks = options.Value.IdleTimeout.Ticks;
var maxIdleSessionCount = options.Value.MaxIdleSessionCount;
// Create two lists that will be reused between runs.
// This assumes that the number of idle sessions is not breached frequently.
// If the idle sessions often breach the maximum, a priority queue could be considered.
var idleSessionsTimestamps = new List<long>();
var idleSessionSessionIds = new List<string>();
while (!stoppingToken.IsCancellationRequested && await timer.WaitForNextTickAsync(stoppingToken))
{
var idleActivityCutoff = idleTimeoutTicks switch
{
< 0 => long.MinValue,
var ticks => timeProvider.GetTimestamp() - ticks,
};
foreach (var (_, session) in handler.Sessions)
{
if (session.IsActive || session.SessionClosed.IsCancellationRequested)
{
// There's a request currently active or the session is already being closed.
continue;
}
if (session.LastActivityTicks < idleActivityCutoff)
{
RemoveAndCloseSession(session.Id);
continue;
}
// Add the timestamp and the session
idleSessionsTimestamps.Add(session.LastActivityTicks);
idleSessionSessionIds.Add(session.Id);
// Emit critical log at most once every 5 seconds the idle count it exceeded,
// since the IdleTimeout will no longer be respected.
if (idleSessionsTimestamps.Count == maxIdleSessionCount + 1)
{
LogMaxSessionIdleCountExceeded(maxIdleSessionCount);
}
}
if (idleSessionsTimestamps.Count > maxIdleSessionCount)
{
var timestamps = CollectionsMarshal.AsSpan(idleSessionsTimestamps);
// Sort only if the maximum is breached and sort solely by the timestamp. Sort both collections.
timestamps.Sort(CollectionsMarshal.AsSpan(idleSessionSessionIds));
var sessionsToPrune = CollectionsMarshal.AsSpan(idleSessionSessionIds)[..^maxIdleSessionCount];
foreach (var id in sessionsToPrune)
{
RemoveAndCloseSession(id);
}
}
idleSessionsTimestamps.Clear();
idleSessionSessionIds.Clear();
await sessions.PruneIdleSessionsAsync(stoppingToken);
}
}
catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
@ -98,17 +40,7 @@ internal sealed partial class IdleTrackingBackgroundService(
{
try
{
List<Task> disposeSessionTasks = [];
foreach (var (sessionKey, _) in handler.Sessions)
{
if (handler.Sessions.TryRemove(sessionKey, out var session))
{
disposeSessionTasks.Add(DisposeSessionAsync(session));
}
}
await Task.WhenAll(disposeSessionTasks);
await sessions.DisposeAllSessionsAsync();
}
finally
{
@ -123,39 +55,6 @@ internal sealed partial class IdleTrackingBackgroundService(
}
}
private void RemoveAndCloseSession(string sessionId)
{
if (!handler.Sessions.TryRemove(sessionId, out var session))
{
return;
}
LogSessionIdle(session.Id);
// Don't slow down the idle tracking loop. DisposeSessionAsync logs. We only await during graceful shutdown.
_ = DisposeSessionAsync(session);
}
private async Task DisposeSessionAsync(HttpMcpSession<StreamableHttpServerTransport> session)
{
try
{
await session.DisposeAsync();
}
catch (Exception ex)
{
LogSessionDisposeError(session.Id, ex);
}
}
[LoggerMessage(Level = LogLevel.Information, Message = "Closing idle session {sessionId}.")]
private partial void LogSessionIdle(string sessionId);
[LoggerMessage(Level = LogLevel.Error, Message = "Error disposing session {sessionId}.")]
private partial void LogSessionDisposeError(string sessionId, Exception ex);
[LoggerMessage(Level = LogLevel.Critical, Message = "Exceeded maximum of {maxIdleSessionCount} idle sessions. Now closing sessions active more recently than configured IdleTimeout.")]
private partial void LogMaxSessionIdleCountExceeded(int maxIdleSessionCount);
[LoggerMessage(Level = LogLevel.Critical, Message = "The IdleTrackingBackgroundService has stopped unexpectedly.")]
private partial void IdleTrackingBackgroundServiceStoppedUnexpectedly();
}