refactor: simplify the init-concurrency limiter to a single shared budget

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"`
 }
 

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))
 	})
 }

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(),
 	}