Share concurrent request limits across the fleet

[neo-sky]

Summary

• admit against a shared lease table so replicas enforce one limit between them instead of one each
• add a shadow mode that counts leases fleet-wide while rejection stays per-replica
• fall back to this replica's own leases when the lease store cannot answer
• renew, sweep and release leases from supervised background tasks

Rollout

FLEET_CONCURRENCY_MODE defaults to off, so merging changes nothing. Replicas
left in off write no leases at all, so move the whole fleet to shadow before
reading anything from the numbers, then compare concurrency.would_reject
against concurrency.admitted to size the impact. Move to enforce one replica
at a time after that. Rollback is setting the mode back to off.

Verification

• cargo fmt --all -- --check
• cargo clippy --all-targets --all-features -- -D warnings
• cargo nextest run --lib --bins (1407 passed)
• cargo nextest run --test integration_tests (8 passed)
• cargo nextest run --test e2e_all (698 passed, fresh database)
• cargo build --release

[claude]

Review: Share concurrent request limits across the fleet

Solid design overall — the advisory-lock admission, LeaseState::PendingWrite/mark_stored orphan handling, and the renew-never-inserts invariant are all correct and well tested. Good that FLEET_CONCURRENCY_MODE=off is a true no-op.

Four things I would want fixed before this rolls out, two of which affect shadow — the mode the rollout plan says to enable first.

🔴 1. shadow mode silently removes the existing per-replica limit

crates/services/src/completions/mod.rs:1880-1911

When the fleet lease is admitted, try_acquire_concurrent_slot returns early with ConcurrentSlot::Lease and never touches concurrent_counts. That cache has exactly one increment site (line 1907), reached only on the LeaseAdmission::Shadowed fall-through.

So on a replica running shadow:

• under the fleet limit → lease admitted, local counter stays at 0
• over the fleet limit → Shadowed → falls to the local path, which sees 0 and admits

The effective ceiling becomes fleet_limit + (replicas × limit) instead of today's replicas × limit. Enabling the "measure-only" mode on one instance raises the concurrency that instance accepts.

shadowing_records_the_verdict_without_rejecting asserts only slots.len() > LIMIT, so it passes either way.

Fix — in shadow mode, use the lease call purely for measurement and let the local counter decide:

match self.try_acquire_lease(fleet, organization_id, model_id, model_name).await {
    Ok(slot) if fleet.enforcing => return Ok(slot),
    Ok(slot) => { slot.release(); }   // measured fleet-wide; local counter still decides
    Err(LeaseAdmission::AtLimit(error)) => return Err(error),
    ...
}

(or increment the local counter alongside the lease and have ConcurrentSlot::Lease carry it). Either way, add an upper-bound assertion to the shadow test.

🔴 2. shadow mode rejects real traffic when the lease store is unreachable

crates/services/src/completions/mod.rs:1888-1899 → admit_from_held_leases

The Unavailable branch calls admit_from_held_leases without consulting fleet.enforcing, and that function returns CompletionError::RateLimitExceeded on Err(in_flight). A DB blip during the shadow bake therefore produces 429s from a mode documented as non-rejecting. Gate the rejection on fleet.enforcing.

🟠 3. pg_advisory_xact_lock on the inference hot path with no timeout

crates/database/src/repositories/concurrency_lease.rs:411-425

Every completion request now opens a transaction and takes a fleet-wide exclusive lock keyed on (org, model) while holding a deadpool connection. There is no lock_timeout/statement_timeout on the transaction (unlike reporting_query.rs, which sets one) and no tokio::time::timeout around try_acquire.

LeaseAdmission::Unavailable only fires on an error — it cannot fire on a hang. If the primary stalls or one holder is slow, requests for a hot (org, model) queue on the lock while holding pooled connections, starving every other repository in the process. retry_db! retries PoolError up to 3× with backoff, which amplifies it.

tx.execute("SET LOCAL lock_timeout = '2s'", &[]).await.map_err(map_db_error)?;

plus a client-side tokio::time::timeout around the repository call, so a hang degrades to the replica-local path instead of blocking the request.

🟡 4. Failed limit lookups are no longer cached

crates/services/src/completions/mod.rs:1252-1272

The get_with → optionally_get_with switch means a failing get_concurrent_limit is no longer memoised. The comment frames this as avoiding a pinned default for the whole TTL, which is fair, but the consequence is that during a lookup outage every request issues a fresh query retried up to 3× against an already-unhealthy DB. A short negative-cache TTL (a few seconds) keeps the intent without the amplification.

