IdleTracking uses lists instead of SortedSet (#629)

This commit is contained in:
Szymon Kulec 2025-07-18 20:36:21 +02:00 • committed by GitHub
commit 929ae58906
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,4 +1,5 @@
using Microsoft.Extensions.Hosting; using System.Runtime.InteropServices;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options; using Microsoft.Extensions.Options;
using ModelContextProtocol.Server; using ModelContextProtocol.Server;
@ -21,6 +22,7 @@ internal sealed partial class IdleTrackingBackgroundService(
{ {
ArgumentOutOfRangeException.ThrowIfLessThan(options.Value.IdleTimeout, TimeSpan.Zero); ArgumentOutOfRangeException.ThrowIfLessThan(options.Value.IdleTimeout, TimeSpan.Zero);
} }
ArgumentOutOfRangeException.ThrowIfLessThan(options.Value.MaxIdleSessionCount, 0); ArgumentOutOfRangeException.ThrowIfLessThan(options.Value.MaxIdleSessionCount, 0);
try try
@ -31,8 +33,11 @@ internal sealed partial class IdleTrackingBackgroundService(
var idleTimeoutTicks = options.Value.IdleTimeout.Ticks; var idleTimeoutTicks = options.Value.IdleTimeout.Ticks;
var maxIdleSessionCount = options.Value.MaxIdleSessionCount; var maxIdleSessionCount = options.Value.MaxIdleSessionCount;
// The default ValueTuple Comparer will check the first item then the second which preserves both order and uniqueness. // Create two lists that will be reused between runs.
var idleSessions = new SortedSet<(long Timestamp, string SessionId)>(); // 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)) while (!stoppingToken.IsCancellationRequested && await timer.WaitForNextTickAsync(stoppingToken))
{ {
@ -56,26 +61,34 @@ internal sealed partial class IdleTrackingBackgroundService(
continue; 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, // Emit critical log at most once every 5 seconds the idle count it exceeded,
// since the IdleTimeout will no longer be respected. // since the IdleTimeout will no longer be respected.
if (idleSessions.Count == maxIdleSessionCount + 1) if (idleSessionsTimestamps.Count == maxIdleSessionCount + 1)
{ {
LogMaxSessionIdleCountExceeded(maxIdleSessionCount); LogMaxSessionIdleCountExceeded(maxIdleSessionCount);
} }
} }
if (idleSessions.Count > maxIdleSessionCount) if (idleSessionsTimestamps.Count > maxIdleSessionCount)
{ {
var sessionsToPrune = idleSessions.ToArray()[..^maxIdleSessionCount]; var timestamps = CollectionsMarshal.AsSpan(idleSessionsTimestamps);
foreach (var (_, id) in sessionsToPrune)
// 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); RemoveAndCloseSession(id);
} }
} }
idleSessions.Clear(); idleSessionsTimestamps.Clear();
idleSessionSessionIds.Clear();
} }
} }
catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested) catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
@ -145,4 +158,4 @@ internal sealed partial class IdleTrackingBackgroundService(
[LoggerMessage(Level = LogLevel.Critical, Message = "The IdleTrackingBackgroundService has stopped unexpectedly.")] [LoggerMessage(Level = LogLevel.Critical, Message = "The IdleTrackingBackgroundService has stopped unexpectedly.")]
private partial void IdleTrackingBackgroundServiceStoppedUnexpectedly(); private partial void IdleTrackingBackgroundServiceStoppedUnexpectedly();
} }