From 8a9ec786e096b8586dd285075d4fd374f4d1551a Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:12:36 -0700 Subject: [PATCH 1/6] refactor: simplify the init-concurrency limiter to a single shared budget The limiter primitive added in #780 carried an optional per-environment gate and a matching INIT_PER_ENV_MAX_PERCENT config option. Drop per-environment fairness: the limiter now bounds concurrent admissions with one shared budget (a fixed number of held slots plus a bounded FIFO queue of waiters), and Acquire no longer takes an environment key. Also revise the INIT_SEND_TIMEOUT documentation and default to describe an absolute cap on how long one gated delivery may hold a slot. No runtime effect on its own: nothing acquires from the limiter yet. --- config/config.go | 13 ++-- config/config_concurrency_test.go | 8 +-- internal/concurrency/limiter.go | 96 +++++++--------------------- internal/concurrency/limiter_test.go | 45 +++++-------- 4 files changed, 48 insertions(+), 114 deletions(-) diff --git a/config/config.go b/config/config.go index e2d64ef9..f4860665 100644 --- a/config/config.go +++ b/config/config.go @@ -161,14 +161,11 @@ type ConcurrencyConfig struct { // of queueing them. MaxQueued ct.OptInt `conf:"INIT_MAX_QUEUED"` - // PerEnvMaxPercent limits the share of the budget that any single environment may - // use, as a percentage of MaxConcurrent plus MaxQueued. This keeps one busy - // environment from starving the others. A value of 0 applies no per-environment - // limit. - PerEnvMaxPercent ct.OptInt `conf:"INIT_PER_ENV_MAX_PERCENT"` - - // SendTimeout releases a delivery slot if a streaming initialization payload cannot - // make progress to its client within this duration. It defaults to 30s. + // SendTimeout is the absolute cap on how long a single gated delivery may hold a slot. A + // throughput floor (64 KB/s) closes a client that stalls or is slower than the floor well + // before this; the cap only backstops a client stuck right at the floor on a very large + // payload. If a delivery exceeds it, the connection is closed to reclaim the slot (and the + // SDK reconnects). It defaults to 2m. SendTimeout ct.OptDuration `conf:"INIT_SEND_TIMEOUT"` } diff --git a/config/config_concurrency_test.go b/config/config_concurrency_test.go index 38c04c6d..3dd7c864 100644 --- a/config/config_concurrency_test.go +++ b/config/config_concurrency_test.go @@ -13,16 +13,14 @@ import ( // the INIT_* environment variables. func TestConcurrencyConfigFromEnvironment(t *testing.T) { withEnvironment(map[string]string{ - "INIT_MAX_CONCURRENT": "200", - "INIT_MAX_QUEUED": "1000", - "INIT_PER_ENV_MAX_PERCENT": "40", - "INIT_SEND_TIMEOUT": "15s", + "INIT_MAX_CONCURRENT": "200", + "INIT_MAX_QUEUED": "1000", + "INIT_SEND_TIMEOUT": "15s", }, func() { var c Config require.NoError(t, LoadConfigFromEnvironment(&c, slog.Default())) assert.Equal(t, 200, c.Concurrency.MaxConcurrent.GetOrElse(-1)) assert.Equal(t, 1000, c.Concurrency.MaxQueued.GetOrElse(-1)) - assert.Equal(t, 40, c.Concurrency.PerEnvMaxPercent.GetOrElse(-1)) assert.Equal(t, 15*time.Second, c.Concurrency.SendTimeout.GetOrElse(0)) }) } diff --git a/internal/concurrency/limiter.go b/internal/concurrency/limiter.go index e0ae0099..35dde192 100644 --- a/internal/concurrency/limiter.go +++ b/internal/concurrency/limiter.go @@ -1,8 +1,7 @@ // Package concurrency provides an admission limiter that bounds how much concurrent // work a burst of requests or connections can impose on Relay. It uses two limits: a // maximum number of slots held at once, and a bounded FIFO queue of callers waiting for -// a slot. An optional per-environment gate keeps one environment from using the whole -// budget. +// a slot. package concurrency import ( @@ -19,13 +18,9 @@ type Params struct { // MaxQueued is the number of callers that may wait in FIFO order for a slot once all // slots are held. A value of 0 rejects callers immediately instead of waiting. MaxQueued int - // PerEnvMax limits how many of a single environment's callers may participate, - // counting both held and waiting, at once. A value of 0 applies no per-environment - // limit. The per-environment gate never blocks; it only rejects. - PerEnvMax int } -// Stats is a point-in-time snapshot of a Limiter's counters, for logging/metrics. +// Stats is a point-in-time snapshot of a Limiter's counters for logging and metrics. type Stats struct { Enabled bool MaxConcurrent int @@ -36,30 +31,23 @@ type Stats struct { Rejected int64 } -// Limiter bounds concurrency with two limits plus an optional per-environment gate that -// never blocks and only rejects. The zero value is not usable; construct one with New. +// Limiter bounds concurrency with two limits: a maximum number of slots held at once and +// a bounded FIFO queue of waiters. The zero value is not usable; construct one with New. type Limiter struct { name string enabled bool - tokens chan struct{} // holds MaxConcurrent slots; receive to acquire, send to release + tokens chan struct{} // holds the free slots; receive to acquire one, send to release one maxQueued int64 waiting int64 shutdown chan struct{} closeOnce sync.Once - perEnvMax int - perEnv sync.Map // envKey -> *envGate - admitted atomic.Int64 rejected atomic.Int64 } -type envGate struct { - slots chan struct{} // cap = perEnvMax; try-receive to enter, send to leave -} - -// New builds a Limiter. name is used only for logging/metrics identification. +// New builds a Limiter. name identifies the limiter in logs and metrics. func New(name string, p Params) *Limiter { l := &Limiter{name: name} if p.MaxConcurrent <= 0 { @@ -72,9 +60,6 @@ func New(name string, p Params) *Limiter { } l.maxQueued = int64(p.MaxQueued) l.shutdown = make(chan struct{}) - if p.PerEnvMax > 0 { - l.perEnvMax = p.PerEnvMax - } return l } @@ -84,84 +69,51 @@ func (l *Limiter) Name() string { return l.name } // Enabled reports whether the limiter is enforcing a limit. func (l *Limiter) Enabled() bool { return l != nil && l.enabled } -// Acquire attempts to admit one unit of work for the given environment key. -// On success it returns a release func (call exactly once) and ok=true. If the -// per-env gate or the global backlog is full, or ctx is cancelled, or the -// limiter is shut down, it returns a no-op release and ok=false. A disabled or -// nil limiter always admits immediately. -func (l *Limiter) Acquire(ctx context.Context, envKey string) (release func(), ok bool) { +// Acquire attempts to admit one unit of work. On success it returns a release function, +// which the caller must call exactly once, and ok is true. It returns a no-op release and +// ok=false if the queue is full, ctx is cancelled, or the limiter is shut down. A disabled +// or nil limiter always admits immediately. +func (l *Limiter) Acquire(ctx context.Context) (release func(), ok bool) { if !l.Enabled() { return func() {}, true } - // Per-environment gate. It never blocks; it rejects when the environment is over its share. - var releaseEnv func() - if l.perEnvMax > 0 { - g := l.gateFor(envKey) - select { - case g.slots <- struct{}{}: - releaseEnv = func() { <-g.slots } - default: - l.rejected.Add(1) - return func() {}, false - } - } - - reject := func() (func(), bool) { - if releaseEnv != nil { - releaseEnv() - } - l.rejected.Add(1) - return func() {}, false - } - - // Fast path: a token is immediately available. + // Take a free slot if one is available. select { case <-l.tokens: - return l.admit(releaseEnv), true + return l.admit(), true default: } - // No token free: enter the bounded FIFO backlog, or reject. + // No free slot. Join the bounded queue, or reject if it is full. if atomic.AddInt64(&l.waiting, 1) > l.maxQueued { atomic.AddInt64(&l.waiting, -1) - return reject() + l.rejected.Add(1) + return func() {}, false } defer atomic.AddInt64(&l.waiting, -1) select { case <-l.tokens: - return l.admit(releaseEnv), true + return l.admit(), true case <-ctx.Done(): - return reject() + l.rejected.Add(1) + return func() {}, false case <-l.shutdown: - return reject() + l.rejected.Add(1) + return func() {}, false } } -func (l *Limiter) admit(releaseEnv func()) func() { +func (l *Limiter) admit() func() { l.admitted.Add(1) var once sync.Once return func() { - once.Do(func() { - l.tokens <- struct{}{} - if releaseEnv != nil { - releaseEnv() - } - }) - } -} - -func (l *Limiter) gateFor(envKey string) *envGate { - if g, ok := l.perEnv.Load(envKey); ok { - return g.(*envGate) + once.Do(func() { l.tokens <- struct{}{} }) } - g := &envGate{slots: make(chan struct{}, l.perEnvMax)} - actual, _ := l.perEnv.LoadOrStore(envKey, g) - return actual.(*envGate) } -// Close unblocks all waiters (they receive ok=false). Idempotent. +// Close unblocks all waiters, which then receive ok=false. It may be called more than once. func (l *Limiter) Close() { if !l.Enabled() { return diff --git a/internal/concurrency/limiter_test.go b/internal/concurrency/limiter_test.go index 650aafd1..42240c68 100644 --- a/internal/concurrency/limiter_test.go +++ b/internal/concurrency/limiter_test.go @@ -13,7 +13,7 @@ func TestDisabledLimiterAlwaysAdmits(t *testing.T) { t.Fatal("expected disabled") } for i := 0; i < 100; i++ { - release, ok := l.Acquire(context.Background(), "env") + release, ok := l.Acquire(context.Background()) if !ok { t.Fatal("disabled limiter must always admit") } @@ -23,7 +23,7 @@ func TestDisabledLimiterAlwaysAdmits(t *testing.T) { func TestNilLimiterAdmits(t *testing.T) { var l *Limiter - release, ok := l.Acquire(context.Background(), "env") + release, ok := l.Acquire(context.Background()) if !ok { t.Fatal("nil limiter must admit") } @@ -32,9 +32,9 @@ func TestNilLimiterAdmits(t *testing.T) { func TestRejectWhenNoBacklog(t *testing.T) { l := New("t", Params{MaxConcurrent: 2, MaxQueued: 0}) - r1, ok1 := l.Acquire(context.Background(), "e") - r2, ok2 := l.Acquire(context.Background(), "e") - _, ok3 := l.Acquire(context.Background(), "e") + r1, ok1 := l.Acquire(context.Background()) + r2, ok2 := l.Acquire(context.Background()) + _, ok3 := l.Acquire(context.Background()) if !ok1 || !ok2 { t.Fatal("first two should be admitted") } @@ -50,13 +50,13 @@ func TestRejectWhenNoBacklog(t *testing.T) { func TestQueueThenAdmitOnRelease(t *testing.T) { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) - r1, ok1 := l.Acquire(context.Background(), "e") + r1, ok1 := l.Acquire(context.Background()) if !ok1 { t.Fatal("first should be admitted") } admitted := make(chan struct{}) go func() { - r2, ok2 := l.Acquire(context.Background(), "e") // must queue, then admit when r1 releases + r2, ok2 := l.Acquire(context.Background()) // must queue, then admit when r1 releases if ok2 { close(admitted) r2() @@ -79,36 +79,23 @@ func TestQueueThenAdmitOnRelease(t *testing.T) { func TestBacklogFullRejects(t *testing.T) { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) - r1, _ := l.Acquire(context.Background(), "e") // holds the only token + r1, _ := l.Acquire(context.Background()) // holds the only token defer r1() - go l.Acquire(context.Background(), "e") // fills the single backlog slot + go l.Acquire(context.Background()) // fills the single backlog slot time.Sleep(50 * time.Millisecond) - if _, ok := l.Acquire(context.Background(), "e"); ok { + if _, ok := l.Acquire(context.Background()); ok { t.Fatal("expected rejection when backlog is full") } } -func TestPerEnvGateIsolatesEnvironments(t *testing.T) { - // Global room for 10, but each env may hold at most 1 (participant cap). - l := New("t", Params{MaxConcurrent: 10, MaxQueued: 10, PerEnvMax: 1}) - rA, okA := l.Acquire(context.Background(), "A") - _, okA2 := l.Acquire(context.Background(), "A") // second A rejected by per-env gate - rB, okB := l.Acquire(context.Background(), "B") // different env still admitted - if !okA || okA2 || !okB { - t.Fatalf("per-env gate failed: okA=%v okA2=%v okB=%v", okA, okA2, okB) - } - rA() - rB() -} - func TestContextCancelUnblocksWaiter(t *testing.T) { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 5}) - r1, _ := l.Acquire(context.Background(), "e") + r1, _ := l.Acquire(context.Background()) defer r1() ctx, cancel := context.WithCancel(context.Background()) done := make(chan bool) go func() { - _, ok := l.Acquire(ctx, "e") + _, ok := l.Acquire(ctx) done <- ok }() time.Sleep(50 * time.Millisecond) @@ -125,15 +112,15 @@ func TestContextCancelUnblocksWaiter(t *testing.T) { func TestReleaseIsIdempotent(t *testing.T) { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 0}) - r, _ := l.Acquire(context.Background(), "e") + r, _ := l.Acquire(context.Background()) r() r() // must not release a second token // Two acquires should now succeed sequentially, proving only one token exists. - r1, ok1 := l.Acquire(context.Background(), "e") + r1, ok1 := l.Acquire(context.Background()) if !ok1 { t.Fatal("expected admit after release") } - if _, ok2 := l.Acquire(context.Background(), "e"); ok2 { + if _, ok2 := l.Acquire(context.Background()); ok2 { t.Fatal("double release leaked a token") } r1() @@ -149,7 +136,7 @@ func TestConcurrentAcquireBoundsHeld(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - release, ok := l.Acquire(context.Background(), "e") + release, ok := l.Acquire(context.Background()) if !ok { return } From 16ad622f5982cdc48651972cddd25a91ccb1404d Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:25:03 -0700 Subject: [PATCH 2/6] make the limiter honor its Acquire and Close contracts, and test them Review of the primitive found three places where the documented contract and the code disagreed, plus untested paths. All are limiter-local: - Close is now an admission barrier. The doc promised that a shut-down limiter rejects, but the free-slot fast path never consulted the shutdown channel, so callers arriving after Close were admitted whenever a slot was free. Shutdown wiring will assume "Close means no new admissions". - An already-cancelled context is rejected up front instead of being handed a free slot: a dead request would only waste the slot ahead of live ones, and the doc already claimed this behavior. - The waiting count is decremented inside each select arm rather than by a deferred decrement spanning admission, so a caller that has been handed a slot no longer counts against the queue bound (a spanning decrement caused spurious queue-full rejections under exactly the herd load the limiter targets, and inflated Stats.Waiting). - "FIFO" softened to "approximate arrival order": a caller counts against the queue before it parks, so strict ordering was never guaranteed. - Name() is nil-safe like every other method; waiting is an atomic.Int64 like its siblings. Tests now cover Close (admission barrier, unblocks waiters with ok=false, idempotent, safe release afterwards), the cancelled-context paths and their rejection counts, Stats fields, and the waiting-count semantic; each new guard was verified to fail with its guard removed. Queue-state setup polls Stats instead of sleeping, and the queue-full test no longer leaks its waiter. --- internal/concurrency/limiter.go | 69 ++++++++---- internal/concurrency/limiter_test.go | 160 ++++++++++++++++++++++++++- 2 files changed, 202 insertions(+), 27 deletions(-) diff --git a/internal/concurrency/limiter.go b/internal/concurrency/limiter.go index 35dde192..87ac6f10 100644 --- a/internal/concurrency/limiter.go +++ b/internal/concurrency/limiter.go @@ -1,7 +1,7 @@ // Package concurrency provides an admission limiter that bounds how much concurrent // work a burst of requests or connections can impose on Relay. It uses two limits: a -// maximum number of slots held at once, and a bounded FIFO queue of callers waiting for -// a slot. +// maximum number of slots held at once, and a bounded queue of callers waiting for a +// slot in approximate arrival order. package concurrency import ( @@ -15,8 +15,9 @@ import ( type Params struct { // MaxConcurrent is the number of slots that may be held at once. MaxConcurrent int - // MaxQueued is the number of callers that may wait in FIFO order for a slot once all - // slots are held. A value of 0 rejects callers immediately instead of waiting. + // MaxQueued is the number of callers that may wait for a slot once all slots are + // held. Waiters are served in approximate arrival order. A value of 0 rejects + // callers immediately instead of waiting. MaxQueued int } @@ -32,14 +33,15 @@ type Stats struct { } // Limiter bounds concurrency with two limits: a maximum number of slots held at once and -// a bounded FIFO queue of waiters. The zero value is not usable; construct one with New. +// a bounded queue of waiters served in approximate arrival order. The zero value is not +// usable; construct one with New. type Limiter struct { name string enabled bool tokens chan struct{} // holds the free slots; receive to acquire one, send to release one maxQueued int64 - waiting int64 + waiting atomic.Int64 shutdown chan struct{} closeOnce sync.Once @@ -63,21 +65,38 @@ func New(name string, p Params) *Limiter { return l } -// Name returns the limiter's identifier. -func (l *Limiter) Name() string { return l.name } +// Name returns the limiter's identifier, or an empty string for a nil limiter. +func (l *Limiter) Name() string { + if l == nil { + return "" + } + return l.name +} // Enabled reports whether the limiter is enforcing a limit. func (l *Limiter) Enabled() bool { return l != nil && l.enabled } // Acquire attempts to admit one unit of work. On success it returns a release function, // which the caller must call exactly once, and ok is true. It returns a no-op release and -// ok=false if the queue is full, ctx is cancelled, or the limiter is shut down. A disabled -// or nil limiter always admits immediately. +// ok=false if the queue is full, ctx is already cancelled or becomes cancelled while +// waiting, or the limiter is shut down. A disabled or nil limiter always admits +// immediately. func (l *Limiter) Acquire(ctx context.Context) (release func(), ok bool) { if !l.Enabled() { return func() {}, true } + // Refuse work that arrives after shutdown or is already abandoned, even when a slot + // is free: Close is an admission barrier, and a dead request would only waste the slot. + select { + case <-l.shutdown: + return l.reject() + default: + } + if ctx.Err() != nil { + return l.reject() + } + // Take a free slot if one is available. select { case <-l.tokens: @@ -86,22 +105,23 @@ func (l *Limiter) Acquire(ctx context.Context) (release func(), ok bool) { } // No free slot. Join the bounded queue, or reject if it is full. - if atomic.AddInt64(&l.waiting, 1) > l.maxQueued { - atomic.AddInt64(&l.waiting, -1) - l.rejected.Add(1) - return func() {}, false + if l.waiting.Add(1) > l.maxQueued { + l.waiting.Add(-1) + return l.reject() } - defer atomic.AddInt64(&l.waiting, -1) + // The waiting count is decremented inside each arm, before returning, so a caller + // that has already been handed a slot is not still counted against the queue bound. select { case <-l.tokens: + l.waiting.Add(-1) return l.admit(), true case <-ctx.Done(): - l.rejected.Add(1) - return func() {}, false + l.waiting.Add(-1) + return l.reject() case <-l.shutdown: - l.rejected.Add(1) - return func() {}, false + l.waiting.Add(-1) + return l.reject() } } @@ -113,7 +133,14 @@ func (l *Limiter) admit() func() { } } -// Close unblocks all waiters, which then receive ok=false. It may be called more than once. +func (l *Limiter) reject() (func(), bool) { + l.rejected.Add(1) + return func() {}, false +} + +// Close stops admissions: callers that arrive afterwards are rejected, and all waiters +// are unblocked with ok=false. Releases of already-held slots remain safe. It may be +// called more than once. func (l *Limiter) Close() { if !l.Enabled() { return @@ -131,7 +158,7 @@ func (l *Limiter) Stats() Stats { MaxConcurrent: cap(l.tokens), MaxQueued: int(l.maxQueued), Held: cap(l.tokens) - len(l.tokens), - Waiting: int(atomic.LoadInt64(&l.waiting)), + Waiting: int(l.waiting.Load()), Admitted: l.admitted.Load(), Rejected: l.rejected.Load(), } diff --git a/internal/concurrency/limiter_test.go b/internal/concurrency/limiter_test.go index 42240c68..7362aa80 100644 --- a/internal/concurrency/limiter_test.go +++ b/internal/concurrency/limiter_test.go @@ -2,11 +2,25 @@ package concurrency import ( "context" + "runtime" "sync" "testing" "time" ) +// waitForWaiting polls until the limiter reports the given queue depth, so tests can +// establish "a caller is parked" without a fixed sleep. +func waitForWaiting(t *testing.T, l *Limiter, n int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for l.Stats().Waiting != n { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for Waiting==%d (stats: %+v)", n, l.Stats()) + } + runtime.Gosched() + } +} + func TestDisabledLimiterAlwaysAdmits(t *testing.T) { l := New("t", Params{MaxConcurrent: 0}) if l.Enabled() { @@ -28,6 +42,10 @@ func TestNilLimiterAdmits(t *testing.T) { t.Fatal("nil limiter must admit") } release() + if l.Name() != "" { + t.Fatal("nil limiter Name should be empty") + } + l.Close() // must not panic } func TestRejectWhenNoBacklog(t *testing.T) { @@ -62,8 +80,7 @@ func TestQueueThenAdmitOnRelease(t *testing.T) { r2() } }() - // Give the goroutine time to enter the backlog. - time.Sleep(50 * time.Millisecond) + waitForWaiting(t, l, 1) select { case <-admitted: t.Fatal("queued caller admitted before release") @@ -80,12 +97,53 @@ func TestQueueThenAdmitOnRelease(t *testing.T) { func TestBacklogFullRejects(t *testing.T) { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) r1, _ := l.Acquire(context.Background()) // holds the only token - defer r1() - go l.Acquire(context.Background()) // fills the single backlog slot - time.Sleep(50 * time.Millisecond) + waiterCtx, cancelWaiter := context.WithCancel(context.Background()) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + if r, ok := l.Acquire(waiterCtx); ok { // fills the single backlog slot + r() + } + }() + waitForWaiting(t, l, 1) if _, ok := l.Acquire(context.Background()); ok { t.Fatal("expected rejection when backlog is full") } + // Unwind without leaking the waiter or its token. + cancelWaiter() + wg.Wait() + r1() +} + +func TestWaitingCountsOnlyParkedCallers(t *testing.T) { + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + r1, _ := l.Acquire(context.Background()) + admitted := make(chan func()) + go func() { + if r, ok := l.Acquire(context.Background()); ok { + admitted <- r + } + }() + waitForWaiting(t, l, 1) + r1() // hand the slot to the waiter + r2 := <-admitted + // An admitted caller must no longer count against the queue bound: with the single + // backlog slot logically empty, a new caller must be able to queue rather than be shed. + waitForWaiting(t, l, 0) + queued := make(chan bool) + go func() { + r, ok := l.Acquire(context.Background()) + if ok { + r() + } + queued <- ok + }() + waitForWaiting(t, l, 1) + r2() + if ok := <-queued; !ok { + t.Fatal("caller was shed while the queue was logically empty") + } } func TestContextCancelUnblocksWaiter(t *testing.T) { @@ -98,7 +156,8 @@ func TestContextCancelUnblocksWaiter(t *testing.T) { _, ok := l.Acquire(ctx) done <- ok }() - time.Sleep(50 * time.Millisecond) + waitForWaiting(t, l, 1) + before := l.Stats().Rejected cancel() select { case ok := <-done: @@ -108,6 +167,62 @@ func TestContextCancelUnblocksWaiter(t *testing.T) { case <-time.After(time.Second): t.Fatal("waiter not unblocked by context cancel") } + if got := l.Stats().Rejected; got != before+1 { + t.Fatalf("cancelled waiter must count as rejected: before=%d after=%d", before, got) + } +} + +func TestAlreadyCancelledContextIsRejected(t *testing.T) { + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + // Even with a slot free, an abandoned request must not consume it. + if _, ok := l.Acquire(ctx); ok { + t.Fatal("an already-cancelled request must be rejected") + } + if s := l.Stats(); s.Held != 0 || s.Rejected != 1 { + t.Fatalf("unexpected stats: %+v", s) + } +} + +func TestCloseIsAnAdmissionBarrier(t *testing.T) { + l := New("t", Params{MaxConcurrent: 2, MaxQueued: 2}) + l.Close() + // All slots are free, but a closed limiter must not admit new work. + for i := 0; i < 4; i++ { + if _, ok := l.Acquire(context.Background()); ok { + t.Fatal("Acquire after Close must be rejected") + } + } + if s := l.Stats(); s.Rejected != 4 || s.Admitted != 0 { + t.Fatalf("unexpected stats: %+v", s) + } +} + +func TestCloseUnblocksWaiters(t *testing.T) { + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 4}) + r1, _ := l.Acquire(context.Background()) + results := make(chan bool, 3) + for i := 0; i < 3; i++ { + go func() { + _, ok := l.Acquire(context.Background()) + results <- ok + }() + } + waitForWaiting(t, l, 3) + l.Close() + for i := 0; i < 3; i++ { + select { + case ok := <-results: + if ok { + t.Fatal("waiter admitted by Close") + } + case <-time.After(time.Second): + t.Fatal("waiter not unblocked by Close") + } + } + r1() // releasing a held slot after Close must not panic or block + l.Close() // idempotent } func TestReleaseIsIdempotent(t *testing.T) { @@ -126,6 +241,39 @@ func TestReleaseIsIdempotent(t *testing.T) { r1() } +func TestStatsSnapshot(t *testing.T) { + l := New("gate", Params{MaxConcurrent: 2, MaxQueued: 3}) + if l.Name() != "gate" { + t.Fatalf("unexpected name %q", l.Name()) + } + r1, _ := l.Acquire(context.Background()) + r2, _ := l.Acquire(context.Background()) + unblocked := make(chan struct{}) + go func() { + defer close(unblocked) + if r, ok := l.Acquire(context.Background()); ok { + r() + } + }() + waitForWaiting(t, l, 1) + + s := l.Stats() + if !s.Enabled || s.MaxConcurrent != 2 || s.MaxQueued != 3 || s.Held != 2 || s.Waiting != 1 || s.Admitted != 2 { + t.Fatalf("unexpected stats: %+v", s) + } + l.Close() // unblock the waiter + <-unblocked + r1() + r2() + if s := l.Stats(); s.Held != 0 { + t.Fatalf("expected all slots free after release: %+v", s) + } + + if ds := (&Limiter{}).Stats(); ds.Enabled { + t.Fatal("disabled limiter must report Enabled=false") + } +} + func TestConcurrentAcquireBoundsHeld(t *testing.T) { const maxC = 4 l := New("t", Params{MaxConcurrent: maxC, MaxQueued: 1000}) From 8633f0fa023c831eb44139bb34685fc94c0a1d43 Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:58:31 -0700 Subject: [PATCH 3/6] count budget occupancy instead of queued callers The previous fix for spurious queue-full rejections moved the waiting-counter decrement between lines inside the woken goroutine, but the window it was meant to close is dominated by that goroutine's wake-up latency: a released slot is handed directly to a parked waiter, who keeps counting against the queue until the scheduler runs it, so every slot turnover shed one admissible caller under exactly the burst load the queue exists for. Measured rejection rates were unchanged by that fix. Replace the queued-callers counter with a single occupancy counter covering slot holders and waiters together, bounded by MaxConcurrent+MaxQueued. A caller reserves its unit once on entry and keeps it from queue to held to released, so the queue-to-held transition no longer passes through a state where the waiter double-counts. This also removes the phantom queue position a caller being rejected briefly held (which could shed an unrelated admissible caller), and makes Stats.Waiting derived and honest instead of over-reading. The release frees occupancy before returning the slot. Close now also covers the release-that-races-Close case: a waiter handed a slot re-checks shutdown and returns the slot instead of keeping it, so once Close returns, a subsequently released slot cannot admit a parked waiter. The Close and Acquire docs state the contract in those terms. Replace two tests that passed unchanged on the code they were written to guard. The new ones fail on both prior revisions: an exact-fit-budget loop (two live callers against a budget of two, so any rejection is spurious) that sheds on most turnovers under either old counter, and a Close-then-release loop asserting the waiter is never admitted. Add guards for occupancy leaks on the cancel, shutdown, and close-unblock paths, and a white-box test for the slot-won-concurrently-with-Close re-check. The release's free-before-send ordering is belt-and-braces; its inversion is not observable at test granularity. Trim the SendTimeout godoc to what exists in this package, pointing the default and throughput-floor behavior at the wiring change that implements them, so published godoc does not describe machinery this branch lacks. --- config/config.go | 12 +-- internal/concurrency/limiter.go | 91 ++++++++++++------- internal/concurrency/limiter_test.go | 128 +++++++++++++++++++++++---- 3 files changed, 177 insertions(+), 54 deletions(-) diff --git a/config/config.go b/config/config.go index f4860665..393621ca 100644 --- a/config/config.go +++ b/config/config.go @@ -161,11 +161,13 @@ type ConcurrencyConfig struct { // of queueing them. MaxQueued ct.OptInt `conf:"INIT_MAX_QUEUED"` - // SendTimeout is the absolute cap on how long a single gated delivery may hold a slot. A - // throughput floor (64 KB/s) closes a client that stalls or is slower than the floor well - // before this; the cap only backstops a client stuck right at the floor on a very large - // payload. If a delivery exceeds it, the connection is closed to reclaim the slot (and the - // SDK reconnects). It defaults to 2m. + // SendTimeout is the longest a single initialization delivery may hold a concurrency + // slot before its connection is closed to reclaim the slot (the SDK then reconnects). + // + // TODO(wiring): enforcement is not in this package. The delivery path that consumes + // this option -- applying its default and a per-write throughput floor that cuts a + // stalled client well before this cap -- lands with the code that wires the limiter + // into the endpoints. SendTimeout ct.OptDuration `conf:"INIT_SEND_TIMEOUT"` } diff --git a/internal/concurrency/limiter.go b/internal/concurrency/limiter.go index 87ac6f10..0fe8f40c 100644 --- a/internal/concurrency/limiter.go +++ b/internal/concurrency/limiter.go @@ -15,13 +15,16 @@ import ( type Params struct { // MaxConcurrent is the number of slots that may be held at once. MaxConcurrent int - // MaxQueued is the number of callers that may wait for a slot once all slots are - // held. Waiters are served in approximate arrival order. A value of 0 rejects - // callers immediately instead of waiting. + // MaxQueued is how many callers may wait for a slot once every slot is held, served + // in approximate arrival order. A value of 0 or less adds no waiting capacity: when + // every slot is held a caller is rejected immediately, though a caller arriving just + // as a slot is being released may briefly wait to take it. MaxQueued int } // Stats is a point-in-time snapshot of a Limiter's counters for logging and metrics. +// Held is exact; Waiting is derived from the occupancy counter and may transiently count +// a caller whose slot handoff is in progress. type Stats struct { Enabled bool MaxConcurrent int @@ -40,10 +43,18 @@ type Limiter struct { enabled bool tokens chan struct{} // holds the free slots; receive to acquire one, send to release one - maxQueued int64 - waiting atomic.Int64 - shutdown chan struct{} - closeOnce sync.Once + maxQueued int + // inFlight counts the callers occupying the budget: slot holders plus queued waiters, + // bounded by MaxConcurrent+MaxQueued. A caller reserves its unit once on entry and + // keeps it from queue to held to released, so handing a released slot to a parked + // waiter never leaves a moment where the waiter still counts against the queue while + // the slot is already spoken for. A separate queue counter has exactly that moment (it + // lasts until the woken goroutine is scheduled) and sheds an admissible caller on + // every slot turnover under burst load. + inFlight atomic.Int64 + maxOccupancy int64 + shutdown chan struct{} + closeOnce sync.Once admitted atomic.Int64 rejected atomic.Int64 @@ -60,7 +71,8 @@ func New(name string, p Params) *Limiter { for i := 0; i < p.MaxConcurrent; i++ { l.tokens <- struct{}{} } - l.maxQueued = int64(p.MaxQueued) + l.maxQueued = max(p.MaxQueued, 0) + l.maxOccupancy = int64(p.MaxConcurrent) + int64(l.maxQueued) l.shutdown = make(chan struct{}) return l } @@ -97,40 +109,55 @@ func (l *Limiter) Acquire(ctx context.Context) (release func(), ok bool) { return l.reject() } + // Reserve one unit of the budget's total capacity (slots plus queue). Rejecting on + // overflow here is what bounds the queue. + if l.inFlight.Add(1) > l.maxOccupancy { + l.inFlight.Add(-1) + return l.reject() + } + // Take a free slot if one is available. select { case <-l.tokens: - return l.admit(), true + return l.admit() default: } - // No free slot. Join the bounded queue, or reject if it is full. - if l.waiting.Add(1) > l.maxQueued { - l.waiting.Add(-1) - return l.reject() - } - - // The waiting count is decremented inside each arm, before returning, so a caller - // that has already been handed a slot is not still counted against the queue bound. + // Every slot is held. Wait for one; the occupancy reservation above bounds how many + // callers may wait here. select { case <-l.tokens: - l.waiting.Add(-1) - return l.admit(), true + return l.admit() case <-ctx.Done(): - l.waiting.Add(-1) + l.inFlight.Add(-1) return l.reject() case <-l.shutdown: - l.waiting.Add(-1) + l.inFlight.Add(-1) return l.reject() } } -func (l *Limiter) admit() func() { +// admit completes an admission after a slot has been received -- unless shutdown has +// landed in the meantime, in which case the slot is returned and the caller is rejected, +// so a slot released concurrently with Close cannot smuggle a waiter past it. +func (l *Limiter) admit() (func(), bool) { + select { + case <-l.shutdown: + l.tokens <- struct{}{} // never blocks: this slot's buffer capacity is unoccupied + l.inFlight.Add(-1) + return l.reject() + default: + } l.admitted.Add(1) var once sync.Once return func() { - once.Do(func() { l.tokens <- struct{}{} }) - } + once.Do(func() { + // Free the occupancy before returning the slot, so the queue capacity this + // release re-opens is visible to arrivals no later than the slot itself. + l.inFlight.Add(-1) + l.tokens <- struct{}{} + }) + }, true } func (l *Limiter) reject() (func(), bool) { @@ -138,9 +165,11 @@ func (l *Limiter) reject() (func(), bool) { return func() {}, false } -// Close stops admissions: callers that arrive afterwards are rejected, and all waiters -// are unblocked with ok=false. Releases of already-held slots remain safe. It may be -// called more than once. +// Close stops admissions: once it returns, callers entering Acquire are rejected, and +// parked waiters are unblocked and rejected. A waiter handed a slot concurrently with +// Close re-checks shutdown and returns the slot instead of keeping it; only an admission +// that fully completed before shutdown landed keeps its slot. Releases of held slots +// remain safe afterwards. It may be called more than once. func (l *Limiter) Close() { if !l.Enabled() { return @@ -153,12 +182,14 @@ func (l *Limiter) Stats() Stats { if !l.Enabled() { return Stats{Enabled: false} } + held := cap(l.tokens) - len(l.tokens) + waiting := max(int(l.inFlight.Load())-held, 0) return Stats{ Enabled: true, MaxConcurrent: cap(l.tokens), - MaxQueued: int(l.maxQueued), - Held: cap(l.tokens) - len(l.tokens), - Waiting: int(l.waiting.Load()), + MaxQueued: l.maxQueued, + Held: held, + Waiting: waiting, Admitted: l.admitted.Load(), Rejected: l.rejected.Load(), } diff --git a/internal/concurrency/limiter_test.go b/internal/concurrency/limiter_test.go index 7362aa80..1c85258b 100644 --- a/internal/concurrency/limiter_test.go +++ b/internal/concurrency/limiter_test.go @@ -8,8 +8,10 @@ import ( "time" ) -// waitForWaiting polls until the limiter reports the given queue depth, so tests can -// establish "a caller is parked" without a fixed sleep. +// waitForWaiting polls until Stats reports n waiting callers. The count is a reservation +// made on entering the queued path, so this establishes that the callers have committed to +// waiting (they reach the select moments later); it is not a strict parked-at-the-select +// barrier. func waitForWaiting(t *testing.T, l *Limiter, n int) { t.Helper() deadline := time.Now().Add(2 * time.Second) @@ -116,33 +118,75 @@ func TestBacklogFullRejects(t *testing.T) { r1() } -func TestWaitingCountsOnlyParkedCallers(t *testing.T) { +func TestSlotTurnoverDoesNotShedExactFitLoad(t *testing.T) { + // The budget is MaxConcurrent+MaxQueued = 2 and there are never more than two live + // callers, so every rejection is spurious by construction. Each iteration hands the + // held slot to a parked waiter and immediately re-acquires: the vacated queue capacity + // must be available even while the woken waiter is still being scheduled. (A separate + // queued-callers counter fails here on most iterations, because the woken waiter keeps + // counting against the queue until it runs.) + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + r, ok := l.Acquire(context.Background()) + if !ok { + t.Fatal("initial acquire") + } + var wg sync.WaitGroup + for i := 0; i < 400; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if r2, ok := l.Acquire(context.Background()); ok { + r2() + } + }() + waitForWaiting(t, l, 1) + r() // hand the slot to the waiter + var reacquired bool + r, reacquired = l.Acquire(context.Background()) + if !reacquired { + t.Fatalf("iteration %d: caller shed while the budget had room", i) + } + } + r() + wg.Wait() +} + +func TestCancelledWaitersFreeTheirQueueCapacity(t *testing.T) { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) r1, _ := l.Acquire(context.Background()) - admitted := make(chan func()) - go func() { - if r, ok := l.Acquire(context.Background()); ok { - admitted <- r + // Repeatedly park a waiter and cancel it: each cancellation must return its budget + // reservation, or the queue capacity erodes until admissible callers are shed. + for i := 0; i < 5; i++ { + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan bool) + go func() { + _, ok := l.Acquire(ctx) + done <- ok + }() + waitForWaiting(t, l, 1) + cancel() + if ok := <-done; ok { + t.Fatal("cancelled waiter should have been rejected") } - }() - waitForWaiting(t, l, 1) - r1() // hand the slot to the waiter - r2 := <-admitted - // An admitted caller must no longer count against the queue bound: with the single - // backlog slot logically empty, a new caller must be able to queue rather than be shed. - waitForWaiting(t, l, 0) - queued := make(chan bool) + } + // Every cancellation must have returned its reservation while the slot is still held: + // leaked occupancy shows up directly as phantom waiters. + if s := l.Stats(); s.Waiting != 0 { + t.Fatalf("cancelled waiters leaked queue occupancy: %+v", s) + } + // And the queue must still have its full capacity: a fresh waiter can park. + admitted := make(chan bool) go func() { r, ok := l.Acquire(context.Background()) if ok { r() } - queued <- ok + admitted <- ok }() waitForWaiting(t, l, 1) - r2() - if ok := <-queued; !ok { - t.Fatal("caller was shed while the queue was logically empty") + r1() + if ok := <-admitted; !ok { + t.Fatal("queue capacity eroded by cancelled waiters") } } @@ -199,7 +243,49 @@ func TestCloseIsAnAdmissionBarrier(t *testing.T) { } } +func TestCloseBeatsARacingRelease(t *testing.T) { + // A slot released after Close must not be handed to a parked waiter: even when the + // waiter's slot arm wins the select, it re-checks shutdown, returns the slot, and is + // rejected. Loop because which select arm wins is random; every iteration must reject. + for i := 0; i < 100; i++ { + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + r, _ := l.Acquire(context.Background()) + got := make(chan bool) + go func() { + _, ok := l.Acquire(context.Background()) + got <- ok + }() + waitForWaiting(t, l, 1) + l.Close() + r() // released after Close: the waiter must not be admitted + if ok := <-got; ok { + t.Fatalf("iteration %d: waiter admitted with a slot released after Close", i) + } + } +} + +func TestSlotWonConcurrentlyWithCloseIsReturned(t *testing.T) { + // White-box: a waiter that has reserved occupancy and won a slot in the same instant + // Close lands must, on its re-check, return the slot and be rejected -- this is the + // specific arm TestCloseBeatsARacingRelease can only reach when scheduling cooperates. + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + l.inFlight.Add(1) // the waiter's entry reservation + <-l.tokens // the waiter's select wins the slot... + l.Close() // ...as shutdown lands + if _, ok := l.admit(); ok { + t.Fatal("a slot won concurrently with Close must not be admitted") + } + if free := len(l.tokens); free != 1 { + t.Fatalf("the slot was not returned: %d free", free) + } + if s := l.Stats(); s.Held != 0 || s.Waiting != 0 || s.Rejected != 1 { + t.Fatalf("occupancy not released or rejection not counted: %+v", s) + } +} + func TestCloseUnblocksWaiters(t *testing.T) { + // The quiescent case: no slot is released while Close runs (the racing case is + // TestCloseBeatsARacingRelease). l := New("t", Params{MaxConcurrent: 1, MaxQueued: 4}) r1, _ := l.Acquire(context.Background()) results := make(chan bool, 3) @@ -221,6 +307,10 @@ func TestCloseUnblocksWaiters(t *testing.T) { t.Fatal("waiter not unblocked by Close") } } + // The unblocked waiters must have returned their occupancy reservations. + if s := l.Stats(); s.Waiting != 0 { + t.Fatalf("waiters unblocked by Close leaked queue occupancy: %+v", s) + } r1() // releasing a held slot after Close must not panic or block l.Close() // idempotent } From 3416dbf0b056dd15d24a7bdd62b40ba05bc6637c Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:24:24 -0700 Subject: [PATCH 4/6] state the SendTimeout enforcement scoping without a TODO marker The godox linter forbids TODO comments; say the same thing declaratively: enforcement, the default, and the throughput floor land with the wiring change, and nothing reads the option until then. --- config/config.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/config/config.go b/config/config.go index 393621ca..7605945e 100644 --- a/config/config.go +++ b/config/config.go @@ -164,10 +164,10 @@ type ConcurrencyConfig struct { // SendTimeout is the longest a single initialization delivery may hold a concurrency // slot before its connection is closed to reclaim the slot (the SDK then reconnects). // - // TODO(wiring): enforcement is not in this package. The delivery path that consumes - // this option -- applying its default and a per-write throughput floor that cuts a - // stalled client well before this cap -- lands with the code that wires the limiter - // into the endpoints. + // Enforcement is not in this package: the delivery path that consumes this option -- + // applying its default and a per-write throughput floor that cuts a stalled client + // well before this cap -- lands with the code that wires the limiter into the + // endpoints. Until then the option is read by nothing. SendTimeout ct.OptDuration `conf:"INIT_SEND_TIMEOUT"` } From 9ed4d6b22b16953f9c0a84fbe42d85e2cf493eb9 Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:45:44 -0700 Subject: [PATCH 5/6] fix the exact-fit test oracle, pin budget erosion, count parked callers The exact-fit regression test's premise was wrong: a spawned waiter from an earlier iteration can still be live (its slot stolen by the main caller's fast path), legitimately filling the budget, so the test failed intermittently on correct code in exactly the CI configuration that runs without the race detector. The oracle now samples the live spawned callers and judges a rejection spurious only when the budget provably had room; spawned callers retry until admitted rather than being silently dropped, and cleanup cancels and drains them on failure. Still fails at iteration 0 on the previous revisions. Add the two killing tests review found missing: an overflow rejection must back out its reservation (without the back-out, every overflow rejection permanently erodes the budget by one unit until the limiter rejects everything), and a negative MaxQueued must normalize to zero rather than poison the occupancy bound. Stats.Waiting now reports a dedicated parked-callers counter instead of a value derived from the occupancy counter, which could read far outside its documented range in both directions (inflated by entrants mid-rejection, clamped to zero during handoffs). The parked counter feeds Stats only and plays no part in admission. The occupancy-leak test assertions observe the occupancy counter directly, since Waiting no longer reflects it. Doc corrections, each matching measured behavior: the reservation scheme's no-double-count guarantee is scoped to the queue-to-held handoff (a cancelled waiter's reservation is returned when its goroutine next runs); Close is not a drain barrier (Held can grow for an instant after it returns); the MaxQueued config doc no longer contradicts the limiter's brief-wait-at-zero semantics; the config section states that nothing consumes it until the wiring lands; the release-ordering comment claims only what is observable. --- config/config.go | 13 ++- internal/concurrency/limiter.go | 48 ++++++---- internal/concurrency/limiter_test.go | 131 +++++++++++++++++++++++---- 3 files changed, 149 insertions(+), 43 deletions(-) diff --git a/config/config.go b/config/config.go index 7605945e..9e87425a 100644 --- a/config/config.go +++ b/config/config.go @@ -148,17 +148,20 @@ type Config struct { // ConcurrencyConfig corresponds to the [Concurrency] section in the configuration file. // -// It limits how many SDK initialization deliveries the Relay Proxy performs at the same -// time, covering both polling and streaming. This caps the memory and egress that a burst -// of connecting SDKs can consume. The limit is disabled unless MaxConcurrent is set. +// It configures a limit on how many SDK initialization deliveries the Relay Proxy may +// perform at the same time, covering both polling and streaming, to cap the memory and +// egress that a burst of connecting SDKs can consume. The limit is disabled unless +// MaxConcurrent is set. The code that wires the limiter into the endpoints consumes this +// section; until that lands, setting it changes nothing. type ConcurrencyConfig struct { // MaxConcurrent is the maximum number of initialization deliveries allowed in flight // at once. A value of 0 or less disables the limit. MaxConcurrent ct.OptInt `conf:"INIT_MAX_CONCURRENT"` // MaxQueued is the maximum number of clients that may wait for a slot once - // MaxConcurrent is reached. A value of 0 rejects excess clients immediately instead - // of queueing them. + // MaxConcurrent is reached. A value of 0 adds no waiting capacity: an excess client + // is rejected rather than queued, though one arriving just as a slot is released may + // briefly wait to take it. MaxQueued ct.OptInt `conf:"INIT_MAX_QUEUED"` // SendTimeout is the longest a single initialization delivery may hold a concurrency diff --git a/internal/concurrency/limiter.go b/internal/concurrency/limiter.go index 0fe8f40c..9ad7c788 100644 --- a/internal/concurrency/limiter.go +++ b/internal/concurrency/limiter.go @@ -23,8 +23,9 @@ type Params struct { } // Stats is a point-in-time snapshot of a Limiter's counters for logging and metrics. -// Held is exact; Waiting is derived from the occupancy counter and may transiently count -// a caller whose slot handoff is in progress. +// Held is exact. Waiting counts callers parked waiting for a slot; a caller being handed +// a slot, or one whose cancellation has not yet been scheduled, may be counted for the +// instant it takes its goroutine to run. type Stats struct { Enabled bool MaxConcurrent int @@ -46,15 +47,22 @@ type Limiter struct { maxQueued int // inFlight counts the callers occupying the budget: slot holders plus queued waiters, // bounded by MaxConcurrent+MaxQueued. A caller reserves its unit once on entry and - // keeps it from queue to held to released, so handing a released slot to a parked - // waiter never leaves a moment where the waiter still counts against the queue while - // the slot is already spoken for. A separate queue counter has exactly that moment (it - // lasts until the woken goroutine is scheduled) and sheds an admissible caller on - // every slot turnover under burst load. + // keeps it from queue to held to released, so the queue-to-held HANDOFF never has a + // moment where the waiter still counts against the queue while the slot is already + // spoken for. (A separate queue counter has exactly that moment -- it lasts until the + // woken goroutine is scheduled -- and sheds an admissible caller on every slot + // turnover under burst load.) The cancellation path keeps a narrower version of the + // window: a cancelled waiter's reservation is returned when its goroutine next runs, + // so a doomed waiter can briefly occupy budget after its cancellation commits. inFlight atomic.Int64 maxOccupancy int64 - shutdown chan struct{} - closeOnce sync.Once + // parked counts callers blocked waiting for a slot. It feeds Stats.Waiting only and + // plays no part in admission, so it cannot reintroduce turnover shedding; unlike a + // value derived from inFlight, it is bounded by the callers actually waiting rather + // than inflated by entrants mid-rejection. + parked atomic.Int64 + shutdown chan struct{} + closeOnce sync.Once admitted atomic.Int64 rejected atomic.Int64 @@ -125,13 +133,17 @@ func (l *Limiter) Acquire(ctx context.Context) (release func(), ok bool) { // Every slot is held. Wait for one; the occupancy reservation above bounds how many // callers may wait here. + l.parked.Add(1) select { case <-l.tokens: + l.parked.Add(-1) return l.admit() case <-ctx.Done(): + l.parked.Add(-1) l.inFlight.Add(-1) return l.reject() case <-l.shutdown: + l.parked.Add(-1) l.inFlight.Add(-1) return l.reject() } @@ -152,8 +164,8 @@ func (l *Limiter) admit() (func(), bool) { var once sync.Once return func() { once.Do(func() { - // Free the occupancy before returning the slot, so the queue capacity this - // release re-opens is visible to arrivals no later than the slot itself. + // Precautionary ordering: free the occupancy before returning the slot. The + // gap is a couple of instructions either way; no test can observe it. l.inFlight.Add(-1) l.tokens <- struct{}{} }) @@ -167,9 +179,11 @@ func (l *Limiter) reject() (func(), bool) { // Close stops admissions: once it returns, callers entering Acquire are rejected, and // parked waiters are unblocked and rejected. A waiter handed a slot concurrently with -// Close re-checks shutdown and returns the slot instead of keeping it; only an admission -// that fully completed before shutdown landed keeps its slot. Releases of held slots -// remain safe afterwards. It may be called more than once. +// Close re-checks shutdown and returns the slot instead of keeping it; an admission +// whose re-check ran before shutdown landed keeps its slot, so Held can still grow for +// an instant after Close returns -- shutdown code must track its outstanding work rather +// than treat Close as a drain barrier. Releases of held slots remain safe afterwards. It +// may be called more than once. func (l *Limiter) Close() { if !l.Enabled() { return @@ -182,14 +196,12 @@ func (l *Limiter) Stats() Stats { if !l.Enabled() { return Stats{Enabled: false} } - held := cap(l.tokens) - len(l.tokens) - waiting := max(int(l.inFlight.Load())-held, 0) return Stats{ Enabled: true, MaxConcurrent: cap(l.tokens), MaxQueued: l.maxQueued, - Held: held, - Waiting: waiting, + Held: cap(l.tokens) - len(l.tokens), + Waiting: int(l.parked.Load()), Admitted: l.admitted.Load(), Rejected: l.rejected.Load(), } diff --git a/internal/concurrency/limiter_test.go b/internal/concurrency/limiter_test.go index 1c85258b..a8544062 100644 --- a/internal/concurrency/limiter_test.go +++ b/internal/concurrency/limiter_test.go @@ -4,6 +4,7 @@ import ( "context" "runtime" "sync" + "sync/atomic" "testing" "time" ) @@ -119,36 +120,58 @@ func TestBacklogFullRejects(t *testing.T) { } func TestSlotTurnoverDoesNotShedExactFitLoad(t *testing.T) { - // The budget is MaxConcurrent+MaxQueued = 2 and there are never more than two live - // callers, so every rejection is spurious by construction. Each iteration hands the - // held slot to a parked waiter and immediately re-acquires: the vacated queue capacity - // must be available even while the woken waiter is still being scheduled. (A separate - // queued-callers counter fails here on most iterations, because the woken waiter keeps - // counting against the queue until it runs.) + // Each iteration hands the held slot to a parked waiter and immediately re-acquires: + // the vacated queue capacity must be available even while the woken waiter is still + // being scheduled. (A separate queued-callers counter fails here on most iterations, + // because the woken waiter keeps counting against the queue until it runs.) + // + // The oracle must account for stragglers: a spawned waiter from an earlier iteration + // can still be live (its slot stolen by main's fast path), legitimately filling the + // budget. A rejection is judged spurious only when the live spawned callers sampled + // before the acquire provably left room; otherwise main retries. l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) - r, ok := l.Acquire(context.Background()) + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + t.Cleanup(wg.Wait) // cleanups run LIFO: cancel first, then wait + t.Cleanup(cancel) + + r, ok := l.Acquire(ctx) if !ok { t.Fatal("initial acquire") } - var wg sync.WaitGroup + var live atomic.Int64 // spawned callers between entry and their release completing for i := 0; i < 400; i++ { wg.Add(1) + live.Add(1) go func() { defer wg.Done() - if r2, ok := l.Acquire(context.Background()); ok { - r2() + defer live.Add(-1) + for ctx.Err() == nil { + if r2, ok := l.Acquire(ctx); ok { + r2() + return + } + runtime.Gosched() } }() waitForWaiting(t, l, 1) r() // hand the slot to the waiter + liveBefore := live.Load() var reacquired bool - r, reacquired = l.Acquire(context.Background()) + r, reacquired = l.Acquire(ctx) if !reacquired { - t.Fatalf("iteration %d: caller shed while the budget had room", i) + // Budget is 2 and this caller needs 1: with at most 1 live spawned caller, + // the budget provably had room, so the shed is spurious. + if liveBefore <= 1 { + t.Fatalf("iteration %d: caller shed while the budget had room (live=%d)", i, liveBefore) + } + for !reacquired { // stragglers filled the budget; retry once they drain + runtime.Gosched() + r, reacquired = l.Acquire(ctx) + } } } r() - wg.Wait() } func TestCancelledWaitersFreeTheirQueueCapacity(t *testing.T) { @@ -170,9 +193,9 @@ func TestCancelledWaitersFreeTheirQueueCapacity(t *testing.T) { } } // Every cancellation must have returned its reservation while the slot is still held: - // leaked occupancy shows up directly as phantom waiters. - if s := l.Stats(); s.Waiting != 0 { - t.Fatalf("cancelled waiters leaked queue occupancy: %+v", s) + // only r1's occupancy unit may remain. + if got := l.inFlight.Load(); got != 1 { + t.Fatalf("cancelled waiters leaked occupancy: inFlight=%d, want 1", got) } // And the queue must still have its full capacity: a fresh waiter can park. admitted := make(chan bool) @@ -243,10 +266,77 @@ func TestCloseIsAnAdmissionBarrier(t *testing.T) { } } +func TestOverflowRejectionReturnsItsReservation(t *testing.T) { + // A caller rejected for overflow must back its reservation out. Without the back-out, + // every overflow rejection permanently erodes the budget by one unit until the limiter + // rejects everything forever -- the worst accounting failure the design can have. + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + r1, _ := l.Acquire(context.Background()) + admitted := make(chan bool) + go func() { + r, ok := l.Acquire(context.Background()) + if ok { + r() + } + admitted <- ok + }() + waitForWaiting(t, l, 1) // the queue is now full + if _, ok := l.Acquire(context.Background()); ok { + t.Fatal("expected an overflow rejection") + } + r1() // drain: the parked waiter is admitted and releases + if ok := <-admitted; !ok { + t.Fatal("parked waiter should have been admitted") + } + if got := l.inFlight.Load(); got != 0 { + t.Fatalf("overflow rejection leaked occupancy: inFlight=%d, want 0", got) + } + // Full capacity is intact: a holder plus a parked waiter fit again. + r2, ok := l.Acquire(context.Background()) + if !ok { + t.Fatal("slot capacity eroded") + } + done := make(chan bool) + go func() { + r, ok := l.Acquire(context.Background()) + if ok { + r() + } + done <- ok + }() + waitForWaiting(t, l, 1) + r2() + if ok := <-done; !ok { + t.Fatal("queue capacity eroded") + } +} + +func TestNegativeMaxQueuedBehavesAsZero(t *testing.T) { + l := New("t", Params{MaxConcurrent: 2, MaxQueued: -5}) + // Unclamped, a negative queue bound would poison the occupancy limit and reject + // callers while slots sit free. + r1, ok1 := l.Acquire(context.Background()) + r2, ok2 := l.Acquire(context.Background()) + if !ok1 || !ok2 { + t.Fatal("callers within MaxConcurrent must be admitted") + } + if _, ok := l.Acquire(context.Background()); ok { + t.Fatal("expected rejection with no queue capacity") + } + if s := l.Stats(); s.MaxQueued != 0 { + t.Fatalf("negative MaxQueued should normalize to 0: %+v", s) + } + r1() + r2() +} + func TestCloseBeatsARacingRelease(t *testing.T) { // A slot released after Close must not be handed to a parked waiter: even when the // waiter's slot arm wins the select, it re-checks shutdown, returns the slot, and is - // rejected. Loop because which select arm wins is random; every iteration must reject. + // rejected. Loop because which select arm wins is random; in most iterations the + // waiter commits to the shutdown arm before the release lands, so the re-check itself + // is only reliably exercised by TestSlotWonConcurrentlyWithCloseIsReturned -- that + // white-box test is the deterministic regression guard for it. for i := 0; i < 100; i++ { l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) r, _ := l.Acquire(context.Background()) @@ -307,9 +397,10 @@ func TestCloseUnblocksWaiters(t *testing.T) { t.Fatal("waiter not unblocked by Close") } } - // The unblocked waiters must have returned their occupancy reservations. - if s := l.Stats(); s.Waiting != 0 { - t.Fatalf("waiters unblocked by Close leaked queue occupancy: %+v", s) + // The unblocked waiters must have returned their occupancy reservations; only r1's + // unit may remain. + if got := l.inFlight.Load(); got != 1 { + t.Fatalf("waiters unblocked by Close leaked occupancy: inFlight=%d, want 1", got) } r1() // releasing a held slot after Close must not panic or block l.Close() // idempotent From acabacb3ae5fd207a200434a0f0e6f3af0a2605b Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:48:57 -0700 Subject: [PATCH 6/6] pin the parked counter's exit paths and bound the exact-fit retry Swapping the leak-guard assertions to direct occupancy reads last round deleted the only coverage of the parked counter's rejection-arm decrements: dropping either one would inflate Stats.Waiting by one for every client that disconnects while queued -- the exact failure the counter was added to fix -- with the suite green. The guards now assert both the occupancy counter and Waiting, the white-box Close test also checks occupancy directly, and a new test pins that Waiting comes from the parked counter rather than being derived from occupancy (reverting to the derived formula now fails). Bound the exact-fit test's straggler-retry loop with a five-second deadline: against a budget-erosion regression the unbounded loop hung to the package timeout, burning ten minutes before the erosion test's clean diagnostic could run; now it fails in seconds. Doc corrections: Stats.Waiting is bounded by MaxConcurrent+MaxQueued rather than MaxQueued alone (a caller arriving as a slot is released may briefly park without queue capacity); Held and Waiting are independent samples whose sum can transiently exceed the budget, so neither is a basis for exact accounting; the cancellation window is rarer than the handoff one, not shorter; the waitForWaiting helper doc describes the parked counter it now reads. --- internal/concurrency/limiter.go | 17 ++++++---- internal/concurrency/limiter_test.go | 46 ++++++++++++++++++++++------ 2 files changed, 48 insertions(+), 15 deletions(-) diff --git a/internal/concurrency/limiter.go b/internal/concurrency/limiter.go index 9ad7c788..0868bb7e 100644 --- a/internal/concurrency/limiter.go +++ b/internal/concurrency/limiter.go @@ -23,9 +23,13 @@ type Params struct { } // Stats is a point-in-time snapshot of a Limiter's counters for logging and metrics. -// Held is exact. Waiting counts callers parked waiting for a slot; a caller being handed -// a slot, or one whose cancellation has not yet been scheduled, may be counted for the -// instant it takes its goroutine to run. +// Waiting counts callers parked waiting for a slot. It is bounded by +// MaxConcurrent+MaxQueued rather than MaxQueued alone, because a caller arriving as a +// slot is released may briefly park without using queue capacity; a caller being handed +// a slot, or one whose cancellation or shutdown wake-up has not yet been scheduled, may +// be counted for the instant it takes its goroutine to run. Held and Waiting are sampled +// independently, so their sum can transiently exceed the budget; neither is a basis for +// exact accounting. type Stats struct { Enabled bool MaxConcurrent int @@ -51,9 +55,10 @@ type Limiter struct { // moment where the waiter still counts against the queue while the slot is already // spoken for. (A separate queue counter has exactly that moment -- it lasts until the // woken goroutine is scheduled -- and sheds an admissible caller on every slot - // turnover under burst load.) The cancellation path keeps a narrower version of the - // window: a cancelled waiter's reservation is returned when its goroutine next runs, - // so a doomed waiter can briefly occupy budget after its cancellation commits. + // turnover under burst load.) The cancellation path keeps a rarer -- though not + // shorter -- version of the window: a cancelled waiter's reservation is returned when + // its goroutine next runs, so a doomed waiter can occupy budget for a scheduling + // delay after its cancellation commits. inFlight atomic.Int64 maxOccupancy int64 // parked counts callers blocked waiting for a slot. It feeds Stats.Waiting only and diff --git a/internal/concurrency/limiter_test.go b/internal/concurrency/limiter_test.go index a8544062..677e8167 100644 --- a/internal/concurrency/limiter_test.go +++ b/internal/concurrency/limiter_test.go @@ -9,10 +9,9 @@ import ( "time" ) -// waitForWaiting polls until Stats reports n waiting callers. The count is a reservation -// made on entering the queued path, so this establishes that the callers have committed to -// waiting (they reach the select moments later); it is not a strict parked-at-the-select -// barrier. +// waitForWaiting polls until Stats reports n waiting callers. The count is incremented one +// statement before the caller parks at its select, so this establishes the callers are at +// (or an instant from) the select; it is not a strict parked-at-the-select barrier. func waitForWaiting(t *testing.T, l *Limiter, n int) { t.Helper() deadline := time.Now().Add(2 * time.Second) @@ -165,7 +164,14 @@ func TestSlotTurnoverDoesNotShedExactFitLoad(t *testing.T) { if liveBefore <= 1 { t.Fatalf("iteration %d: caller shed while the budget had room (live=%d)", i, liveBefore) } - for !reacquired { // stragglers filled the budget; retry once they drain + // Stragglers filled the budget; retry once they drain. Healthy code needs at + // most a handful of rounds, so a bounded deadline turns an accounting + // regression into a fast failure rather than a package-timeout hang. + deadline := time.Now().Add(5 * time.Second) + for !reacquired { + if time.Now().After(deadline) { + t.Fatalf("iteration %d: could not re-acquire within 5s; budget eroded?", i) + } runtime.Gosched() r, reacquired = l.Acquire(ctx) } @@ -192,11 +198,15 @@ func TestCancelledWaitersFreeTheirQueueCapacity(t *testing.T) { t.Fatal("cancelled waiter should have been rejected") } } - // Every cancellation must have returned its reservation while the slot is still held: - // only r1's occupancy unit may remain. + // Every cancellation must have returned its reservation while the slot is still held + // (only r1's occupancy unit may remain) AND its parked count -- a leaked parked count + // inflates Stats.Waiting forever, one per disconnected queued client. if got := l.inFlight.Load(); got != 1 { t.Fatalf("cancelled waiters leaked occupancy: inFlight=%d, want 1", got) } + if s := l.Stats(); s.Waiting != 0 { + t.Fatalf("cancelled waiters leaked the parked count: %+v", s) + } // And the queue must still have its full capacity: a fresh waiter can park. admitted := make(chan bool) go func() { @@ -368,11 +378,26 @@ func TestSlotWonConcurrentlyWithCloseIsReturned(t *testing.T) { if free := len(l.tokens); free != 1 { t.Fatalf("the slot was not returned: %d free", free) } + if got := l.inFlight.Load(); got != 0 { + t.Fatalf("the rejected admission leaked occupancy: inFlight=%d, want 0", got) + } if s := l.Stats(); s.Held != 0 || s.Waiting != 0 || s.Rejected != 1 { t.Fatalf("occupancy not released or rejection not counted: %+v", s) } } +func TestWaitingReportsParkedCallersNotOccupancy(t *testing.T) { + // Waiting must come from the parked counter, not be derived from occupancy: a + // reservation mid-admission or mid-rejection is not a waiting caller, and deriving + // from occupancy is what let Waiting read far beyond the queue bound. + l := New("t", Params{MaxConcurrent: 2, MaxQueued: 2}) + l.inFlight.Add(1) // occupancy without a parked caller + if s := l.Stats(); s.Waiting != 0 { + t.Fatalf("occupancy alone must not report as Waiting: %+v", s) + } + l.inFlight.Add(-1) +} + func TestCloseUnblocksWaiters(t *testing.T) { // The quiescent case: no slot is released while Close runs (the racing case is // TestCloseBeatsARacingRelease). @@ -397,11 +422,14 @@ func TestCloseUnblocksWaiters(t *testing.T) { t.Fatal("waiter not unblocked by Close") } } - // The unblocked waiters must have returned their occupancy reservations; only r1's - // unit may remain. + // The unblocked waiters must have returned their occupancy reservations (only r1's + // unit may remain) and their parked counts. if got := l.inFlight.Load(); got != 1 { t.Fatalf("waiters unblocked by Close leaked occupancy: inFlight=%d, want 1", got) } + if s := l.Stats(); s.Waiting != 0 { + t.Fatalf("waiters unblocked by Close leaked the parked count: %+v", s) + } r1() // releasing a held slot after Close must not panic or block l.Close() // idempotent }