Prune idle sessions before starting new ones (#701)
This commit is contained in:
parent
70960e1f1b
commit
b067261ca7
1 changed files with 472 additions and 281 deletions
|
|
@ -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();
|
||||
}
|
||||
Loading…
Reference in a new issue