From 773252b44a50add9b3c9afcd3da7bb59d2c139e6 Mon Sep 17 00:00:00 2001 From: maliming Date: Fri, 6 Mar 2026 17:48:56 +0800 Subject: [PATCH] feat: implement early break in Phase 2 for multi-rule policies and add corresponding tests --- .../Checker/OperationRateLimitingChecker.cs | 23 +++- ...nRateLimitingPhase2EarlyBreakTestModule.cs | 102 ++++++++++++++++++ ...OperationRateLimitingCheckerFixes_Tests.cs | 53 +++++++++ 3 files changed, 174 insertions(+), 4 deletions(-) create mode 100644 framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/AbpOperationRateLimitingPhase2EarlyBreakTestModule.cs diff --git a/framework/src/Volo.Abp.OperationRateLimiting/Volo/Abp/OperationRateLimiting/Checker/OperationRateLimitingChecker.cs b/framework/src/Volo.Abp.OperationRateLimiting/Volo/Abp/OperationRateLimiting/Checker/OperationRateLimitingChecker.cs index 3b3006b248..56f22a79c6 100644 --- a/framework/src/Volo.Abp.OperationRateLimiting/Volo/Abp/OperationRateLimiting/Checker/OperationRateLimitingChecker.cs +++ b/framework/src/Volo.Abp.OperationRateLimiting/Volo/Abp/OperationRateLimiting/Checker/OperationRateLimitingChecker.cs @@ -65,16 +65,31 @@ public class OperationRateLimitingChecker : IOperationRateLimitingChecker, ITran ThrowRateLimitException(policy, aggregatedResult, context); } - // Phase 2: All rules pass - now increment all counters. - // Also guard against a concurrent race where another request consumed the last quota + // Phase 2: All rules passed in Phase 1 - now increment counters. + // Guard against concurrent races where another request consumed the last quota // between Phase 1 and Phase 2. + // Once any rule fails during increment, stop incrementing subsequent rules + // to minimize wasted quota. Remaining rules use read-only check instead. var incrementResults = new List(); + var phase2Failed = false; foreach (var rule in rules) { - incrementResults.Add(await rule.AcquireAsync(context)); + if (phase2Failed) + { + incrementResults.Add(await rule.CheckAsync(context)); + } + else + { + var result = await rule.AcquireAsync(context); + incrementResults.Add(result); + if (!result.IsAllowed) + { + phase2Failed = true; + } + } } - if (incrementResults.Any(r => !r.IsAllowed)) + if (phase2Failed) { var aggregatedResult = AggregateResults(incrementResults, policy); ThrowRateLimitException(policy, aggregatedResult, context); diff --git a/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/AbpOperationRateLimitingPhase2EarlyBreakTestModule.cs b/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/AbpOperationRateLimitingPhase2EarlyBreakTestModule.cs new file mode 100644 index 0000000000..b36b9778cd --- /dev/null +++ b/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/AbpOperationRateLimitingPhase2EarlyBreakTestModule.cs @@ -0,0 +1,102 @@ +using System; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.DependencyInjection.Extensions; +using Volo.Abp.Autofac; +using Volo.Abp.ExceptionHandling; +using Volo.Abp.Modularity; + +namespace Volo.Abp.OperationRateLimiting; + +/// +/// A mock store that simulates a multi-rule Phase 2 race condition: +/// - GetAsync always reports quota available (Phase 1 passes for all rules). +/// - IncrementAsync succeeds for the first call, fails on the second call +/// (simulating a concurrent race on Rule2), and tracks total increment calls +/// so tests can verify that Rule3 was never incremented (early break). +/// +internal class MultiRuleRaceConditionSimulatorStore : IOperationRateLimitingStore +{ + private int _incrementCallCount; + + /// + /// Total number of IncrementAsync calls made. + /// + public int IncrementCallCount => _incrementCallCount; + + public Task GetAsync(string key, TimeSpan duration, int maxCount) + { + return Task.FromResult(new OperationRateLimitingStoreResult + { + IsAllowed = true, + CurrentCount = 0, + MaxCount = maxCount + }); + } + + public Task IncrementAsync(string key, TimeSpan duration, int maxCount) + { + var callIndex = Interlocked.Increment(ref _incrementCallCount); + + if (callIndex == 2) + { + // Second rule: simulate concurrent race - another request consumed the last slot. + return Task.FromResult(new OperationRateLimitingStoreResult + { + IsAllowed = false, + CurrentCount = maxCount, + MaxCount = maxCount, + RetryAfter = duration + }); + } + + // First rule (and any others if early break fails): succeed. + return Task.FromResult(new OperationRateLimitingStoreResult + { + IsAllowed = true, + CurrentCount = 1, + MaxCount = maxCount + }); + } + + public Task ResetAsync(string key) + { + return Task.CompletedTask; + } +} + +[DependsOn( + typeof(AbpOperationRateLimitingModule), + typeof(AbpExceptionHandlingModule), + typeof(AbpTestBaseModule), + typeof(AbpAutofacModule) +)] +public class AbpOperationRateLimitingPhase2EarlyBreakTestModule : AbpModule +{ + public override void ConfigureServices(ServiceConfigurationContext context) + { + context.Services.Replace( + ServiceDescriptor.Singleton()); + + Configure(options => + { + // 3-rule composite policy: all PartitionByParameter with different durations + // so they generate unique cache keys and don't trigger duplicate rule validation. + options.AddPolicy("TestMultiRuleRacePolicy", policy => + { + policy.AddRule(rule => rule + .WithFixedWindow(TimeSpan.FromHours(1), maxCount: 5) + .PartitionByParameter()); + + policy.AddRule(rule => rule + .WithFixedWindow(TimeSpan.FromHours(2), maxCount: 5) + .PartitionByParameter()); + + policy.AddRule(rule => rule + .WithFixedWindow(TimeSpan.FromHours(3), maxCount: 5) + .PartitionByParameter()); + }); + }); + } +} diff --git a/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/OperationRateLimitingCheckerFixes_Tests.cs b/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/OperationRateLimitingCheckerFixes_Tests.cs index 8d15d2a2b9..fce15fa466 100644 --- a/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/OperationRateLimitingCheckerFixes_Tests.cs +++ b/framework/test/Volo.Abp.OperationRateLimiting.Tests/Volo/Abp/OperationRateLimiting/OperationRateLimitingCheckerFixes_Tests.cs @@ -97,6 +97,59 @@ public class OperationRateLimitingCheckerPhase1_Tests : OperationRateLimitingTes } } +/// +/// Tests for Phase 2 early break: when a multi-rule policy encounters a race condition +/// in Phase 2 (Rule2 fails), Rule3 should NOT be incremented. +/// Uses a mock store where IncrementAsync fails on the 2nd call. +/// +public class OperationRateLimitingCheckerPhase2EarlyBreak_Tests + : AbpIntegratedTest +{ + protected override void SetAbpApplicationCreationOptions(AbpApplicationCreationOptions options) + { + options.UseAutofac(); + } + + [Fact] + public async Task Should_Not_Increment_Remaining_Rules_After_Phase2_Failure() + { + // 3-rule policy. Mock store: Rule1 increment succeeds, Rule2 increment fails (race), + // Rule3 should NOT be incremented due to early break. + var checker = GetRequiredService(); + var store = (MultiRuleRaceConditionSimulatorStore)GetRequiredService(); + var context = new OperationRateLimitingContext { Parameter = "early-break-test" }; + + var exception = await Assert.ThrowsAsync(async () => + { + await checker.CheckAsync("TestMultiRuleRacePolicy", context); + }); + + exception.PolicyName.ShouldBe("TestMultiRuleRacePolicy"); + exception.Result.IsAllowed.ShouldBeFalse(); + + // Key assertion: only 2 IncrementAsync calls were made (Rule1 + Rule2). + // Rule3 was skipped (used CheckAsync instead) due to early break. + store.IncrementCallCount.ShouldBe(2); + } + + [Fact] + public async Task Should_Include_All_Rule_Results_Despite_Early_Break() + { + // Even with early break, the aggregated result should contain all 3 rules + // (Rule3 via CheckAsync instead of AcquireAsync). + var checker = GetRequiredService(); + var context = new OperationRateLimitingContext { Parameter = $"all-results-{Guid.NewGuid()}" }; + + var exception = await Assert.ThrowsAsync(async () => + { + await checker.CheckAsync("TestMultiRuleRacePolicy", context); + }); + + exception.Result.RuleResults.ShouldNotBeNull(); + exception.Result.RuleResults!.Count.ShouldBe(3); + } +} + /// /// Tests for Fix #1: Phase 2 in CheckAsync now checks the result of AcquireAsync. /// Uses a mock store that simulates a concurrent race condition: