Address PR feedback on StatefulSessionManager

The main thing is disposing pruned sessions outside the lock. This
improved session creation performance when at max idle sessions
by about 10% at a 10,000 idle session max and 100% at a 1,000 max

Author: halter73

src/ModelContextProtocol.AspNetCore/IdleTrackingBackgroundService.cs

@@ -48,8 +48,8 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
                 {
                     // Something went terribly wrong. A very unexpected exception must be bubbling up, but let's ensure we also stop the application,
                     // so that it hopefully gets looked at and restarted. This shouldn't really be reachable.
-                    IdleTrackingBackgroundServiceStoppedUnexpectedly();
                     appLifetime.StopApplication();
+                    IdleTrackingBackgroundServiceStoppedUnexpectedly();
                 }
             }
         }

src/ModelContextProtocol.AspNetCore/StatefulSessionManager.cs

@@ -1,25 +1,28 @@
-using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.Options;
-using System.Collections.Concurrent;
+using System.Collections.Concurrent;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.InteropServices;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
 
 namespace ModelContextProtocol.AspNetCore;
 
 internal sealed partial class StatefulSessionManager(
     IOptions<HttpServerTransportOptions> httpServerTransportOptions,
-    ILogger<StatefulSessionManager> logger) : ConcurrentDictionary<string, StreamableHttpSession>(StringComparer.Ordinal)
+    ILogger<StatefulSessionManager> logger)
 {
     // Workaround for https://github.com/dotnet/runtime/issues/91121. This is fixed in .NET 9 and later.
     private readonly ILogger _logger = logger;
 
+    private readonly ConcurrentDictionary<string, StreamableHttpSession> _sessions = new(StringComparer.Ordinal);
+
     private readonly TimeProvider _timeProvider = httpServerTransportOptions.Value.TimeProvider;
     private readonly TimeSpan _idleTimeout = httpServerTransportOptions.Value.IdleTimeout;
     private readonly long _idleTimeoutTicks = httpServerTransportOptions.Value.IdleTimeout.Ticks;
     private readonly int _maxIdleSessionCount = httpServerTransportOptions.Value.MaxIdleSessionCount;
 
-    private readonly SemaphoreSlim _idlePruningLock = new(1, 1);
-    private readonly List<long> _idleTimeStamps = [];
+    private readonly object _idlePruningLock = new();
+    private readonly List<long> _idleTimestamps = [];
     private readonly List<string> _idleSessionIds = [];
     private int _nextIndexToPrune;
 
@@ -30,62 +33,72 @@ internal sealed partial class StatefulSessionManager(
     public void IncrementIdleSessionCount() => Interlocked.Increment(ref _currentIdleSessionCount);
     public void DecrementIdleSessionCount() => Interlocked.Decrement(ref _currentIdleSessionCount);
 
+    public bool TryGetValue(string key, [NotNullWhen(true)] out StreamableHttpSession? value) => _sessions.TryGetValue(key, out value);
+    public bool TryRemove(string key, [NotNullWhen(true)] out StreamableHttpSession? value) => _sessions.TryRemove(key, out value);
+
     public async ValueTask StartNewSessionAsync(StreamableHttpSession newSession, CancellationToken cancellationToken)
     {
-        if (TryAddSessionImmediately(newSession))
+        while (!TryAddSessionImmediately(newSession))
         {
-            return;
-        }
+            StreamableHttpSession? sessionToPrune = null;
 
-        await _idlePruningLock.WaitAsync(cancellationToken);
-        try
-        {
-            while (!TryAddSessionImmediately(newSession))
+            lock (_idlePruningLock)
             {
                 EnsureIdleSessionsSortedUnsynchronized();
 
                 while (_nextIndexToPrune < _idleSessionIds.Count)
                 {
                     var pruneId = _idleSessionIds[_nextIndexToPrune++];
-                    if (TryGetValue(pruneId, out var sessionToPrune))
+                    if (TryGetValue(pruneId, out sessionToPrune))
                     {
-                        if (!sessionToPrune.IsActive && TryRemove(pruneId, out _))
+                        if (!sessionToPrune.IsActive && _sessions.TryRemove(pruneId, out sessionToPrune))
                         {
                             LogIdleSessionLimit(pruneId, _maxIdleSessionCount);
-
-                            // We're intentionally waiting for the idle session to be disposed before releasing the _idlePruningLock to
-                            // ensure new sessions are not created faster than they're removed when we're at or above the maximum idle session count.
-                            await DisposeSessionAsync(sessionToPrune);
-
-                            // Take one last chance to check if the initialize request was aborted before we incur the cost of managing a new session.
-                            cancellationToken.ThrowIfCancellationRequested();
-                            AddSession(newSession);
-                            return;
+                            break;
                         }
                     }
                 }
 
-                // If we couldn't find any active idle sessions to dispose, start another full prune to repopulate _idleSessionIds.
-                PruneIdleSessionsUnsynchronized();
-
-                if (_idleSessionIds.Count == 0 && Interlocked.Read(ref _currentIdleSessionCount) >= _maxIdleSessionCount)
+                if (sessionToPrune is null)
                 {
-                    // This indicates all idle sessions are in the process of being disposed which should not happen during normal operation.
-                    // Since there are no idle sessions to prune right now, log a critical error and create the new session anyway.
-                    LogTooManyIdleSessionsClosingConcurrently(newSession.Id, _maxIdleSessionCount, Interlocked.Read(ref _currentIdleSessionCount));
-                    AddSession(newSession);
-                    return;
+                    // If we couldn't find any active idle sessions to dispose, start another full prune to repopulate _idleSessionIds.
+                    PruneIdleSessionsUnsynchronized();
+
+                    if (_idleSessionIds.Count > 0)
+                    {
+                        continue;
+                    }
+                    else
+                    {
+                        if (Volatile.Read(ref _currentIdleSessionCount) >= _maxIdleSessionCount)
+                        {
+                            // This indicates all idle sessions are in the process of being disposed which should not happen during normal operation.
+                            // Since there are no idle sessions to prune right now, log a critical error and create the new session anyway.
+                            LogTooManyIdleSessionsClosingConcurrently(newSession.Id, _maxIdleSessionCount, Volatile.Read(ref _currentIdleSessionCount));
+                        }
+
+                        AddSession(newSession);
+                        return;
+                    }
                 }
             }
-        }
-        catch
-        {
-            await newSession.DisposeAsync();
-            throw;
-        }
-        finally
-        {
-            _idlePruningLock.Release();
+
+            try
+            {
+                // Since we're at or above the maximum idle session count, we're intentionally waiting for the idle session to be disposed
+                // before adding a new session to the dictionary to ensure sessions not created faster than they're removed.
+                await DisposeSessionAsync(sessionToPrune);
+
+                // Take one last chance to check if the initialize request was aborted before we incur the cost of managing a new session.
+                cancellationToken.ThrowIfCancellationRequested();
+                AddSession(newSession);
+                return;
+            }
+            catch
+            {
+                await newSession.DisposeAsync();
+                throw;
+            }
         }
     }
 
@@ -95,15 +108,10 @@ public async ValueTask StartNewSessionAsync(StreamableHttpSession newSession, Ca
     /// </summary>
     public async Task PruneIdleSessionsAsync(CancellationToken cancellationToken)
     {
-        await _idlePruningLock.WaitAsync(cancellationToken);
-        try
+        lock (_idlePruningLock)
         {
             PruneIdleSessionsUnsynchronized();
         }
-        finally
-        {
-            _idlePruningLock.Release();
-        }
     }
 
     private void PruneIdleSessionsUnsynchronized()
@@ -116,11 +124,11 @@ private void PruneIdleSessionsUnsynchronized()
 
         // We clear the lists at the start of pruning rather than the end so we can use them between runs
         // to find the most idle sessions to remove one-at-a-time if necessary to make room for new sessions.
-        _idleTimeStamps.Clear();
+        _idleTimestamps.Clear();
         _idleSessionIds.Clear();
         _nextIndexToPrune = -1;
 
-        foreach (var (_, session) in this)
+        foreach (var (_, session) in _sessions)
         {
             if (session.IsActive || session.SessionClosed.IsCancellationRequested)
             {
@@ -136,11 +144,11 @@ private void PruneIdleSessionsUnsynchronized()
             }
 
             // Add the timestamp and the session
-            _idleTimeStamps.Add(session.LastActivityTicks);
+            _idleTimestamps.Add(session.LastActivityTicks);
             _idleSessionIds.Add(session.Id);
         }
 
-        if (_idleTimeStamps.Count > _maxIdleSessionCount)
+        if (_idleTimestamps.Count > _maxIdleSessionCount)
         {
             // Sort only if the maximum is breached and sort solely by the timestamp.
             EnsureIdleSessionsSortedUnsynchronized();
@@ -163,7 +171,7 @@ private void EnsureIdleSessionsSortedUnsynchronized()
             return;
         }
 
-        var timestamps = CollectionsMarshal.AsSpan(_idleTimeStamps);
+        var timestamps = CollectionsMarshal.AsSpan(_idleTimestamps);
         timestamps.Sort(CollectionsMarshal.AsSpan(_idleSessionIds));
         _nextIndexToPrune = 0;
     }
@@ -175,9 +183,9 @@ public async Task DisposeAllSessionsAsync()
     {
         List<Task> disposeSessionTasks = [];
 
-        foreach (var (sessionKey, _) in this)
+        foreach (var (sessionKey, _) in _sessions)
         {
-            if (TryRemove(sessionKey, out var session))
+            if (_sessions.TryRemove(sessionKey, out var session))
             {
                 disposeSessionTasks.Add(DisposeSessionAsync(session));
             }
@@ -188,7 +196,7 @@ public async Task DisposeAllSessionsAsync()
 
     private bool TryAddSessionImmediately(StreamableHttpSession session)
     {
-        if (Interlocked.Read(ref _currentIdleSessionCount) < _maxIdleSessionCount)
+        if (Volatile.Read(ref _currentIdleSessionCount) < _maxIdleSessionCount)
         {
             AddSession(session);
             return true;
@@ -199,15 +207,15 @@ private bool TryAddSessionImmediately(StreamableHttpSession session)
 
     private void AddSession(StreamableHttpSession session)
     {
-        if (!TryAdd(session.Id, session))
+        if (!_sessions.TryAdd(session.Id, session))
         {
             throw new UnreachableException($"Unreachable given good entropy! Session with ID '{session.Id}' has already been created.");
         }
     }
 
     private void RemoveAndCloseSession(string sessionId)
     {
-        if (!TryRemove(sessionId, out var session))
+        if (!_sessions.TryRemove(sessionId, out var session))
         {
             return;
         }

src/ModelContextProtocol.AspNetCore/StreamableHttpSession.cs

@@ -132,7 +132,7 @@ private async ValueTask DisposeServerThenTransportAsync()
         }
         finally
         {
-            await Transport.DisposeAsync();
+            await transport.DisposeAsync();
         }
     }