diff --git a/config/config.go b/config/config.go index e2d64ef9..9e87425a 100644 --- a/config/config.go +++ b/config/config.go @@ -148,27 +148,29 @@ 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"` - // 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 longest a single initialization delivery may hold a concurrency + // slot before its connection is closed to reclaim the slot (the SDK then reconnects). + // + // 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"` } 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..0868bb7e 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. +// maximum number of slots held at once, and a bounded queue of callers waiting for a +// slot in approximate arrival order. package concurrency import ( @@ -16,16 +15,21 @@ 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 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 - // 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. +// 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 @@ -36,30 +40,40 @@ 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 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 MaxConcurrent slots; receive to acquire, send to release - maxQueued int64 - waiting int64 + tokens chan struct{} // holds the free slots; receive to acquire one, send to release one + 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 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 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 + // 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 - 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 { @@ -70,98 +84,111 @@ 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{}) - if p.PerEnvMax > 0 { - l.perEnvMax = p.PerEnvMax - } 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 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 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 } - // 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 - } + // 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() } - reject := func() (func(), bool) { - if releaseEnv != nil { - releaseEnv() - } - l.rejected.Add(1) - return func() {}, false + // 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() } - // 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() default: } - // No token free: enter the bounded FIFO backlog, or reject. - if atomic.AddInt64(&l.waiting, 1) > l.maxQueued { - atomic.AddInt64(&l.waiting, -1) - return reject() - } - defer atomic.AddInt64(&l.waiting, -1) - + // 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: - return l.admit(releaseEnv), true + l.parked.Add(-1) + return l.admit() case <-ctx.Done(): - return reject() + l.parked.Add(-1) + l.inFlight.Add(-1) + return l.reject() case <-l.shutdown: - return reject() + l.parked.Add(-1) + l.inFlight.Add(-1) + return l.reject() } } -func (l *Limiter) admit(releaseEnv func()) 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() { + // 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{}{} - if releaseEnv != nil { - releaseEnv() - } }) - } + }, true } -func (l *Limiter) gateFor(envKey string) *envGate { - if g, ok := l.perEnv.Load(envKey); ok { - return g.(*envGate) - } - g := &envGate{slots: make(chan struct{}, l.perEnvMax)} - actual, _ := l.perEnv.LoadOrStore(envKey, g) - return actual.(*envGate) +func (l *Limiter) reject() (func(), bool) { + l.rejected.Add(1) + return func() {}, false } -// Close unblocks all waiters (they receive ok=false). Idempotent. +// 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; 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 @@ -177,9 +204,9 @@ func (l *Limiter) Stats() Stats { return Stats{ Enabled: true, MaxConcurrent: cap(l.tokens), - MaxQueued: int(l.maxQueued), + MaxQueued: l.maxQueued, Held: cap(l.tokens) - len(l.tokens), - Waiting: int(atomic.LoadInt64(&l.waiting)), + 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 650aafd1..677e8167 100644 --- a/internal/concurrency/limiter_test.go +++ b/internal/concurrency/limiter_test.go @@ -2,18 +2,34 @@ package concurrency import ( "context" + "runtime" "sync" + "sync/atomic" "testing" "time" ) +// 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) + 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() { 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,18 +39,22 @@ 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") } release() + if l.Name() != "" { + t.Fatal("nil limiter Name should be empty") + } + l.Close() // must not panic } 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,20 +70,19 @@ 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() } }() - // 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") @@ -79,39 +98,143 @@ 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 - defer r1() - go l.Acquire(context.Background(), "e") // fills the single backlog slot - time.Sleep(50 * time.Millisecond) - if _, ok := l.Acquire(context.Background(), "e"); ok { + r1, _ := l.Acquire(context.Background()) // holds the only token + 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 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) +func TestSlotTurnoverDoesNotShedExactFitLoad(t *testing.T) { + // 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}) + 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 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() + 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(ctx) + if !reacquired { + // 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) + } + // 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) + } + } + } + r() +} + +func TestCancelledWaitersFreeTheirQueueCapacity(t *testing.T) { + l := New("t", Params{MaxConcurrent: 1, MaxQueued: 1}) + r1, _ := l.Acquire(context.Background()) + // 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") + } + } + // 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() { + r, ok := l.Acquire(context.Background()) + if ok { + r() + } + admitted <- ok + }() + waitForWaiting(t, l, 1) + r1() + if ok := <-admitted; !ok { + t.Fatal("queue capacity eroded by cancelled waiters") } - 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) + waitForWaiting(t, l, 1) + before := l.Stats().Rejected cancel() select { case ok := <-done: @@ -121,24 +244,245 @@ 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 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; 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()) + 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 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). + 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") + } + } + // 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 } 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() } +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}) @@ -149,7 +493,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 }