diff --git a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs index 29b16440cd8e03..3329d3dd1f6cbc 100644 --- a/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs +++ b/src/libraries/System.Threading.RateLimiting/src/System/Threading/RateLimiting/SlidingWindowRateLimiter.cs @@ -24,6 +24,13 @@ public sealed class SlidingWindowRateLimiter : ReplenishingRateLimiter private long _failedLeasesCount; private long _successfulLeasesCount; + /// + /// Function to calculate elapsed time from a given tick value. + /// Defaults to . + /// In tests, this field can be reassigned via reflection to inject custom time behavior without modifying the public API. + /// + private readonly Func _getElapsedTime = RateLimiterHelper.GetElapsedTime; + private readonly Timer? _renewTimer; private readonly SlidingWindowRateLimiterOptions _options; private readonly TimeSpan _replenishmentPeriod; @@ -124,8 +131,19 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount) return SuccessfulLease; } - Interlocked.Increment(ref _failedLeasesCount); - return FailedLease; + lock (Lock) + { + // Edge case where the check before the lock showed 0 available permits but when we got the lock some permits were now available + if (_permitCount > 0) + { + Interlocked.Increment(ref _successfulLeasesCount); + return SuccessfulLease; + } + + Interlocked.Increment(ref _failedLeasesCount); + // Still holding Lock from the top of this block - required by CreateFailedSlidingWindowLease. + return CreateFailedSlidingWindowLease(0); + } } lock (Lock) @@ -135,9 +153,9 @@ protected override RateLimitLease AttemptAcquireCore(int permitCount) return lease; } - // TODO: Acquire additional metadata during a failed lease decision Interlocked.Increment(ref _failedLeasesCount); - return FailedLease; + // Still holding Lock from the top of this block - required by CreateFailedSlidingWindowLease. + return CreateFailedSlidingWindowLease(permitCount); } } @@ -204,8 +222,9 @@ protected override ValueTask AcquireAsyncCore(int permitCount, C else { Interlocked.Increment(ref _failedLeasesCount); - // Don't queue if queue limit reached and QueueProcessingOrder is OldestFirst - return new ValueTask(FailedLease); + // Don't queue if queue limit reached and QueueProcessingOrder is OldestFirst. + // Still holding Lock from the top of this block - required by CreateFailedSlidingWindowLease. + return new ValueTask(CreateFailedSlidingWindowLease(permitCount)); } } @@ -218,6 +237,50 @@ protected override ValueTask AcquireAsyncCore(int permitCount, C } } + // All callers of this method must hold Lock. _requestsPerSegment and _currentSegmentIndex are only otherwise + // mutated under Lock (in ReplenishInternal), so reading them here without it would risk a torn read (a mix of + // pre- and post-replenish state) rather than a clean snapshot. If a new call site is ever added, it must + // acquire Lock before calling this method. + // + // Note: like FixedWindowRateLimiter.CreateFailedWindowLease, this estimate ignores queue state. If permits are + // available (_permitCount >= permitCount) but the request still fails because it can't jump an OldestFirst + // queue, neededPermits floors at 1 and the estimate is based on segment replenishment timing even though the + // real blocker is queue draining, not permit scarcity. + private SlidingWindowLease CreateFailedSlidingWindowLease(int permitCount) + { + // How many permits still need to free up to satisfy this request. Math.Max(1, ...) covers the + // permitCount == 0 fast-path case, where we just need *any* permit to become available + // (i.e. _permitCount > 0). + int neededPermits = Math.Max(1, permitCount - _permitCount); + + // Walk forward from the oldest segment (the next to be freed) toward the newest, accumulating how many + // permits each segment will release, until we've accounted for enough to satisfy the request. + int cumulativeFreed = 0; + int periodsUntilEnough = _options.SegmentsPerWindow; + for (int i = 1; i <= _options.SegmentsPerWindow; i++) + { + int segmentIndex = (_currentSegmentIndex + i) % _options.SegmentsPerWindow; + cumulativeFreed += _requestsPerSegment[segmentIndex]; + if (cumulativeFreed >= neededPermits) + { + periodsUntilEnough = i; + break; + } + } + + TimeSpan? elapsedInCurrentPeriod = _getElapsedTime(_lastReplenishmentTick); + long remainingTicksInCurrentPeriod = _replenishmentPeriod.Ticks - (elapsedInCurrentPeriod?.Ticks ?? 0); + long remainingTicks = _replenishmentPeriod.Ticks * (periodsUntilEnough - 1) + remainingTicksInCurrentPeriod; + + // Clamp to zero if negative (segment(s) expired but not yet replenished) + if (remainingTicks < 0) + { + remainingTicks = 0; + } + + return new SlidingWindowLease(false, TimeSpan.FromTicks(remainingTicks)); + } + private bool TryLeaseUnsynchronized(int permitCount, [NotNullWhen(true)] out RateLimitLease? lease) { ThrowIfDisposed(); @@ -298,7 +361,7 @@ private void ReplenishInternal(long nowTicks) return; } - _lastReplenishmentTick = nowTicks; + Volatile.Write(ref _lastReplenishmentTick, nowTicks); // Increment the current segment index while move the window // We need to know the no. of requests that were acquired in a segment previously to ensure that we don't acquire more than the permit limit. diff --git a/src/libraries/System.Threading.RateLimiting/tests/SlidingWindowRateLimiterTests.cs b/src/libraries/System.Threading.RateLimiting/tests/SlidingWindowRateLimiterTests.cs index d2e92bd28fb23b..1e03a43fb17ae4 100644 --- a/src/libraries/System.Threading.RateLimiting/tests/SlidingWindowRateLimiterTests.cs +++ b/src/libraries/System.Threading.RateLimiting/tests/SlidingWindowRateLimiterTests.cs @@ -253,7 +253,7 @@ public override async Task CanAcquireResourceAsync_QueuesAndGrabsNewest() [Fact] public override async Task FailsWhenQueuingMoreThanLimit_OldestFirst() { - var limiter = new SlidingWindowRateLimiter(new SlidingWindowRateLimiterOptions + var options = new SlidingWindowRateLimiterOptions { PermitLimit = 1, QueueProcessingOrder = QueueProcessingOrder.OldestFirst, @@ -261,12 +261,18 @@ public override async Task FailsWhenQueuingMoreThanLimit_OldestFirst() Window = TimeSpan.FromMilliseconds(1), SegmentsPerWindow = 2, AutoReplenishment = false - }); + }; + var limiter = new SlidingWindowRateLimiter(options); + using var lease = limiter.AttemptAcquire(1); var wait = limiter.AcquireAsync(1); + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + var failedLease = await limiter.AcquireAsync(1); Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var timeSpan)); + Assert.Equal(options.Window, timeSpan); } [Fact] @@ -1308,6 +1314,283 @@ public void AutoReplenishIgnoresTimerJitter() Assert.Equal(7, limiter.GetStatistics().CurrentAvailablePermits); } + [Fact] + public async Task RetryMetadataOnFailedAcquireAsync() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 2, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + using var lease = limiter.AttemptAcquire(2); + + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + + var failedLease = await limiter.AcquireAsync(2); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter.Name, out var metadata)); + var metaDataTime = Assert.IsType(metadata); + Assert.Equal(options.Window.Ticks, metaDataTime.Ticks); + + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata)); + Assert.Equal(options.Window.Ticks, typedMetadata.Ticks); + Assert.Collection(failedLease.MetadataNames, item => Assert.Equal(MetadataName.RetryAfter.Name, item)); + } + + [Fact] + public async Task CorrectRetryMetadataWithQueuedItem() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 2, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + using var lease = limiter.AttemptAcquire(2); + // A queued item is present to confirm queue occupancy doesn't factor into the RetryAfter estimate + // (see CreateFailedSlidingWindowLease) - the failed request below is rejected immediately because 2 permits + // exceeds QueueLimit, not because of the queued item itself. + var wait = limiter.AcquireAsync(1); + Assert.False(wait.IsCompleted); + + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + + var failedLease = await limiter.AcquireAsync(2); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata)); + Assert.Equal(options.Window.Ticks, typedMetadata.Ticks); + } + + [Fact] + public async Task CorrectRetryMetadataWithNonZeroAvailableItems() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 3, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + // 2 of the 3 permits are consumed, leaving 1 available - but that 1 isn't enough to satisfy the 3-permit + // request below, so it should still fail and estimate against the segment holding the other 2. + using var lease = limiter.AttemptAcquire(2); + + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + + var failedLease = await limiter.AcquireAsync(3); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var typedMetadata)); + Assert.Equal(options.Window.Ticks, typedMetadata.Ticks); + } + + [Fact] + public async Task RetryAfterWithPartiallyElapsedSegment() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 1, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 0, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + using var lease = limiter.AttemptAcquire(1); + + // 2 seconds elapsed since the last replenishment tick. The single permit is in the current segment, so + // the deficit isn't resolved until the walk wraps all the way around: (SegmentsPerWindow - 1) full + // periods, plus the remainder of the current one. + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.FromSeconds(2)); + + var failedLease = await limiter.AcquireAsync(1); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var timeSpan)); + Assert.Equal(TimeSpan.FromSeconds(18), timeSpan); // 20 - 2 = 18 + } + + [Fact] + public async Task RetryAfterClampsToZeroWhenWindowExpired() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 1, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 0, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + using var lease = limiter.AttemptAcquire(1); + + // 25 seconds elapsed - longer than the full window, so even after walking all the way around to the + // segment holding the permit, the estimate goes negative and should clamp to zero rather than report a + // negative TimeSpan. + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.FromSeconds(25)); + + var failedLease = await limiter.AcquireAsync(1); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var timeSpan)); + Assert.Equal(TimeSpan.Zero, timeSpan); // Clamped to zero, not negative + } + + [Fact] + public void RetryMetadataOnFailedAttemptAcquire() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 2, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + using var lease = limiter.AttemptAcquire(2); + + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + + // Exercises the general locked branch in AttemptAcquireCore + var failedLease = limiter.AttemptAcquire(2); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var timeSpan)); + Assert.Equal(options.Window, timeSpan); + } + + [Fact] + public void RetryMetadataOnFailedAttemptAcquireZeroPermitCount() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 1, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + Window = TimeSpan.FromSeconds(20), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + using var lease = limiter.AttemptAcquire(1); + + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + + // permitCount == 0 with no permits available hits the fast-path failure branch that was just moved inside + // the lock for the race condition fix. + var failedLease = limiter.AttemptAcquire(0); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var timeSpan)); + Assert.Equal(options.Window, timeSpan); + } + + [Fact] + public async Task RetryAfterAccountsForPermitsAcrossMultipleSegments() + { + var options = new SlidingWindowRateLimiterOptions + { + PermitLimit = 4, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 0, + Window = TimeSpan.FromSeconds(8), + SegmentsPerWindow = 4, + AutoReplenishment = false + }; + var limiter = new SlidingWindowRateLimiter(options); + + // Segment 0 holds 3 permits + using var lease1 = limiter.AttemptAcquire(3); + Assert.True(lease1.IsAcquired); + + // Advance one period (2s) so segment 0 becomes an older segment relative to the new current one + Replenish(limiter, 2000L); + + // Segment 1 (now current) holds the last permit - all 4 are now consumed + using var lease2 = limiter.AttemptAcquire(1); + Assert.True(lease2.IsAcquired); + Assert.Equal(0, limiter.GetStatistics().CurrentAvailablePermits); + + SetElapsedTimeSinceLastReplenishment(limiter, TimeSpan.Zero); + + // Only 2 permits are needed - segment 0's 3 permits alone satisfy that, one period before the walk would + // otherwise wrap all the way back to the current segment. + var failedLease = await limiter.AcquireAsync(2); + Assert.False(failedLease.IsAcquired); + Assert.True(failedLease.TryGetMetadata(MetadataName.RetryAfter, out var timeSpan)); + + // periodsUntilEnough = 3 (two empty segments walked, then segment 0 satisfies the deficit) + var periodTicks = options.Window.Ticks / options.SegmentsPerWindow; + Assert.Equal(TimeSpan.FromTicks(periodTicks * 3), timeSpan); + } + + [Fact] + public async Task AttemptAcquireZero_PermitBecomesAvailableWhileWaitingForLock() + { + var limiter = new SlidingWindowRateLimiter(new SlidingWindowRateLimiterOptions + { + PermitLimit = 1, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 1, + Window = TimeSpan.FromMilliseconds(2), + SegmentsPerWindow = 2, + AutoReplenishment = false + }); + + using var lease = limiter.AttemptAcquire(1); + Assert.True(lease.IsAcquired); + Assert.Equal(0, limiter.GetStatistics().CurrentAvailablePermits); + + var lockProperty = typeof(SlidingWindowRateLimiter).GetProperty("Lock", Reflection.BindingFlags.NonPublic | Reflection.BindingFlags.Instance)!; + var lockObject = lockProperty.GetValue(limiter)!; + + Task acquireTask; + + // Hold the limiter's own lock on this thread, so a concurrent AttemptAcquire(0) can pass + // its outer, unlocked _permitCount check (which sees 0) but then has to block on the lock + // before it reaches the recheck. That's exactly the race window the fix is meant to close. + Monitor.Enter(lockObject); + try + { + acquireTask = Task.Run(() => limiter.AttemptAcquire(0)); + + // Two segments in this window - the first Replenish call advances to the (empty) + // second segment; the second call wraps back around and frees the permit that was + // recorded against the first segment. + Replenish(limiter, 1L); + Replenish(limiter, 1L); + Assert.True(limiter.GetStatistics().CurrentAvailablePermits > 0); + + // Do not await acquireTask here - see comment above about why. + } + finally + { + Monitor.Exit(lockObject); + } + + var result = await acquireTask; + Assert.True(result.IsAcquired); + Assert.Equal(0, limiter.GetStatistics().TotalFailedLeases); // the bug would have incremented this + } + private static readonly double TickFrequency = (double)TimeSpan.TicksPerSecond / Stopwatch.Frequency; static internal void Replenish(SlidingWindowRateLimiter limiter, long addMilliseconds) @@ -1317,5 +1600,14 @@ static internal void Replenish(SlidingWindowRateLimiter limiter, long addMillise var currentTick = (long)internalTick.GetValue(limiter); replenishInternalMethod.Invoke(limiter, new object[] { currentTick + addMilliseconds * (long)(TimeSpan.TicksPerMillisecond / TickFrequency) }); } + + // Function that replaces the _getElapsedTime function in SlidingWindowRateLimiter to return a specified + // TimeSpan. Used for testing the RetryAfter metadata on failed leases. + static internal void SetElapsedTimeSinceLastReplenishment(SlidingWindowRateLimiter limiter, TimeSpan elapsedTimeSinceLastReplenishment) + { + var _getElapsedTimeField = typeof(SlidingWindowRateLimiter).GetField("_getElapsedTime", Reflection.BindingFlags.NonPublic | Reflection.BindingFlags.Instance)!; + Func overrideFunc = (_) => elapsedTimeSinceLastReplenishment; + _getElapsedTimeField.SetValue(limiter, overrideFunc); + } } }