diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs index b2966bf85c05ae..f9cf3393698c1b 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs +++ b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/DefaultPartitionedRateLimiter.cs @@ -16,14 +16,14 @@ internal sealed class DefaultPartitionedRateLimiter : Partition private static readonly TimeSpan s_idleTimeLimit = TimeSpan.FromSeconds(10); // TODO: Look at ConcurrentDictionary to try and avoid a global lock - private readonly Dictionary> _limiters; + private readonly Dictionary> _limiters; private bool _disposed; private readonly TaskCompletionSource _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>> _cachedLimiters = new(); + private readonly List>> _cachedLimiters = new(); private bool _cacheInvalid; private readonly List _limitersToDispose = new(); private readonly TimerAwaitable _timer; @@ -42,7 +42,7 @@ public DefaultPartitionedRateLimiter(Func> p private DefaultPartitionedRateLimiter(Func> partitioner, IEqualityComparer? equalityComparer, TimeSpan timerInterval) { - _limiters = new Dictionary>(equalityComparer); + _limiters = new Dictionary>(equalityComparer); _partitioner = partitioner; _timer = new TimerAwaitable(timerInterval, timerInterval); @@ -87,20 +87,36 @@ protected override ValueTask AcquireAsyncCore(TResource resource private RateLimiter GetRateLimiter(TResource resource) { RateLimitPartition partition = _partitioner(resource); - Lazy? limiter; + Lazy? 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(() => 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(() => 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> limiter in _limiters) + foreach (KeyValuePair> 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? exceptions = null; - foreach (KeyValuePair> limiter in _limiters) + foreach (KeyValuePair> 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? aggregateExceptions = null; // cachedLimiters is safe to use outside the lock because it is only updated by the Timer - foreach (KeyValuePair> rateLimiter in _cachedLimiters) + foreach (KeyValuePair> 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; + } } } diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/RateLimiterHelper.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/RateLimiterHelper.cs index e538fc720072dd..df3f1ecb3e0e66 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/RateLimiterHelper.cs +++ b/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); diff --git a/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs b/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs index ff6e7c4264941b..2c86be1bb781e6 100644 --- a/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs +++ b/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(resource => + { + return RateLimitPartition.GetNoLimiter(1); + }); + + var lease = limiter.AttemptAcquire(""); + Assert.True(lease.IsAcquired); + Assert.NotNull(GetLazyLimiterEntry(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(limiter, key: 1); + + await Utils.RunTimerFunc(limiter); + Assert.Null(GetLazyLimiterEntry(limiter, key: 1)); + + lease = limiter.AttemptAcquire(""); + Assert.True(lease.IsAcquired); + Assert.NotNull(GetLazyLimiterEntry(limiter, key: 1)); + } + + [Fact] + public async Task LimiterWithNullIdleDurationIsNotCleanedUp() + { + var factoryCallCount = 0; + using var limiter = Utils.CreatePartitionedLimiterWithoutTimer(resource => + { + return RateLimitPartition.Get(1, _ => + { + factoryCallCount++; + return new CustomizableLimiter + { + IdleDurationImpl = () => null + }; + }); + }); + + var lease = limiter.AttemptAcquire(""); + Assert.True(lease.IsAcquired); + Assert.Equal(1, factoryCallCount); + + BackdateLastAccessTimestamp(limiter, key: 1); + + await Utils.RunTimerFunc(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(PartitionedRateLimiter 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(limiter, key); + long backdatedTimestamp = System.Diagnostics.Stopwatch.GetTimestamp() - (System.Diagnostics.Stopwatch.Frequency * 60); + lastAccessField.SetValue(entry, backdatedTimestamp); + } + + private static object GetLimiterEntry(PartitionedRateLimiter limiter, TKey key) + { + var lazyEntry = GetLazyLimiterEntry(limiter, key); + Assert.NotNull(lazyEntry); + + // Force creation of the Lazy.Value. Use Type.GetType with a literal string + // so trimming can see the LimiterEntry type, then construct the closed Lazy. + 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(PartitionedRateLimiter limiter, TKey key) + { + // Use Type.GetType so trimming can see what type we're reflecting on, then construct the + // closed DefaultPartitionedRateLimiter 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() {