Fix NoopLimiter disposal in DefaultPartitionedRateLimiter Heartbeat

src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs

@@ -16,14 +16,14 @@ internal sealed class DefaultPartitionedRateLimiter<TResource, TKey> : Partition
         private static readonly TimeSpan s_idleTimeLimit = TimeSpan.FromSeconds(10);
 
         // TODO: Look at ConcurrentDictionary to try and avoid a global lock
-        private readonly Dictionary<TKey, Lazy<RateLimiter>> _limiters;
+        private readonly Dictionary<TKey, Lazy<LimiterEntry>> _limiters;
         private bool _disposed;
         private readonly TaskCompletionSource<object?> _disposeComplete = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
         // Used by the Timer to call TryRelenish on ReplenishingRateLimiters
         // We use a separate list to avoid running TryReplenish (which might be user code) inside our lock
         // And we cache the list to amortize the allocation cost to as close to 0 as we can get
-        private readonly List<KeyValuePair<TKey, Lazy<RateLimiter>>> _cachedLimiters = new();
+        private readonly List<KeyValuePair<TKey, Lazy<LimiterEntry>>> _cachedLimiters = new();
         private bool _cacheInvalid;
         private readonly List<RateLimiter> _limitersToDispose = new();
         private readonly TimerAwaitable _timer;