Minor

• sweep_expired deletes at most 1000 rows per tick and ticks once per TTL (60s default). If replicas die holding more than ~1000 leases/minute, the table grows faster than it is reclaimed. Consider looping until a tick reclaims fewer than the batch size.

Verified as correct

• renew updates in place and cannot resurrect a released row; mark_stored skips ids no longer held; the orphan-release path after persist covers the release-during-write race.
• organizations.rate_limit + is_active = true in the admission query matches PgOrganizationRepository::get_concurrent_limit, and effective_limit reproduces the > 0 semantics.
• database.pool() is the Patroni write pool — correct for a read-then-insert transaction.
• Migration V0075 is additive with appropriate indexes; safe for rolling updates.
• No customer data in any of the new log statements (ids, counts, and error types only).

⚠️

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: 9cb3366c-2601-44c7-a92c-408735cfe8e8
• Base: main at 54b4a7c
• Head: feat/fleet-concurrency-leases at f783811
• Created: 2026-08-25 01:39 UTC
• Updated: 2026-08-25 01:46 UTC

_{Automatic trigger · attempt 1 of 3 · completed in 6m 54s}

[ironloopai]

Review · Summary

Found three actionable fleet-concurrency issues.

Findings: 🔴 High 1 · 🟠 Medium 2

_{Code-specific findings are attached to the diff.}

Validation

• ✅ Captured CI — The captured unit, integration, E2E, lint, test-suite, and release-build checks passed.

Review details

• Run: 9cb3366c-2601-44c7-a92c-408735cfe8e8
• Attempts: 1

[ironloopai]

🔴 High · Bound the asynchronous lease-release queue

Every completed request sends an ID into an unbounded channel, while the single receiver awaits database release calls. During a lease-store stall or outage, fallback traffic can continue completing faster than retries drain, so this shared queue grows without bound and can exhaust process memory. Use a bounded/coalescing nonblocking queue and allow dropped release notifications to expire safely.

[ironloopai]

🟠 Medium · Make acquisition retries idempotent before checking capacity

A connection loss after the server commits can make retry_db! rerun this operation with the same lease ID. If other requests fill the remaining capacity before the retry, the count check returns AtLimit before reaching the conflict-tolerant insert. The request is then rejected and never tracks or releases its already-committed lease, consuming capacity until its TTL. Detect an existing lease_id (or exclude it from the count) before the limit check.

[ironloopai]

🟠 Medium · Do not use a single shadow replica as fleet-wide rollout evidence

Instances left in off mode never create FleetConcurrency, so they continue using only local counters and write no leases. Consequently, the documented one-instance shadow rollout measures only that instance, and enabling enforce one instance at a time still ignores traffic on off replicas. Require all replicas to write shadow leases before relying on the metric or performing a gradual enforce rollout.

[github-actions]

🔍 OpenCodeReview found 5 issue(s) in this PR.

• ✅ 5 posted as inline comment(s)
• 📝 0 posted as summary

[github-actions]

In containerized environments (Kubernetes, Docker), std::process::id() is almost always 1, and HOSTNAME is typically the pod/container name. If HOSTNAME is unset in some environments, the fallback produces unknown-1 — identical across all replicas. While this does not cause functional lease corruption (leases are tracked by UUID in the concurrency_leases table, not by instance_id), it makes the instance_id column useless for identifying which replica holds a lease during fleet-wide concurrency incidents. Consider generating a random suffix (e.g., a short UUID) to guarantee uniqueness even when HOSTNAME is missing.

Suggestion:

Suggested change

	let host = env::var("HOSTNAME").unwrap_or_else(|_| "unknown".to_string());
	format!("{host}-{}", std::process::id())
	let host = env::var("HOSTNAME").unwrap_or_else(|_| "unknown".to_string());
	format!("{host}-{}-{}", std::process::id(), uuid::Uuid::new_v4().simple())

[github-actions]

Unlike FLEET_CONCURRENCY_MODE which warns on unrecognized values, an invalid FLEET_CONCURRENCY_LEASE_TTL_SECONDS (non-numeric, zero, or negative) silently falls back to 60 with no warning. An operator who sets a typo or invalid value may be unaware the default is being used. Consider adding an eprintln! warning consistent with the mode handling for unrecognized values.

