diff --git a/src/ModelContextProtocol.AspNetCore/IdleTrackingBackgroundService.cs b/src/ModelContextProtocol.AspNetCore/IdleTrackingBackgroundService.cs index 26ffd44bb..c4a5f11ee 100644 --- a/src/ModelContextProtocol.AspNetCore/IdleTrackingBackgroundService.cs +++ b/src/ModelContextProtocol.AspNetCore/IdleTrackingBackgroundService.cs @@ -1,4 +1,5 @@ -using Microsoft.Extensions.Hosting; +using System.Runtime.InteropServices; +using Microsoft.Extensions.Hosting; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using ModelContextProtocol.Server; @@ -21,6 +22,7 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken) { ArgumentOutOfRangeException.ThrowIfLessThan(options.Value.IdleTimeout, TimeSpan.Zero); } + ArgumentOutOfRangeException.ThrowIfLessThan(options.Value.MaxIdleSessionCount, 0); try @@ -31,8 +33,11 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken) var idleTimeoutTicks = options.Value.IdleTimeout.Ticks; var maxIdleSessionCount = options.Value.MaxIdleSessionCount; - // The default ValueTuple Comparer will check the first item then the second which preserves both order and uniqueness. - var idleSessions = new SortedSet<(long Timestamp, string SessionId)>(); + // 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(); + var idleSessionSessionIds = new List(); while (!stoppingToken.IsCancellationRequested && await timer.WaitForNextTickAsync(stoppingToken)) { @@ -56,26 +61,34 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken) continue; } - idleSessions.Add((session.LastActivityTicks, session.Id)); + // 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 (idleSessions.Count == maxIdleSessionCount + 1) + if (idleSessionsTimestamps.Count == maxIdleSessionCount + 1) { LogMaxSessionIdleCountExceeded(maxIdleSessionCount); } } - if (idleSessions.Count > maxIdleSessionCount) + if (idleSessionsTimestamps.Count > maxIdleSessionCount) { - var sessionsToPrune = idleSessions.ToArray()[..^maxIdleSessionCount]; - foreach (var (_, id) in sessionsToPrune) + 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); } } - idleSessions.Clear(); + idleSessionsTimestamps.Clear(); + idleSessionSessionIds.Clear(); } } catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested) @@ -145,4 +158,4 @@ private async Task DisposeSessionAsync(HttpMcpSession