@@ -42,7 +42,7 @@ public DefaultPartitionedRateLimiter(Func<TResource, RateLimitPartition<TKey>> p
         private DefaultPartitionedRateLimiter(Func<TResource, RateLimitPartition<TKey>> partitioner,
             IEqualityComparer<TKey>? equalityComparer, TimeSpan timerInterval)
         {
-            _limiters = new Dictionary<TKey, Lazy<RateLimiter>>(equalityComparer);
+            _limiters = new Dictionary<TKey, Lazy<LimiterEntry>>(equalityComparer);
             _partitioner = partitioner;
 
             _timer = new TimerAwaitable(timerInterval, timerInterval);
@@ -87,20 +87,36 @@ protected override ValueTask<RateLimitLease> AcquireAsyncCore(TResource resource
         private RateLimiter GetRateLimiter(TResource resource)
         {
             RateLimitPartition<TKey> partition = _partitioner(resource);
-            Lazy<RateLimiter>? limiter;
+            Lazy<LimiterEntry>? entry;
             lock (Lock)
             {
                 ThrowIfDisposed();
-                if (!_limiters.TryGetValue(partition.PartitionKey, out limiter))
+                if (!_limiters.TryGetValue(partition.PartitionKey, out entry))
                 {
-                    // Using Lazy avoids calling user code (partition.Factory) inside the lock
-                    limiter = new Lazy<RateLimiter>(() => partition.Factory(partition.PartitionKey));
-                    _limiters.Add(partition.PartitionKey, limiter);
+                    // Using Lazy avoids calling user code (partition.Factory) inside the lock.
+                    // The LimiterEntry constructor initializes LastAccessTimestamp to Stopwatch.GetTimestamp()
+                    // when the factory runs (on the first access of entry.Value below). Until then,
+                    // Lazy.IsValueCreated is false and Heartbeat skips this entry, so there is no window
+                    // in which the entry could be observed without a timestamp.
+                    entry = new Lazy<LimiterEntry>(() => new LimiterEntry(partition.Factory(partition.PartitionKey)));
+                    _limiters.Add(partition.PartitionKey, entry);
                     // Cache is invalid now
                     _cacheInvalid = true;
                 }
+                else if (entry.IsValueCreated)
+                {
+                    LimiterEntry limiterEntry = entry.Value;
+
+                    if (limiterEntry.Limiter is NoopLimiter)
+                    {
+                        // Refresh the timestamp under the lock so Heartbeat won't evict a limiter that's actively being used.
+                        // Use Volatile.Write so the write is atomic on 32-bit platforms where the outside-lock read in Heartbeat may otherwise tear.
+                        Volatile.Write(ref limiterEntry.LastAccessTimestamp, Stopwatch.GetTimestamp());
+                    }
+                }
             }
-            return limiter.Value;
+
+            return entry.Value.Limiter;
         }
 
         protected override void Dispose(bool disposing)
@@ -125,11 +141,11 @@ protected override void Dispose(bool disposing)
 
             // Safe to access _limiters outside the lock
             // The timer is no longer running and _disposed is set so anyone trying to access fields will be checking that first
-            foreach (KeyValuePair<TKey, Lazy<RateLimiter>> limiter in _limiters)
+            foreach (KeyValuePair<TKey, Lazy<LimiterEntry>> limiter in _limiters)
             {
                 try
                 {
-                    limiter.Value.Value.Dispose();
+                    limiter.Value.Value.Limiter.Dispose();
                 }
                 catch (Exception ex)
                 {
@@ -160,11 +176,11 @@ protected override async ValueTask DisposeAsyncCore()
             }
 
             List<Exception>? exceptions = null;
-            foreach (KeyValuePair<TKey, Lazy<RateLimiter>> limiter in _limiters)
+            foreach (KeyValuePair<TKey, Lazy<LimiterEntry>> limiter in _limiters)
             {
                 try
                 {
-                    await limiter.Value.Value.DisposeAsync().ConfigureAwait(false);
+                    await limiter.Value.Value.Limiter.DisposeAsync().ConfigureAwait(false);
                 }
                 catch (Exception ex)
                 {
@@ -225,18 +241,19 @@ private async Task Heartbeat()
             List<Exception>? aggregateExceptions = null;
 
             // cachedLimiters is safe to use outside the lock because it is only updated by the Timer
-            foreach (KeyValuePair<TKey, Lazy<RateLimiter>> rateLimiter in _cachedLimiters)
+            foreach (KeyValuePair<TKey, Lazy<LimiterEntry>> rateLimiter in _cachedLimiters)
             {
                 if (!rateLimiter.Value.IsValueCreated)
                 {
                     continue;
                 }
-                if (rateLimiter.Value.Value.IdleDuration is TimeSpan idleDuration && idleDuration > s_idleTimeLimit)
+                LimiterEntry limiterEntry = rateLimiter.Value.Value;
+                if (GetIdleDuration(limiterEntry) is TimeSpan idleDuration && idleDuration > s_idleTimeLimit)
                 {
                     lock (Lock)
                     {
                         // Check time again under lock to make sure no one calls Acquire or WaitAsync after checking the time and removing the limiter
-                        idleDuration = rateLimiter.Value.Value.IdleDuration ?? TimeSpan.Zero;
+                        idleDuration = GetIdleDuration(limiterEntry) ?? TimeSpan.Zero;
                         if (idleDuration > s_idleTimeLimit)
                         {
                             // Remove limiter from the lookup table and mark cache as invalid
@@ -246,12 +263,12 @@ private async Task Heartbeat()
                             _limiters.Remove(rateLimiter.Key);
 
                             // We don't want to dispose inside the lock so we need to defer it
-                            _limitersToDispose.Add(rateLimiter.Value.Value);
+                            _limitersToDispose.Add(limiterEntry.Limiter);
                         }
                     }
                 }
                 // We know the limiter can be replenished so let's attempt to replenish tokens
-                else if (rateLimiter.Value.Value is ReplenishingRateLimiter replenishingRateLimiter)
+                else if (limiterEntry.Limiter is ReplenishingRateLimiter replenishingRateLimiter)
                 {
                     try
                     {
@@ -284,5 +301,31 @@ private async Task Heartbeat()
                 throw new AggregateException(aggregateExceptions);
             }
         }
+
+        private static TimeSpan? GetIdleDuration(LimiterEntry limiterEntry)
+        {
+            // NoopLimiter always reports IdleDuration == null, but it is also always safe to evict.
+            // Fall back to our internally tracked last-access timestamp only for that known case.
+            // Use Volatile.Read so the value is read atomically on 32-bit platforms where 64-bit
+            // reads are not guaranteed to be atomic.
+            return limiterEntry.Limiter.IdleDuration ?? (limiterEntry.Limiter is NoopLimiter
+                ? RateLimiterHelper.GetElapsedTime(Volatile.Read(ref limiterEntry.LastAccessTimestamp))
+                : null);
+        }
+
+        // Wraps a RateLimiter with a timestamp of when it was last accessed by the partitioned limiter.
+        // The timestamp is used by the Heartbeat to evict NoopLimiter partitions, whose IdleDuration
+        // is always null and would otherwise never be cleaned up.
+        private sealed class LimiterEntry
+        {
+            public LimiterEntry(RateLimiter limiter)
+            {
+                Limiter = limiter;
+                LastAccessTimestamp = Stopwatch.GetTimestamp();
+            }
+
+            public RateLimiter Limiter { get; }
+            internal long LastAccessTimestamp;
+        }
     }
 }

src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/RateLimiterHelper.cs

@@ -17,6 +17,11 @@ internal static class RateLimiterHelper
             return Stopwatch.GetElapsedTime(startTimestamp.Value);
         }
 
+        public static TimeSpan GetElapsedTime(long startTimestamp)
+        {
+            return Stopwatch.GetElapsedTime(startTimestamp);
+        }
+
         public static TimeSpan GetElapsedTime(long startTimestamp, long endTimestamp)
         {
             return Stopwatch.GetElapsedTime(startTimestamp, endTimestamp);

src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs

@@ -10,6 +10,9 @@ namespace System.Threading.RateLimiting.Tests
 {
     public class PartitionedRateLimiterTests
     {
+        private const string DefaultPartitionedRateLimiterTypeName = "System.Threading.RateLimiting.DefaultPartitionedRateLimiter`2, System.Threading.RateLimiting";
+        private const string DefaultPartitionedRateLimiterEntryTypeName = "System.Threading.RateLimiting.DefaultPartitionedRateLimiter`2+LimiterEntry, System.Threading.RateLimiting";
+
         [Fact]
         public void ThrowsWhenAcquiringLessThanZero()
         {
@@ -529,6 +532,113 @@ public async Task IdleLimiterIsCleanedUp()
             Assert.Equal(2, factoryCallCount);
         }
 
+        [Fact]
+        public async Task NoopLimiterPartitionIsCleanedUp()
+        {
+            using var limiter = Utils.CreatePartitionedLimiterWithoutTimer<string, int>(resource =>
+            {
+                return RateLimitPartition.GetNoLimiter(1);
+            });
+
+            var lease = limiter.AttemptAcquire("");
+            Assert.True(lease.IsAcquired);
+            Assert.NotNull(GetLazyLimiterEntry<string, int>(limiter, key: 1));
+
+            // Backdate the internally-tracked last-access timestamp so the heartbeat treats the
+            // partition as idle even though NoopLimiter reports IdleDuration == null.
+            BackdateLastAccessTimestamp<string, int>(limiter, key: 1);
+
+            await Utils.RunTimerFunc<string, int>(limiter);
+            Assert.Null(GetLazyLimiterEntry<string, int>(limiter, key: 1));
+
+            lease = limiter.AttemptAcquire("");
+            Assert.True(lease.IsAcquired);
+            Assert.NotNull(GetLazyLimiterEntry<string, int>(limiter, key: 1));
+        }
+
+        [Fact]
+        public async Task LimiterWithNullIdleDurationIsNotCleanedUp()
+        {
+            var factoryCallCount = 0;
+            using var limiter = Utils.CreatePartitionedLimiterWithoutTimer<string, int>(resource =>
+            {
+                return RateLimitPartition.Get(1, _ =>
+                {
+                    factoryCallCount++;
+                    return new CustomizableLimiter
+                    {
+                        IdleDurationImpl = () => null
+                    };
+                });
+            });
+
+            var lease = limiter.AttemptAcquire("");
+            Assert.True(lease.IsAcquired);
+            Assert.Equal(1, factoryCallCount);
+
+            BackdateLastAccessTimestamp<string, int>(limiter, key: 1);
+
+            await Utils.RunTimerFunc<string, int>(limiter);
+
+            lease = limiter.AttemptAcquire("");
+            Assert.True(lease.IsAcquired);
+            Assert.Equal(1, factoryCallCount);
+        }
+
+        // Reaches into the internal _limiters dictionary of DefaultPartitionedRateLimiter and
+        // backdates the LastAccessTimestamp on the LimiterEntry for the given key so the test
+        // doesn't have to wait the real idle window before the heartbeat evicts the limiter.
+        private static void BackdateLastAccessTimestamp<TResource, TKey>(PartitionedRateLimiter<TResource> limiter, TKey key)
+        {
+            // Use Type.GetType with literal strings so trimming can see what types we're reflecting on.
+            var limiterEntryTypeDef = Type.GetType(DefaultPartitionedRateLimiterEntryTypeName);
+            Assert.NotNull(limiterEntryTypeDef);
+            var limiterEntryType = limiterEntryTypeDef.MakeGenericType(typeof(TResource), typeof(TKey));
+
+            // LastAccessTimestamp is a non-public mutable field on the internal LimiterEntry type.
+            var lastAccessField = limiterEntryType.GetField("LastAccessTimestamp", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
+            Assert.NotNull(lastAccessField);
+
+            var entry = GetLimiterEntry<TResource, TKey>(limiter, key);
+            long backdatedTimestamp = System.Diagnostics.Stopwatch.GetTimestamp() - (System.Diagnostics.Stopwatch.Frequency * 60);
+            lastAccessField.SetValue(entry, backdatedTimestamp);
+        }
+
+        private static object GetLimiterEntry<TResource, TKey>(PartitionedRateLimiter<TResource> limiter, TKey key)
+        {
+            var lazyEntry = GetLazyLimiterEntry<TResource, TKey>(limiter, key);
+            Assert.NotNull(lazyEntry);
+
+            // Force creation of the Lazy<LimiterEntry>.Value. Use Type.GetType with a literal string
+            // so trimming can see the LimiterEntry type, then construct the closed Lazy<LimiterEntry>.
+            var limiterEntryTypeDef = Type.GetType(DefaultPartitionedRateLimiterEntryTypeName);
+            Assert.NotNull(limiterEntryTypeDef);
+            var limiterEntryType = limiterEntryTypeDef.MakeGenericType(typeof(TResource), typeof(TKey));
+            var lazyType = typeof(Lazy<>).MakeGenericType(limiterEntryType);
+            var valueProperty = lazyType.GetProperty("Value");
+            Assert.NotNull(valueProperty);
+            var entry = valueProperty.GetValue(lazyEntry);
+            Assert.NotNull(entry);
+            return entry;
+        }
+
+        private static object GetLazyLimiterEntry<TResource, TKey>(PartitionedRateLimiter<TResource> limiter, TKey key)
+        {
+            // Use Type.GetType so trimming can see what type we're reflecting on, then construct the
+            // closed DefaultPartitionedRateLimiter<TResource, TKey> type for the requested arguments.
+            var limiterTypeDef = Type.GetType(DefaultPartitionedRateLimiterTypeName);
+            Assert.NotNull(limiterTypeDef);
+            var limiterType = limiterTypeDef.MakeGenericType(typeof(TResource), typeof(TKey));
+
+            var limitersField = limiterType.GetField("_limiters", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
+            Assert.NotNull(limitersField);
+            Assert.True(limiterType.IsInstanceOfType(limiter));
+            var limitersDict = (System.Collections.IDictionary)limitersField.GetValue(limiter);
+            Assert.NotNull(limitersDict);
+
+            return limitersDict[key];
+        }
+
         [Fact]
         public async Task AllIdleLimitersCleanedUp_DisposeThrows()
         {