Suggestion:

Suggested change

	lease_ttl_seconds: env::var("FLEET_CONCURRENCY_LEASE_TTL_SECONDS")
	.ok()
	.and_then(|value| value.parse::<u64>().ok())
	.filter(|seconds| *seconds > 0)
	.unwrap_or(60),
	lease_ttl_seconds: match env::var("FLEET_CONCURRENCY_LEASE_TTL_SECONDS") {
	Ok(value) => match value.parse::<u64>() {
	Ok(seconds) if seconds > 0 => seconds,
	_ => {
	eprintln!(
	"WARN: invalid FLEET_CONCURRENCY_LEASE_TTL_SECONDS '{value}', \
	falling back to 60"
	);
	60
	}
	},
	Err(_) => 60,
	},

[github-actions]

The migration lacks a descriptive header comment explaining the table's purpose. Other recent migrations (e.g. V0074) include a brief comment. Adding one here would help future maintainers understand that this table tracks fleet-wide concurrency leases for rate-limit enforcement across replicas.

Suggestion:

Suggested change

	CREATE TABLE concurrency_leases (
	id UUID PRIMARY KEY,
	organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
	model_id UUID NOT NULL REFERENCES models(id) ON DELETE CASCADE,
	instance_id TEXT NOT NULL,
	acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
	expires_at TIMESTAMPTZ NOT NULL
	);
	-- Fleet-wide concurrency leases: each in-flight request acquires a lease
	-- with a TTL; the active count per (organization, model) is checked against
	-- the org's rate_limit before admission. Expired leases are swept periodically.
	CREATE TABLE concurrency_leases (
	id UUID PRIMARY KEY,
	organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
	model_id UUID NOT NULL REFERENCES models(id) ON DELETE CASCADE,
	instance_id TEXT NOT NULL,
	acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
	expires_at TIMESTAMPTZ NOT NULL,
	CHECK (expires_at > acquired_at)
	);

[github-actions]

Consider adding a CHECK constraint (expires_at > acquired_at) as defense-in-depth. The application code always sets expires_at = NOW() + positive TTL, so the invariant holds today, but a database-level constraint would prevent accidental insertion of already-expired leases and make the intent explicit.

Suggestion:

Suggested change

	acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
	expires_at TIMESTAMPTZ NOT NULL
	);
	acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
	expires_at TIMESTAMPTZ NOT NULL,
	CHECK (expires_at > acquired_at)
	);

[github-actions]

Using Fn as the closure bound is unnecessarily restrictive for a function that calls the closure multiple times (once per instance). Fn prohibits the closure from mutating captured state, which means a future test author cannot assign unique instance_id values per replica — a realistic fleet scenario — using a simple mutable counter:

let mut i = 0;
setup_test_fleet(INSTANCES, |config| {
    config.fleet_concurrency.instance_id = format!("instance-{i}");
    i += 1; // requires FnMut, not Fn
}).await;

FnMut is the idiomatic bound for closures that are called multiple times and may need to track state across invocations. All closures that satisfy Fn also satisfy FnMut, so existing callers are unaffected.

Suggestion:

Suggested change

	where
	F: Fn(&mut config::ApiConfig),
	{
	where
	F: FnMut(&mut config::ApiConfig),
	{

[PierreLeGuen]

In shadow mode a fleet-admitted request returns before touching the per-replica counter, so a replica can hold its full lease allowance plus a full local allowance — up to 2x the organization limit — in the first rollout step that is meant to change no admission decision.

Blocking findings:

• crates/services/src/completions/mod.rs:1912 — In shadow mode a fleet-admitted request returns early from try_acquire_concurrent_slot and never increments the concurrent_counts atomic. Impact: With an organization limit L, one replica in shadow mode admits up to 2L concurrent backend requests for that… Fix: Do not let a fleet admission bypass the per-replica limiter when !fleet.enforcing.

Checks: cargo +1.92.0 fmt --all -- --check — passed; cargo +1.92.0 build -p services --tests — exit 0; cargo +1.92.0 clippy -p services -p config -p database --all-targets — clean, no new warnings

[neo-sky]

Fixed. While shadowing the fleet only observes now, so the per-replica counter decides admission and a lease can't bypass it. The store-outage path had the same shape and defers too.

shadowing_does_not_admit_past_the_replica_limit pins it, admitting 6 against a limit of 3 on the old code.