[RateLimiting] Fix test race with timer (#72118)

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

@@ -34,12 +34,19 @@ internal sealed class DefaultPartitionedRateLimiter<TResource, TKey> : Partition
 
         public DefaultPartitionedRateLimiter(Func<TResource, RateLimitPartition<TKey>> partitioner,
             IEqualityComparer<TKey>? 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<TResource, RateLimitPartition<TKey>> partitioner,
+            IEqualityComparer<TKey>? equalityComparer, TimeSpan timerInterval)
         {
             _limiters = new Dictionary<TKey, Lazy<RateLimiter>>(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();
         }
 

src/libraries/System.Threading.RateLimiting/tests/Infrastructure/Utils.cs

@@ -1,37 +1,42 @@
 // 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;
 
 namespace System.Threading.RateLimiting.Tests
 {
     internal static class Utils
     {
-        internal static Func<Task> StopTimerAndGetTimerFunc<T>(PartitionedRateLimiter<T> limiter)
+        // Creates a DefaultPartitionedRateLimiter with the timer effectively disabled
+        internal static PartitionedRateLimiter<TResource> CreatePartitionedLimiterWithoutTimer<TResource, TKey>(Func<TResource, RateLimitPartition<TKey>> 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<TResource>)limiterCtor.Invoke(new object[] { partitioner, null, TimeSpan.FromMinutes(10) });
+        }
+
+        // Gets and runs the Heartbeat function on the DefaultPartitionedRateLimiter
+        internal static Task RunTimerFunc<T>(PartitionedRateLimiter<T> 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<object>());
-
-            // 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<object>());
+
+            return (Task)timerLoopMethod.Invoke(limiter, Array.Empty<object>());
         }
     }
 

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<string, int>(resource =>
+            using var limiter = Utils.CreatePartitionedLimiterWithoutTimer<string, int>(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<string, int>(resource =>
+            using var limiter = Utils.CreatePartitionedLimiterWithoutTimer<string, int>(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<AggregateException>(() => timerLoopMethod());
+            var ex = await Assert.ThrowsAsync<AggregateException>(() => 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<string, int>(resource =>
+            using var limiter = Utils.CreatePartitionedLimiterWithoutTimer<string, int>(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<AggregateException>(() => timerLoopMethod());
+            var ex = await Assert.ThrowsAsync<AggregateException>(() => 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