From 4eba10f7a47779358bd39a7e7df796ab08b58b8f Mon Sep 17 00:00:00 2001 From: Jacob Moody Date: Wed, 6 May 2026 14:07:08 -0500 Subject: [PATCH] fix: per-key rate limit blocks at priority level, not whole queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When findAndWakeMatches found a rate-limited task it set a retry timer and returned immediately, stalling every task behind it regardless of fairness key. A high-priority key1 task being throttled would block a ready key2 task from dispatching until the timer fired. The fix moves the rate-limit check inside findMatch. The inner poller loop now saves the first compatible poller in `matched` rather than returning immediately. After finding a poller, the task's rate-limit delay is checked: if positive the task is skipped and its delay contributes to a running `minDelay`; if zero the match is returned. findMatch gains a third return value carrying minDelay back to findAndWakeMatches, which uses it to set the retry timer only when there are no ready matches at all — and only to the earliest time any blocked task (that actually has a compatible poller) becomes ready. --- service/matching/matcher_data.go | 41 +++++--- service/matching/matcher_data_test.go | 140 +++++++++++++++++++++++++- 2 files changed, 165 insertions(+), 16 deletions(-) diff --git a/service/matching/matcher_data.go b/service/matching/matcher_data.go index 45ebd31bcb2..978d14c18ba 100644 --- a/service/matching/matcher_data.go +++ b/service/matching/matcher_data.go @@ -415,13 +415,15 @@ func (d *matcherData) ReprocessTasks(pred func(*internalTask) bool) []*internalT return reprocess } -// findMatch should return the highest priority task+poller match even if the per-task rate -// limit doesn't allow the task to be matched yet. +// findMatch returns the highest-priority ready task+poller pair, skipping tasks that are +// rate-limited. If no task is ready, returns (nil, nil, minDelay) where minDelay is the +// shortest wait across all rate-limited tasks that have a matching poller available. // call with lock held // nolint:revive // will improve later -func (d *matcherData) findMatch(allowForwarding bool) (*internalTask, *waitingPoller) { +func (d *matcherData) findMatch(allowForwarding bool, now int64) (*internalTask, *waitingPoller, time.Duration) { // TODO(pri): optimize so it's not O(d*n) worst case // TODO(pri): this iterates over heap as slice, which isn't quite correct, but okay for now + var minDelay time.Duration for _, task := range d.tasks.heap { // disallow normal poll forwarding when allowForwarding is false, but allow the // "priority backlog poll forwarders". @@ -429,6 +431,7 @@ func (d *matcherData) findMatch(allowForwarding bool) (*internalTask, *waitingPo continue } + var matched *waitingPoller for _, poller := range d.pollers.heap { // can't match cases: if poller.queryOnly && !task.isQuery() && !task.isPollForwarder() { @@ -449,11 +452,23 @@ func (d *matcherData) findMatch(allowForwarding bool) (*internalTask, *waitingPo // their priority above "1". that's inaccurate but it's just a temporary situation. continue } + matched = poller + break + } + if matched == nil { + continue + } - return task, poller + delay := d.rateLimitManager.readyTimeForTask(task).delay(now) + if delay > 0 { + if minDelay == 0 || delay < minDelay { + minDelay = delay + } + continue } + return task, matched, 0 } - return nil, nil + return nil, nil, minDelay } // call with lock held @@ -498,20 +513,16 @@ func (d *matcherData) findAndWakeMatches() { for { // search for highest priority match - task, poller := d.findMatch(allowForwarding) + task, poller, minDelay := d.findMatch(allowForwarding, now) if task == nil || poller == nil { - // no more current matches, stop rate limit timer if was running - d.rateLimitTimer.unset() + if minDelay > 0 { + d.rateLimitTimer.set(d.timeSource, d.rematchAfterTimer, minDelay) + } else { + d.rateLimitTimer.unset() + } return } - // check ready time - delay := d.rateLimitManager.readyTimeForTask(task).delay(now) - d.rateLimitTimer.set(d.timeSource, d.rematchAfterTimer, delay) - if delay > 0 { - return // not ready yet, timer will call match later - } - // ready to signal match d.tasks.Remove(task) d.pollers.Remove(poller) diff --git a/service/matching/matcher_data_test.go b/service/matching/matcher_data_test.go index 09bcb579ec7..43c14d62b42 100644 --- a/service/matching/matcher_data_test.go +++ b/service/matching/matcher_data_test.go @@ -411,6 +411,144 @@ func (s *MatcherDataSuite) TestPerKeyRateLimit() { s.Less(elapsed, 20*time.Second) } +func (s *MatcherDataSuite) TestPerKeyRateLimitDoesNotBlockOtherKeys() { + // Set per-key rate limit: 1 RPS, no burst. After consuming one token for a key, + // the next token for that key is unavailable for 1 second. + s.md.rateLimitManager.SetFairnessKeyRateLimitDefaultForTesting(1.0, enumspb.RATE_LIMIT_SOURCE_API) + s.md.rateLimitManager.UpdatePerKeySimpleRateLimitWithBurstForTesting(0) + + // Dispatch one task from "key1" to exhaust its rate limit token. + key1Task1 := s.newBacklogTaskWithPriority(1, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(key1Task1)) + firstPoll := s.pollFakeTime(time.Second) + s.Require().NoError(firstPoll.ctxErr) + s.Equal(key1Task1, firstPoll.task) + + // key1 is now rate-limited for 1 second of fake time. + // key2 has no consumption yet and is immediately ready. + // + // Enqueue a new key1 task at priority 1 (higher urgency, comes first in the heap) + // and a key2 task at priority 2 (lower urgency, but not rate-limited). + key1Task2 := s.newBacklogTaskWithPriority(2, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + key2Task := s.newBacklogTaskWithPriority(3, 0, nil, &commonpb.Priority{PriorityKey: 2, FairnessKey: "key2"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(key1Task2)) + s.Require().NoError(s.md.EnqueueTaskNoWait(key2Task)) + + // The rate limit should block only key1, not the whole queue. + // key2Task should be dispatchable immediately even though the higher-priority + // key1Task2 is rate-limited and sits at the front of the task heap. + immediateMatch := s.pollImmediately(nil) + s.Require().NotNil(immediateMatch, "key2Task should be dispatched immediately; rate limit on key1 should not block the whole queue") + s.Equal(key2Task, immediateMatch.task) +} + +func (s *MatcherDataSuite) pollFakeTimeAsync(timeout time.Duration) <-chan *matchResult { + ch := make(chan *matchResult, 1) + go func() { ch <- s.pollFakeTime(timeout) }() + return ch +} + +// advanceFakeTimeUntil advances fake time in 10ms steps (with goroutine yields) until +// the provided channel has a result ready, then returns that result. +// Fails the test if the channel does not receive within 10 seconds of fake time. +func (s *MatcherDataSuite) advanceFakeTimeUntil(ch <-chan *matchResult) *matchResult { + const maxSteps = 1000 // 10s of fake time at 10ms/step + for range maxSteps { + gosched(3) + select { + case r := <-ch: + return r + default: + s.ts.Advance(10 * time.Millisecond) + } + } + s.FailNow("advanceFakeTimeUntil: timed out after 10s of fake time") + return nil +} + +// TestPerKeyRateLimitTimerFiresForSkippedKey verifies the end-to-end flow: +// a rate-limited high-priority task is skipped so a lower-priority ready task +// dispatches immediately, and then the rateLimitTimer fires and dispatches the +// originally-skipped task once its rate-limit interval elapses. +func (s *MatcherDataSuite) TestPerKeyRateLimitTimerFiresForSkippedKey() { + s.md.rateLimitManager.SetFairnessKeyRateLimitDefaultForTesting(1.0, enumspb.RATE_LIMIT_SOURCE_API) + s.md.rateLimitManager.UpdatePerKeySimpleRateLimitWithBurstForTesting(0) + + // Exhaust key1's token so it is rate-limited for 1s of fake time. + key1Seed := s.newBacklogTaskWithPriority(1, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(key1Seed)) + s.Require().NoError(s.pollFakeTime(time.Second).ctxErr) + + key1Blocked := s.newBacklogTaskWithPriority(2, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + key2Ready := s.newBacklogTaskWithPriority(3, 0, nil, &commonpb.Priority{PriorityKey: 2, FairnessKey: "key2"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(key1Blocked)) + s.Require().NoError(s.md.EnqueueTaskNoWait(key2Ready)) + + // key2 dispatches immediately despite rate-limited key1 sitting first in the heap. + got := s.pollImmediately(nil) + s.Require().NotNil(got) + s.Equal(key2Ready, got.task) + + // key1 must not dispatch before its interval elapses. + s.Nil(s.pollImmediately(nil), "key1 should still be rate-limited") + + // Start a poll, then advance fake time until the rateLimitTimer fires and dispatches key1. + result := s.advanceFakeTimeUntil(s.pollFakeTimeAsync(2 * time.Second)) + s.Require().NoError(result.ctxErr) + s.Equal(key1Blocked, result.task) +} + +// TestWholeQueueRateLimitBlocksAllTasks is a regression test: a whole-queue rate +// limit (not per-key) must block ALL tasks, including lower-priority ones from +// different fairness keys. findMatch must not allow a lower-priority task to +// bypass a limit that applies to every task in the queue. +func (s *MatcherDataSuite) TestWholeQueueRateLimitBlocksAllTasks() { + s.md.rateLimitManager.SetEffectiveRPSAndSourceForTesting(1.0, enumspb.RATE_LIMIT_SOURCE_API) + s.md.rateLimitManager.UpdateSimpleRateLimitWithBurstForTesting(0) + + // Consume the whole-queue token. + seed := s.newBacklogTaskWithPriority(1, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(seed)) + s.Require().NoError(s.pollFakeTime(time.Second).ctxErr) + + t1 := s.newBacklogTaskWithPriority(2, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + t2 := s.newBacklogTaskWithPriority(3, 0, nil, &commonpb.Priority{PriorityKey: 2, FairnessKey: "key2"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(t1)) + s.Require().NoError(s.md.EnqueueTaskNoWait(t2)) + + // Both tasks share the whole-queue token; neither should dispatch immediately. + s.Nil(s.pollImmediately(nil), "whole-queue limit must block all tasks, not just the high-priority key") + + // After the whole-queue token refills, t1 (priority 1) dispatches first. + r1 := s.advanceFakeTimeUntil(s.pollFakeTimeAsync(2 * time.Second)) + s.Require().NoError(r1.ctxErr) + s.Equal(t1, r1.task) + + r2 := s.advanceFakeTimeUntil(s.pollFakeTimeAsync(2 * time.Second)) + s.Require().NoError(r2.ctxErr) + s.Equal(t2, r2.task) +} + +// TestPerKeyRateLimitSameKeyTasksShareToken verifies that all tasks for the same +// fairness key share one token bucket. After one key1 task dispatches, a second +// key1 task must wait for the full interval even though no other key is involved. +func (s *MatcherDataSuite) TestPerKeyRateLimitSameKeyTasksShareToken() { + s.md.rateLimitManager.SetFairnessKeyRateLimitDefaultForTesting(1.0, enumspb.RATE_LIMIT_SOURCE_API) + s.md.rateLimitManager.UpdatePerKeySimpleRateLimitWithBurstForTesting(0) + + t1 := s.newBacklogTaskWithPriority(1, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(t1)) + s.Require().NoError(s.pollFakeTime(time.Second).ctxErr) + + t2 := s.newBacklogTaskWithPriority(2, 0, nil, &commonpb.Priority{PriorityKey: 1, FairnessKey: "key1"}) + s.Require().NoError(s.md.EnqueueTaskNoWait(t2)) + s.Nil(s.pollImmediately(nil), "second same-key task must wait; key1 token is exhausted") + + r := s.advanceFakeTimeUntil(s.pollFakeTimeAsync(2 * time.Second)) + s.Require().NoError(r.ctxErr) + s.Equal(t2, r.task) +} + func (s *MatcherDataSuite) TestOrder() { t1 := s.newBacklogTaskWithPriority(1, 0, nil, &commonpb.Priority{PriorityKey: 1}) t2 := s.newBacklogTaskWithPriority(2, 0, nil, &commonpb.Priority{PriorityKey: 2}) @@ -844,7 +982,7 @@ func (s *MatcherDataSuite) TestFindMatch() { // Call findMatch s.md.lock.Lock() - foundTask, foundPoller := s.md.findMatch(tc.allowForwarding) + foundTask, foundPoller, _ := s.md.findMatch(tc.allowForwarding, s.ts.Now().UnixNano()) s.md.lock.Unlock() if tc.shouldMatch {