Improve rate limiter and connection disposal logic

Refactored InMemoryRateLimiter's RequestWindow to use primary constructor and simplified field initialization. Enhanced ConnectionTracker disposal logic to handle race conditions by re-adding connections if new activity is detected, and added object disposal checks for lock methods. Introduced MarkAsActiveUnsafe to reset removal status under lock.

Author: Taiizor

src/Zetian/Extensions/RateLimiting/InMemoryRateLimiter.cs

@@ -112,44 +112,35 @@ public void Dispose()
             _windows.Clear();
         }
 
-        private class RequestWindow
+        private class RequestWindow(RateLimitConfiguration configuration)
         {
-            private readonly RateLimitConfiguration _configuration;
             private readonly object _lock = new();
-            private readonly Queue<DateTime> _requests;
-            private DateTime _windowStart;
-            private int _requestCount;
-
-            public RequestWindow(RateLimitConfiguration configuration)
-            {
-                _configuration = configuration;
-                _requests = new Queue<DateTime>();
-                _windowStart = DateTime.UtcNow;
-                _requestCount = 0;
-            }
+            private readonly Queue<DateTime> _requests = new();
+            private DateTime _windowStart = DateTime.UtcNow;
+            private int _requestCount = 0;
 
             public bool IsAllowed()
             {
                 lock (_lock)
                 {
                     CleanupOldRequests();
 
-                    if (_configuration.UseSlidingWindow)
+                    if (configuration.UseSlidingWindow)
                     {
-                        return _requests.Count < _configuration.MaxRequests;
+                        return _requests.Count < configuration.MaxRequests;
                     }
                     else
                     {
                         // Fixed window
                         DateTime now = DateTime.UtcNow;
-                        if (now - _windowStart > _configuration.Window)
+                        if (now - _windowStart > configuration.Window)
                         {
                             // New window
                             _windowStart = now;
                             _requestCount = 0;
                         }
 
-                        return _requestCount < _configuration.MaxRequests;
+                        return _requestCount < configuration.MaxRequests;
                     }
                 }
             }
@@ -160,15 +151,15 @@ public void RecordRequest()
                 {
                     DateTime now = DateTime.UtcNow;
 
-                    if (_configuration.UseSlidingWindow)
+                    if (configuration.UseSlidingWindow)
                     {
                         CleanupOldRequests();
                         _requests.Enqueue(now);
                     }
                     else
                     {
                         // Fixed window
-                        if (now - _windowStart > _configuration.Window)
+                        if (now - _windowStart > configuration.Window)
                         {
                             // New window
                             _windowStart = now;
@@ -188,19 +179,19 @@ public int GetRemaining()
                 {
                     CleanupOldRequests();
 
-                    if (_configuration.UseSlidingWindow)
+                    if (configuration.UseSlidingWindow)
                     {
-                        return Math.Max(0, _configuration.MaxRequests - _requests.Count);
+                        return Math.Max(0, configuration.MaxRequests - _requests.Count);
                     }
                     else
                     {
                         DateTime now = DateTime.UtcNow;
-                        if (now - _windowStart > _configuration.Window)
+                        if (now - _windowStart > configuration.Window)
                         {
-                            return _configuration.MaxRequests;
+                            return configuration.MaxRequests;
                         }
 
-                        return Math.Max(0, _configuration.MaxRequests - _requestCount);
+                        return Math.Max(0, configuration.MaxRequests - _requestCount);
                     }
                 }
             }
@@ -211,27 +202,27 @@ public bool IsExpired()
                 {
                     DateTime now = DateTime.UtcNow;
 
-                    if (_configuration.UseSlidingWindow)
+                    if (configuration.UseSlidingWindow)
                     {
                         CleanupOldRequests();
                         return _requests.Count == 0 &&
-                               (_requests.Count == 0 || now - _windowStart > _configuration.Window * 2);
+                               (_requests.Count == 0 || now - _windowStart > configuration.Window * 2);
                     }
                     else
                     {
-                        return now - _windowStart > _configuration.Window * 2;
+                        return now - _windowStart > configuration.Window * 2;
                     }
                 }
             }
 
             private void CleanupOldRequests()
             {
-                if (!_configuration.UseSlidingWindow)
+                if (!configuration.UseSlidingWindow)
                 {
                     return;
                 }
 
-                DateTime cutoff = DateTime.UtcNow - _configuration.Window;
+                DateTime cutoff = DateTime.UtcNow - configuration.Window;
 
                 while (_requests.Count > 0 && _requests.Peek() < cutoff)
                 {

src/Zetian/Internal/ConnectionTracker.cs

@@ -204,18 +204,32 @@ private void CleanupExpired(object? state)
                 {
                     if (_connections.TryRemove(key, out ConnectionInfo? removed))
                     {
-                        // Only dispose if still marked for removal
-                        removed.Lock();
+                        // Double-check under lock before disposing
                         try
                         {
-                            if (removed.IsMarkedForRemovalUnsafe())
+                            removed.Lock();
+                            try
                             {
-                                removed.Dispose();
+                                // Re-verify no new connections were acquired
+                                if (removed.IsMarkedForRemovalUnsafe() && removed.GetCountUnsafe() == 0)
+                                {
+                                    removed.Dispose();
+                                }
+                                else
+                                {
+                                    // Race condition: new connection acquired, put it back
+                                    removed.MarkAsActiveUnsafe();
+                                    _connections.TryAdd(key, removed);
+                                }
+                            }
+                            finally
+                            {
+                                removed.ReleaseLock();
                             }
                         }
-                        finally
+                        catch (ObjectDisposedException)
                         {
-                            removed.ReleaseLock();
+                            // Already disposed by another thread, ignore
                         }
                     }
                 }
@@ -280,6 +294,14 @@ internal sealed class ConnectionInfo(int maxConnections) : IDisposable
             /// </summary>
             public Task LockAsync(CancellationToken cancellationToken = default)
             {
+#if NET6_0
+                if (_disposed)
+                {
+                    throw new ObjectDisposedException(nameof(ConnectionInfo));
+                }
+#elif NET7_0_OR_GREATER
+                ObjectDisposedException.ThrowIf(_disposed, nameof(ConnectionInfo));
+#endif
                 return _syncLock.WaitAsync(cancellationToken);
             }
 
@@ -288,6 +310,14 @@ public Task LockAsync(CancellationToken cancellationToken = default)
             /// </summary>
             public void Lock()
             {
+#if NET6_0
+                if (_disposed)
+                {
+                    throw new ObjectDisposedException(nameof(ConnectionInfo));
+                }
+#elif NET7_0_OR_GREATER
+                ObjectDisposedException.ThrowIf(_disposed, nameof(ConnectionInfo));
+#endif
                 _syncLock.Wait();
             }
 
@@ -362,6 +392,14 @@ public bool IsMarkedForRemovalUnsafe()
                 return _markedForRemoval;
             }
 
+            /// <summary>
+            /// Marks as active again (must be called under lock)
+            /// </summary>
+            public void MarkAsActiveUnsafe()
+            {
+                _markedForRemoval = false;
+            }
+
             /// <summary>
             /// Releases all resources used by the current instance.
             /// </summary>