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 84aa3ae9feb96..0822fe6a3f7d3 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 @@ -34,12 +34,19 @@ internal sealed class DefaultPartitionedRateLimiter : Partition public DefaultPartitionedRateLimiter(Func> partitioner, IEqualityComparer? equalityComparer = null) + : this(partitioner, equalityComparer, TimeSpan.FromMilliseconds(100)) + { + } + + // Extra ctor for testing purposes, primarily used when wanting to test the timer manually + private DefaultPartitionedRateLimiter(Func> partitioner, + IEqualityComparer? equalityComparer, TimeSpan timerInterval) { _limiters = new Dictionary>(equalityComparer); _partitioner = partitioner; // TODO: Figure out what interval we should use - _timer = new TimerAwaitable(TimeSpan.FromMilliseconds(100), TimeSpan.FromMilliseconds(100)); + _timer = new TimerAwaitable(timerInterval, timerInterval); _timerTask = RunTimer(); } diff --git a/src/libraries/System.Threading.RateLimiting/tests/Infrastructure/Utils.cs b/src/libraries/System.Threading.RateLimiting/tests/Infrastructure/Utils.cs index e2a4c1150f07c..d69b7cc47a139 100644 --- a/src/libraries/System.Threading.RateLimiting/tests/Infrastructure/Utils.cs +++ b/src/libraries/System.Threading.RateLimiting/tests/Infrastructure/Utils.cs @@ -1,11 +1,10 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. -using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Linq; -using System.Text; +using System.Reflection; using System.Threading.Tasks; using Xunit; @@ -13,25 +12,31 @@ namespace System.Threading.RateLimiting.Tests { internal static class Utils { - internal static Func StopTimerAndGetTimerFunc(PartitionedRateLimiter limiter) + // Creates a DefaultPartitionedRateLimiter with the timer effectively disabled + internal static PartitionedRateLimiter CreatePartitionedLimiterWithoutTimer(Func> partitioner) { - var innerTimer = limiter.GetType().GetField("_timer", Reflection.BindingFlags.NonPublic | Reflection.BindingFlags.Instance); + var limiterType = Assembly.GetAssembly(typeof(PartitionedRateLimiter<>)).GetType("System.Threading.RateLimiting.DefaultPartitionedRateLimiter`2"); + Assert.NotNull(limiterType); + + var genericLimiterType = limiterType.MakeGenericType(typeof(TResource), typeof(TKey)); + Assert.NotNull(genericLimiterType); + + var limiterCtor = genericLimiterType.GetConstructors(BindingFlags.Instance | BindingFlags.NonPublic)[0]; + Assert.NotNull(limiterCtor); + + return (PartitionedRateLimiter)limiterCtor.Invoke(new object[] { partitioner, null, TimeSpan.FromMinutes(10) }); + } + + // Gets and runs the Heartbeat function on the DefaultPartitionedRateLimiter + internal static Task RunTimerFunc(PartitionedRateLimiter limiter) + { + var innerTimer = limiter.GetType().GetField("_timer", BindingFlags.NonPublic | BindingFlags.Instance); Assert.NotNull(innerTimer); - var timerStopMethod = innerTimer.FieldType.GetMethod("Stop"); - Assert.NotNull(timerStopMethod); - // Stop the current Timer so it doesn't fire unexpectedly - timerStopMethod.Invoke(innerTimer.GetValue(limiter), Array.Empty()); - - // Create a new Timer object so that disposing the PartitionedRateLimiter doesn't fail with an ODE, but this new Timer wont actually do anything - var timerCtor = innerTimer.FieldType.GetConstructor(new Type[] { typeof(TimeSpan), typeof(TimeSpan) }); - Assert.NotNull(timerCtor); - var newTimer = timerCtor.Invoke(new object[] { TimeSpan.FromMinutes(10), TimeSpan.FromMinutes(10) }); - Assert.NotNull(newTimer); - innerTimer.SetValue(limiter, newTimer); - - var timerLoopMethod = limiter.GetType().GetMethod("Heartbeat", Reflection.BindingFlags.NonPublic | Reflection.BindingFlags.Instance); + + var timerLoopMethod = limiter.GetType().GetMethod("Heartbeat", BindingFlags.NonPublic | BindingFlags.Instance); Assert.NotNull(timerLoopMethod); - return () => (Task)timerLoopMethod.Invoke(limiter, Array.Empty()); + + return (Task)timerLoopMethod.Invoke(limiter, Array.Empty()); } } diff --git a/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs b/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs index e99b603b08bb6..40b3d719ef071 100644 --- a/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs +++ b/src/libraries/System.Threading.RateLimiting/tests/PartitionedRateLimiterTests.cs @@ -452,7 +452,7 @@ public async Task IdleLimiterIsCleanedUp() { CustomizableLimiter innerLimiter = null; var factoryCallCount = 0; - using var limiter = PartitionedRateLimiter.Create(resource => + using var limiter = Utils.CreatePartitionedLimiterWithoutTimer(resource => { return RateLimitPartition.Get(1, _ => { @@ -462,8 +462,6 @@ public async Task IdleLimiterIsCleanedUp() }); }); - var timerLoopMethod = Utils.StopTimerAndGetTimerFunc(limiter); - var lease = limiter.Acquire(""); Assert.True(lease.IsAcquired); @@ -477,7 +475,7 @@ public async Task IdleLimiterIsCleanedUp() }; innerLimiter.IdleDurationImpl = () => TimeSpan.FromMinutes(1); - await timerLoopMethod(); + await Utils.RunTimerFunc(limiter); // Limiter is disposed when timer runs and sees that IdleDuration is greater than idle limit await tcs.Task; @@ -494,7 +492,7 @@ public async Task AllIdleLimitersCleanedUp_DisposeThrows() { CustomizableLimiter innerLimiter1 = null; CustomizableLimiter innerLimiter2 = null; - using var limiter = PartitionedRateLimiter.Create(resource => + using var limiter = Utils.CreatePartitionedLimiterWithoutTimer(resource => { if (resource == "1") { @@ -514,8 +512,6 @@ public async Task AllIdleLimitersCleanedUp_DisposeThrows() } }); - var timerLoopMethod = Utils.StopTimerAndGetTimerFunc(limiter); - var lease = limiter.Acquire("1"); Assert.True(lease.IsAcquired); Assert.NotNull(innerLimiter1); @@ -538,7 +534,7 @@ public async Task AllIdleLimitersCleanedUp_DisposeThrows() innerLimiter2.IdleDurationImpl = () => TimeSpan.FromMinutes(1); // Run Timer - var ex = await Assert.ThrowsAsync(() => timerLoopMethod()); + var ex = await Assert.ThrowsAsync(() => Utils.RunTimerFunc(limiter)); Assert.True(dispose1Called); Assert.True(dispose2Called); @@ -552,7 +548,7 @@ public async Task ThrowingTryReplenishDoesNotPreventIdleLimiterBeingCleanedUp() CustomizableReplenishingLimiter replenishLimiter = new CustomizableReplenishingLimiter(); CustomizableLimiter idleLimiter = null; var factoryCallCount = 0; - using var limiter = PartitionedRateLimiter.Create(resource => + using var limiter = Utils.CreatePartitionedLimiterWithoutTimer(resource => { if (resource == "1") { @@ -569,8 +565,6 @@ public async Task ThrowingTryReplenishDoesNotPreventIdleLimiterBeingCleanedUp() }); }); - var timerLoopMethod = Utils.StopTimerAndGetTimerFunc(limiter); - // Add the replenishing limiter to the internal storage limiter.Acquire("2"); var lease = limiter.Acquire("1"); @@ -589,7 +583,7 @@ public async Task ThrowingTryReplenishDoesNotPreventIdleLimiterBeingCleanedUp() }; idleLimiter.IdleDurationImpl = () => TimeSpan.FromMinutes(1); - var ex = await Assert.ThrowsAsync(() => timerLoopMethod()); + var ex = await Assert.ThrowsAsync(() => Utils.RunTimerFunc(limiter)); Assert.Single(ex.InnerExceptions); // Wait for Timer to run again which will see the throwing TryReplenish and an idle limiter it needs to clean-up