From 5cff7e2da248c47ddd9fb0cd24e2a580750d2bc4 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Wed, 12 Aug 2026 19:31:20 -0400 Subject: [PATCH 1/8] wip(concurrency): fleet-wide lease admission behind a flag Admission goes through a Postgres lease table serialized per organization and model by an advisory lock, so replicas share one count instead of each enforcing the limit locally. Off by default via FLEET_CONCURRENCY_MODE, falling back to the existing per-process counter on any database error. Not finished: leases have no heartbeat, so a request outliving its TTL frees capacity while still in flight, and there are no metrics yet. --- crates/api/src/lib.rs | 21 +- crates/api/tests/common/mod.rs | 32 +++ crates/api/tests/e2e_all/fleet_concurrency.rs | 134 +++++++++++ crates/api/tests/e2e_all/main.rs | 1 + crates/config/src/types.rs | 58 +++++ .../sql/V0073__add_concurrency_leases.sql | 17 ++ .../src/repositories/concurrency_lease.rs | 139 +++++++++++ crates/database/src/repositories/mod.rs | 1 + crates/database/tests/concurrency_leases.rs | 214 +++++++++++++++++ crates/inference_providers/src/mock.rs | 18 +- crates/services/src/completions/mod.rs | 217 ++++++++++++++---- crates/services/src/completions/ports.rs | 25 ++ 12 files changed, 834 insertions(+), 43 deletions(-) create mode 100644 crates/api/tests/e2e_all/fleet_concurrency.rs create mode 100644 crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql create mode 100644 crates/database/src/repositories/concurrency_lease.rs create mode 100644 crates/database/tests/concurrency_leases.rs diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index b972d9016..7018e4e9b 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -424,14 +424,29 @@ pub async fn init_domain_services_with_pool( as Arc; // Create completion service with usage tracking (needs usage_service) - let completion_service = Arc::new(services::CompletionServiceImpl::new( + let mut completion_service_impl = services::CompletionServiceImpl::new( inference_provider_pool.clone(), attestation_service.clone(), usage_service.clone(), metrics_service.clone(), models_repo.clone() as Arc, org_limit_repository, - )); + ); + + if config.fleet_concurrency.mode == config::FleetConcurrencyMode::Enforce { + let lease_repository = Arc::new( + database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository::new( + database.pool().clone(), + ), + ); + completion_service_impl = completion_service_impl.with_fleet_concurrency( + lease_repository, + config.fleet_concurrency.instance_id.clone(), + std::time::Duration::from_secs(config.fleet_concurrency.lease_ttl_seconds), + ); + } + + let completion_service = Arc::new(completion_service_impl); let brave_search_provider = Arc::new(services::responses::tools::brave::BraveWebSearchProvider::new()); @@ -2656,6 +2671,7 @@ mod tests { aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), ita: config::ItaAttestationConfig::default(), + fleet_concurrency: config::FleetConcurrencyConfig::default(), }; // Initialize services @@ -2765,6 +2781,7 @@ mod tests { aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), ita: config::ItaAttestationConfig::default(), + fleet_concurrency: config::FleetConcurrencyConfig::default(), }; let auth_components = init_auth_services(database.clone(), &config); diff --git a/crates/api/tests/common/mod.rs b/crates/api/tests/common/mod.rs index b1b5a9a64..e85cd2214 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -130,6 +130,7 @@ pub fn test_config() -> ApiConfig { ..config::UsageReportingConfig::default() }, ita: config::ItaAttestationConfig::default(), + fleet_concurrency: config::FleetConcurrencyConfig::default(), } } @@ -433,6 +434,37 @@ where (server, database) } +/// Build several independent servers that share one database, standing in for +/// the replicas of a deployed fleet. Each gets its own services and provider +/// pool, so anything they agree on has to travel through the database. +pub async fn setup_test_fleet( + instances: usize, + mutate: F, +) -> ( + Vec, + Vec>, + Arc, +) +where + F: Fn(&mut config::ApiConfig), +{ + let infra = setup_test_infrastructure().await; + let database = infra.database.clone(); + + let mut servers = Vec::with_capacity(instances); + let mut mocks = Vec::with_capacity(instances); + for _ in 0..instances { + let mut config = infra.config.clone(); + mutate(&mut config); + let (server, _pool, mock, _router) = + build_test_server_components(database.clone(), config).await; + servers.push(server); + mocks.push(mock); + } + + (servers, mocks, database) +} + pub async fn setup_test_server_with_database() -> (axum_test::TestServer, Arc) { let (server, _, _, database) = setup_test_server_with_pool().await; (server, database) diff --git a/crates/api/tests/e2e_all/fleet_concurrency.rs b/crates/api/tests/e2e_all/fleet_concurrency.rs new file mode 100644 index 000000000..8daa372cb --- /dev/null +++ b/crates/api/tests/e2e_all/fleet_concurrency.rs @@ -0,0 +1,134 @@ +use crate::common; +use inference_providers::mock::ResponseTemplate; +use std::time::Duration; +use uuid::Uuid; + +const INSTANCES: usize = 4; +const LIMIT: i32 = 2; +const ATTEMPTS: usize = 8; +const HOLD: Duration = Duration::from_secs(2); +const MODEL: &str = "Qwen/Qwen3-30B-A3B-Instruct-2507"; + +async fn set_rate_limit(database: &database::Database, organization_id: &str, limit: i32) { + let organization_id = Uuid::parse_str(organization_id).expect("organization id is a uuid"); + database + .pool() + .get() + .await + .expect("database connection") + .execute( + "UPDATE organizations SET rate_limit = $1 WHERE id = $2", + &[&limit, &organization_id], + ) + .await + .expect("rate limit update"); +} + +async fn live_lease_count(database: &database::Database, organization_id: &str) -> i64 { + let organization_id = Uuid::parse_str(organization_id).expect("organization id is a uuid"); + database + .pool() + .get() + .await + .expect("database connection") + .query_one( + "SELECT COUNT(*) FROM concurrency_leases + WHERE organization_id = $1 AND expires_at > NOW()", + &[&organization_id], + ) + .await + .expect("lease count") + .get(0) +} + +/// Four replicas sharing one database must not admit more than the +/// organization's limit between them. Without fleet-wide accounting each +/// replica counts only itself, so this admits INSTANCES * LIMIT. +#[tokio::test] +async fn fleet_limit_is_shared_across_replicas() { + let (servers, mocks, database) = common::setup_test_fleet(INSTANCES, |config| { + config.fleet_concurrency.mode = config::FleetConcurrencyMode::Enforce; + }) + .await; + + for mock in &mocks { + mock.set_default_response(ResponseTemplate::new("held").with_hold(HOLD)) + .await; + } + + let org = common::setup_org_with_credits(&servers[0], 100_000_000_000).await; + let api_key = common::get_api_key_for_org(&servers[0], org.id.clone()).await; + set_rate_limit(&database, &org.id, LIMIT).await; + + let attempts = (0..ATTEMPTS).map(|attempt| { + let server = &servers[attempt % INSTANCES]; + let api_key = api_key.clone(); + async move { + server + .post("/v1/chat/completions") + .add_header("Authorization", format!("Bearer {api_key}")) + .json(&serde_json::json!({ + "model": MODEL, + "messages": [{ "role": "user", "content": "hello" }] + })) + .await + .status_code() + .as_u16() + } + }); + + let statuses = futures::future::join_all(attempts).await; + let admitted = statuses.iter().filter(|status| **status == 200).count(); + let rejected = statuses.iter().filter(|status| **status == 429).count(); + + assert_eq!( + admitted, LIMIT as usize, + "expected the fleet to admit exactly {LIMIT} of {ATTEMPTS}, got {statuses:?}" + ); + assert_eq!( + rejected, + ATTEMPTS - LIMIT as usize, + "every attempt over the limit should be rejected, got {statuses:?}" + ); +} + +/// A finished request must give its slot back, or the organization stays +/// locked out until the leases expire. +#[tokio::test] +async fn completed_requests_release_their_lease() { + let (servers, mocks, database) = common::setup_test_fleet(1, |config| { + config.fleet_concurrency.mode = config::FleetConcurrencyMode::Enforce; + }) + .await; + + mocks[0] + .set_default_response(ResponseTemplate::new("done")) + .await; + + let org = common::setup_org_with_credits(&servers[0], 100_000_000_000).await; + let api_key = common::get_api_key_for_org(&servers[0], org.id.clone()).await; + set_rate_limit(&database, &org.id, LIMIT).await; + + for _ in 0..(ATTEMPTS as i32 + LIMIT) { + let status = servers[0] + .post("/v1/chat/completions") + .add_header("Authorization", format!("Bearer {api_key}")) + .json(&serde_json::json!({ + "model": MODEL, + "messages": [{ "role": "user", "content": "hello" }] + })) + .await + .status_code(); + assert_eq!( + status, 200, + "a released slot should let the next request straight through" + ); + } + + tokio::time::sleep(Duration::from_millis(500)).await; + assert_eq!( + live_lease_count(&database, &org.id).await, + 0, + "leases outlived the requests that took them" + ); +} diff --git a/crates/api/tests/e2e_all/main.rs b/crates/api/tests/e2e_all/main.rs index cb7844523..c14ae42bf 100644 --- a/crates/api/tests/e2e_all/main.rs +++ b/crates/api/tests/e2e_all/main.rs @@ -41,6 +41,7 @@ mod error_msg; mod external_providers; mod feature_requests; mod files; +mod fleet_concurrency; mod function_tools; mod general; mod glm52_tier_routing; diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index 2598a6cb3..01fd55d53 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -28,6 +28,63 @@ pub struct ApiConfig { pub aml: AmlConfig, pub usage_reporting: UsageReportingConfig, pub ita: ItaAttestationConfig, + pub fleet_concurrency: FleetConcurrencyConfig, +} + +/// How concurrent-request limits are enforced across replicas. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum FleetConcurrencyMode { + /// Each replica counts only its own in-flight requests. + #[default] + Off, + /// Admission is decided by the fleet-wide lease count. + Enforce, +} + +#[derive(Debug, Clone)] +pub struct FleetConcurrencyConfig { + pub mode: FleetConcurrencyMode, + pub lease_ttl_seconds: u64, + pub instance_id: String, +} + +impl Default for FleetConcurrencyConfig { + fn default() -> Self { + Self { + mode: FleetConcurrencyMode::Off, + lease_ttl_seconds: 30, + instance_id: String::new(), + } + } +} + +impl FleetConcurrencyConfig { + pub fn from_env() -> Self { + let mode = match env::var("FLEET_CONCURRENCY_MODE") + .unwrap_or_default() + .to_ascii_lowercase() + .as_str() + { + "enforce" => FleetConcurrencyMode::Enforce, + _ => FleetConcurrencyMode::Off, + }; + + Self { + mode, + lease_ttl_seconds: env::var("FLEET_CONCURRENCY_LEASE_TTL_SECONDS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .unwrap_or(30), + instance_id: env::var("FLEET_CONCURRENCY_INSTANCE_ID") + .ok() + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| { + let host = env::var("HOSTNAME").unwrap_or_else(|_| "unknown".to_string()); + format!("{host}-{}", std::process::id()) + }), + } + } } impl ApiConfig { @@ -60,6 +117,7 @@ impl ApiConfig { infra: InfraConfig::from_env(), aml: AmlConfig::from_env()?, ita: ItaAttestationConfig::from_env()?, + fleet_concurrency: FleetConcurrencyConfig::from_env(), usage_reporting: UsageReportingConfig::from_env()?, }) } diff --git a/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql b/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql new file mode 100644 index 000000000..3d55473d8 --- /dev/null +++ b/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql @@ -0,0 +1,17 @@ +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 +); + +CREATE INDEX idx_concurrency_leases_org_model + ON concurrency_leases (organization_id, model_id, expires_at); + +CREATE INDEX idx_concurrency_leases_expires_at + ON concurrency_leases (expires_at); + +CREATE INDEX idx_concurrency_leases_instance + ON concurrency_leases (instance_id, expires_at); diff --git a/crates/database/src/repositories/concurrency_lease.rs b/crates/database/src/repositories/concurrency_lease.rs new file mode 100644 index 000000000..019203c86 --- /dev/null +++ b/crates/database/src/repositories/concurrency_lease.rs @@ -0,0 +1,139 @@ +use crate::pool::DbPool; +use crate::repositories::utils::map_db_error; +use crate::retry_db; +use anyhow::{Context, Result}; +use async_trait::async_trait; +use services::common::RepositoryError; +use services::completions::ports::{ConcurrencyLeaseRepository, LeaseOutcome}; +use std::time::Duration; +use uuid::Uuid; + +#[derive(Debug, Clone)] +pub struct PostgresConcurrencyLeaseRepository { + pool: DbPool, +} + +impl PostgresConcurrencyLeaseRepository { + pub fn new(pool: DbPool) -> Self { + Self { pool } + } + + fn effective_limit(stored: Option, default_limit: u32) -> u32 { + stored + .and_then(|value| u32::try_from(value).ok()) + .unwrap_or(default_limit) + } +} + +#[async_trait] +impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { + async fn try_acquire( + &self, + lease_id: Uuid, + organization_id: Uuid, + model_id: Uuid, + instance_id: &str, + default_limit: u32, + ttl: Duration, + ) -> Result { + let ttl_seconds = ttl.as_secs() as f64; + let lock_key = format!("{organization_id}:{model_id}"); + + let outcome = retry_db!("try_acquire_concurrency_lease", { + let mut client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + let tx = client.transaction().await.map_err(map_db_error)?; + + tx.execute( + "SELECT pg_advisory_xact_lock(hashtextextended($1::text, 0))", + &[&lock_key], + ) + .await + .map_err(map_db_error)?; + + let row = tx + .query_one( + r#" + SELECT + ( + SELECT rate_limit FROM organizations + WHERE id = $1 AND is_active = true + ) AS rate_limit, + ( + SELECT COUNT(*) FROM concurrency_leases + WHERE organization_id = $1 + AND model_id = $2 + AND expires_at > NOW() + ) AS in_flight + "#, + &[&organization_id, &model_id], + ) + .await + .map_err(map_db_error)?; + + let limit = + Self::effective_limit(row.get::<_, Option>("rate_limit"), default_limit); + let in_flight: i64 = row.get("in_flight"); + + if in_flight >= i64::from(limit) { + tx.commit().await.map_err(map_db_error)?; + return Ok::(LeaseOutcome::AtLimit { + limit, + in_flight, + }); + } + + tx.execute( + r#" + INSERT INTO concurrency_leases + (id, organization_id, model_id, instance_id, expires_at) + VALUES ($1, $2, $3, $4, NOW() + make_interval(secs => $5)) + ON CONFLICT (id) DO NOTHING + "#, + &[ + &lease_id, + &organization_id, + &model_id, + &instance_id, + &ttl_seconds, + ], + ) + .await + .map_err(map_db_error)?; + + tx.commit().await.map_err(map_db_error)?; + Ok(LeaseOutcome::Admitted) + })?; + + Ok(outcome) + } + + async fn release(&self, lease_ids: &[Uuid]) -> Result<()> { + if lease_ids.is_empty() { + return Ok(()); + } + + retry_db!("release_concurrency_leases", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + client + .execute( + "DELETE FROM concurrency_leases WHERE id = ANY($1)", + &[&lease_ids], + ) + .await + .map_err(map_db_error) + })?; + + Ok(()) + } +} diff --git a/crates/database/src/repositories/mod.rs b/crates/database/src/repositories/mod.rs index f5411510a..52aeb8eec 100644 --- a/crates/database/src/repositories/mod.rs +++ b/crates/database/src/repositories/mod.rs @@ -4,6 +4,7 @@ pub mod aml; pub mod analytics; pub mod api_key; pub mod attestation; +pub mod concurrency_lease; pub mod conversation; pub mod feature_request; pub mod file; diff --git a/crates/database/tests/concurrency_leases.rs b/crates/database/tests/concurrency_leases.rs new file mode 100644 index 000000000..772a6d641 --- /dev/null +++ b/crates/database/tests/concurrency_leases.rs @@ -0,0 +1,214 @@ +mod support; + +use database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository; +use services::completions::ports::{ConcurrencyLeaseRepository, LeaseOutcome}; +use std::sync::Arc; +use std::time::Duration; +use uuid::Uuid; + +const LIMIT: i32 = 3; +const ATTEMPTS: usize = 24; +const TTL: Duration = Duration::from_secs(30); +const DEFAULT_LIMIT: u32 = 64; + +/// Every replica runs this same read-then-insert, so without the advisory lock +/// they all read the same count before any of them has inserted and the limit +/// is exceeded by however many arrive together. +#[tokio::test] +async fn concurrent_acquires_never_exceed_the_limit() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = $1 WHERE id = $2", + &[&LIMIT, &org.org_id], + ) + .await + .expect("rate limit update"); + + let repository = Arc::new(PostgresConcurrencyLeaseRepository::new(pool.clone())); + + let attempts: Vec<_> = (0..ATTEMPTS) + .map(|attempt| { + let repository = repository.clone(); + let org_id = org.org_id; + let model_id = model.id; + tokio::spawn(async move { + repository + .try_acquire( + Uuid::new_v4(), + org_id, + model_id, + &format!("instance-{}", attempt % 4), + DEFAULT_LIMIT, + TTL, + ) + .await + }) + }) + .collect(); + + let mut admitted = 0usize; + let mut at_limit = 0usize; + for attempt in attempts { + match attempt.await.expect("task").expect("acquire") { + LeaseOutcome::Admitted => admitted += 1, + LeaseOutcome::AtLimit { .. } => at_limit += 1, + } + } + + assert_eq!( + admitted, LIMIT as usize, + "{ATTEMPTS} racing acquires admitted {admitted}, limit is {LIMIT}" + ); + assert_eq!(at_limit, ATTEMPTS - LIMIT as usize); + + let live: i64 = pool + .get() + .await + .expect("connection") + .query_one( + "SELECT COUNT(*) FROM concurrency_leases WHERE organization_id = $1", + &[&org.org_id], + ) + .await + .expect("lease count") + .get(0); + assert_eq!(live, i64::from(LIMIT), "rows written must match admissions"); +} + +#[tokio::test] +async fn released_leases_free_capacity_again() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = $1 WHERE id = $2", + &[&LIMIT, &org.org_id], + ) + .await + .expect("rate limit update"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + let mut held = Vec::new(); + for _ in 0..LIMIT { + let lease_id = Uuid::new_v4(); + let outcome = repository + .try_acquire( + lease_id, + org.org_id, + model.id, + "instance-a", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + assert_eq!(outcome, LeaseOutcome::Admitted); + held.push(lease_id); + } + + let blocked = repository + .try_acquire( + Uuid::new_v4(), + org.org_id, + model.id, + "instance-b", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + assert!(matches!(blocked, LeaseOutcome::AtLimit { .. })); + + repository.release(&held[..1]).await.expect("release"); + + let readmitted = repository + .try_acquire( + Uuid::new_v4(), + org.org_id, + model.id, + "instance-b", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + assert_eq!( + readmitted, + LeaseOutcome::Admitted, + "releasing a lease must free exactly one slot" + ); +} + +/// An expired lease must not keep counting, or a replica that died holding +/// leases would lock the organization out permanently. +#[tokio::test] +async fn expired_leases_stop_counting() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = 1 WHERE id = $1", + &[&org.org_id], + ) + .await + .expect("rate limit update"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + let outcome = repository + .try_acquire( + Uuid::new_v4(), + org.org_id, + model.id, + "instance-dead", + DEFAULT_LIMIT, + Duration::from_secs(1), + ) + .await + .expect("acquire"); + assert_eq!(outcome, LeaseOutcome::Admitted); + + tokio::time::sleep(Duration::from_millis(1200)).await; + + let after_expiry = repository + .try_acquire( + Uuid::new_v4(), + org.org_id, + model.id, + "instance-live", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + assert_eq!( + after_expiry, + LeaseOutcome::Admitted, + "a lease past its TTL must not hold capacity" + ); +} diff --git a/crates/inference_providers/src/mock.rs b/crates/inference_providers/src/mock.rs index 1b14b9f34..7f45ef2f0 100644 --- a/crates/inference_providers/src/mock.rs +++ b/crates/inference_providers/src/mock.rs @@ -21,7 +21,7 @@ use bytes::Bytes; use futures_util::stream; use sha2::{Digest, Sha256}; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::{Mutex, RwLock}; /// Lightweight PII detector used only by [`MockProvider::privacy_classify_raw`] @@ -222,6 +222,7 @@ pub struct ResponseTemplate { /// request's model param — simulates external backends that answer with /// their upstream model name (`provider_config.model_name` overrides). model_override: Option, + hold: Option, } impl ResponseTemplate { @@ -235,9 +236,16 @@ impl ResponseTemplate { tool_calls: None, cache_tokens: None, model_override: None, + hold: None, } } + /// Hold the request in flight for `duration` before responding. + pub fn with_hold(mut self, duration: Duration) -> Self { + self.hold = Some(duration); + self + } + /// Echo `model` in responses instead of the request's model param /// (simulates upstream model-name overrides on external providers). pub fn with_model(mut self, model: impl Into) -> Self { @@ -1013,6 +1021,10 @@ impl crate::InferenceProvider for MockProvider { .unwrap_or_else(|| config.default_response.clone()) }; + if let Some(hold) = response_template.hold { + tokio::time::sleep(hold).await; + } + // Calculate input tokens from messages (rough estimate: 1 word ≈ 1 token) let input_tokens: i32 = params .messages @@ -1133,6 +1145,10 @@ impl crate::InferenceProvider for MockProvider { .unwrap_or_else(|| config.default_response.clone()) }; + if let Some(hold) = response_template.hold { + tokio::time::sleep(hold).await; + } + // Calculate input tokens from messages (rough estimate: 1 word ≈ 1 token) let input_tokens: i32 = params .messages diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 47e29fef2..f4ace7f57 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -9,6 +9,7 @@ use inference_providers::{ChatMessage, MessageRole, SSEEvent, StreamChunk, Strea use moka::future::Cache; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; +use tokio::sync::mpsc; use uuid::Uuid; // Create a new stream that intercepts messages, but passes the original ones through @@ -95,7 +96,7 @@ where total_itl_ms: f64, // Pre-allocated low-cardinality metric tags (for Datadog/OTLP) metric_tags: Vec, - concurrent_counter: Option>, + concurrent_slot: Option, /// Last received usage stats from streaming chunks last_usage_stats: Option, /// Last chat ID from streaming chunks (for attestation and inference_id) @@ -507,9 +508,8 @@ where S: Stream> + Unpin, { fn drop(&mut self) { - // Decrement concurrent counter if present - if let Some(counter) = &self.concurrent_counter { - counter.fetch_sub(1, Ordering::Release); + if let Some(slot) = &self.concurrent_slot { + slot.release(); } // Always record usage in Drop (async, fire-and-forget) @@ -517,32 +517,52 @@ where } } +#[derive(Clone)] +pub(crate) enum ConcurrentSlot { + Local(Arc), + Lease { + id: Uuid, + release: mpsc::UnboundedSender, + }, +} + +impl ConcurrentSlot { + fn release(&self) { + match self { + Self::Local(counter) => { + counter.fetch_sub(1, Ordering::Release); + } + Self::Lease { id, release } => { + let _ = release.send(*id); + } + } + } +} + /// RAII guard for concurrent request slots. /// Automatically releases the slot when dropped, ensuring proper cleanup even if the request panics. -/// Use `disarm()` to take ownership of the counter without decrementing (e.g., to transfer it -/// to an `InterceptStream` that will handle decrement on drop). +/// Use `disarm()` to take ownership of the slot without releasing it (e.g., to transfer it +/// to an `InterceptStream` that will handle release on drop). struct ConcurrentSlotGuard { - counter: Option>, + slot: Option, } impl ConcurrentSlotGuard { - fn new(counter: Arc) -> Self { - Self { - counter: Some(counter), - } + fn new(slot: ConcurrentSlot) -> Self { + Self { slot: Some(slot) } } - /// Disarm the guard and return the counter without decrementing. - /// Used when transferring counter ownership to `InterceptStream`. - fn disarm(&mut self) -> Option> { - self.counter.take() + /// Disarm the guard and return the slot without releasing it. + /// Used when transferring slot ownership to `InterceptStream`. + fn disarm(&mut self) -> Option { + self.slot.take() } } impl Drop for ConcurrentSlotGuard { fn drop(&mut self) { - if let Some(counter) = &self.counter { - counter.fetch_sub(1, Ordering::Release); + if let Some(slot) = &self.slot { + slot.release(); } } } @@ -559,6 +579,21 @@ pub struct CompletionServiceImpl { org_concurrent_limits: Cache, /// Repository for fetching organization concurrent limits organization_limit_repository: Arc, + fleet_concurrency: Option, +} + +enum LeaseAdmission { + AtLimit(ports::CompletionError), + /// The lease store could not answer, so admission falls back to the + /// per-process limit rather than rejecting the request. + Unavailable(anyhow::Error), +} + +struct FleetConcurrency { + repository: Arc, + release: mpsc::UnboundedSender, + instance_id: String, + ttl: Duration, } /// TTL for organization concurrent limit cache (5 minutes) @@ -674,9 +709,44 @@ impl CompletionServiceImpl { concurrent_limit: DEFAULT_CONCURRENT_LIMIT, org_concurrent_limits, organization_limit_repository, + fleet_concurrency: None, } } + /// Released leases are deleted by a background task because `Drop` is + /// synchronous and the delete is not. + pub fn with_fleet_concurrency( + mut self, + repository: Arc, + instance_id: String, + ttl: Duration, + ) -> Self { + let (release, mut released) = mpsc::unbounded_channel::(); + let releaser = repository.clone(); + + tokio::spawn(async move { + let mut batch = Vec::new(); + while released.recv_many(&mut batch, 128).await > 0 { + if let Err(error) = releaser.release(&batch).await { + tracing::warn!( + released = batch.len(), + error = %error, + "Failed to release concurrency leases; they will expire instead" + ); + } + batch.clear(); + } + }); + + self.fleet_concurrency = Some(FleetConcurrency { + repository, + release, + instance_id, + ttl, + }); + self + } + /// Extract tools and tool_choice from the extra HashMap if present and /// parseable as the typed `ToolDefinition` / `ToolChoice` shapes. /// @@ -1237,12 +1307,79 @@ impl CompletionServiceImpl { .collect() } + async fn try_acquire_lease( + &self, + fleet: &FleetConcurrency, + organization_id: Uuid, + model_id: Uuid, + model_name: &str, + ) -> Result { + let lease_id = Uuid::new_v4(); + let outcome = fleet + .repository + .try_acquire( + lease_id, + organization_id, + model_id, + &fleet.instance_id, + self.concurrent_limit, + fleet.ttl, + ) + .await + .map_err(LeaseAdmission::Unavailable)?; + + match outcome { + ports::LeaseOutcome::Admitted => Ok(ConcurrentSlot::Lease { + id: lease_id, + release: fleet.release.clone(), + }), + ports::LeaseOutcome::AtLimit { limit, in_flight } => { + tracing::warn!( + organization_id = %organization_id, + model_id = %model_id, + model_name = %model_name, + current_count = in_flight, + limit = limit, + "Organization concurrent request limit exceeded for model across the fleet" + ); + let message = format!( + "Concurrent request limit exceeded for model {model_name}. Organization limit: {limit} concurrent requests per model." + ); + self.record_error( + &ports::CompletionError::RateLimitExceeded(message.clone()), + Some(model_name), + ); + Err(LeaseAdmission::AtLimit( + ports::CompletionError::RateLimitExceeded(message), + )) + } + } + } + async fn try_acquire_concurrent_slot( &self, organization_id: Uuid, model_id: Uuid, model_name: &str, - ) -> Result, ports::CompletionError> { + ) -> Result { + if let Some(fleet) = &self.fleet_concurrency { + match self + .try_acquire_lease(fleet, organization_id, model_id, model_name) + .await + { + Ok(slot) => return Ok(slot), + Err(LeaseAdmission::AtLimit(error)) => return Err(error), + Err(LeaseAdmission::Unavailable(error)) => { + tracing::warn!( + organization_id = %organization_id, + model_id = %model_id, + error = %error, + "Fleet concurrency unavailable, falling back to per-process limit" + ); + } + } + } + // Get the dynamic limit for this organization (cached with 5-min TTL) let limit = self.get_org_concurrent_limit(organization_id).await; @@ -1278,7 +1415,7 @@ impl CompletionServiceImpl { .compare_exchange_weak(current, current + 1, Ordering::AcqRel, Ordering::Acquire) .is_ok() { - return Ok(counter); + return Ok(ConcurrentSlot::Local(counter)); } } } @@ -1297,7 +1434,7 @@ impl CompletionServiceImpl { inference_type: crate::usage::ports::InferenceType, service_start_time: Instant, provider_start_time: Instant, - concurrent_counter: Option>, + concurrent_slot: Option, response_id: Option, attestation_supported: bool, store_provider_chat_signature: bool, @@ -1336,7 +1473,7 @@ impl CompletionServiceImpl { last_token_time: None, total_itl_ms: 0.0, metric_tags, - concurrent_counter, + concurrent_slot, last_usage_stats: None, last_chat_id: None, stream_completed: false, @@ -1457,13 +1594,13 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { } Self::apply_deepseek_v4_flash_thinking_compat(canonical_name, &mut chat_params); - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model.id, canonical_name) .await?; // RAII guard protects against panics during stream creation. // On success, disarm and transfer counter ownership to InterceptStream. - let mut guard = ConcurrentSlotGuard::new(counter); + let mut guard = ConcurrentSlotGuard::new(slot); Self::reject_e2ee_if_unsupported( model.attestation_supported, @@ -1643,12 +1780,12 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { Self::apply_deepseek_v4_flash_thinking_compat(canonical_name, &mut chat_params); let organization_id = request.organization_id; - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model.id, canonical_name) .await?; // RAII guard ensures slot is released on drop (panic, error, or success) - let _guard = ConcurrentSlotGuard::new(counter); + let _guard = ConcurrentSlotGuard::new(slot); Self::reject_e2ee_if_unsupported( model.attestation_supported, @@ -1815,12 +1952,12 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { request_hash: String, ) -> Result { // Acquire concurrent request slot to enforce organization limits - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model_id, model_name) .await?; // RAII guard ensures slot is released on drop (panic, error, or success) - let _guard = ConcurrentSlotGuard::new(counter); + let _guard = ConcurrentSlotGuard::new(slot); // Call inference provider pool with timeout protection let timeout_duration = std::time::Duration::from_secs(120); // 2 minute timeout for audio @@ -1879,12 +2016,12 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { params: inference_providers::RerankParams, ) -> Result { // Acquire concurrent request slot to enforce organization limits - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model_id, model_name) .await?; // Create RAII guard to ensure slot is released on drop (panic, error, or success) - let _guard = ConcurrentSlotGuard::new(counter); + let _guard = ConcurrentSlotGuard::new(slot); // Call inference provider pool // The guard will automatically release the slot when this function returns or panics @@ -1931,10 +2068,10 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { body: bytes::Bytes, extra: std::collections::HashMap, ) -> Result { - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model_id, model_name) .await?; - let _guard = ConcurrentSlotGuard::new(counter); + let _guard = ConcurrentSlotGuard::new(slot); self.inference_provider_pool .embeddings(model_name, body, extra) @@ -1979,10 +2116,10 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { body: bytes::Bytes, extra: std::collections::HashMap, ) -> Result { - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model_id, model_name) .await?; - let _guard = ConcurrentSlotGuard::new(counter); + let _guard = ConcurrentSlotGuard::new(slot); self.inference_provider_pool .privacy_classify(model_name, body, extra) @@ -2028,12 +2165,12 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { params: inference_providers::ScoreParams, ) -> Result { // Acquire concurrent request slot to enforce organization limits - let counter = self + let slot = self .try_acquire_concurrent_slot(organization_id, model_id, model_name) .await?; // Create RAII guard to ensure slot is released on drop (panic, error, or success) - let _guard = ConcurrentSlotGuard::new(counter); + let _guard = ConcurrentSlotGuard::new(slot); // Call inference provider pool // The guard will automatically release the slot when this function returns or panics @@ -2191,7 +2328,7 @@ mod tests { last_token_time: None, total_itl_ms: 0.0, metric_tags, - concurrent_counter: None, + concurrent_slot: None, last_usage_stats: None, last_chat_id: None, stream_completed: false, @@ -2321,7 +2458,7 @@ mod tests { last_token_time: None, total_itl_ms: 0.0, metric_tags: CompletionServiceImpl::create_metric_tags("test-model"), - concurrent_counter: None, + concurrent_slot: None, last_usage_stats: None, last_chat_id: None, stream_completed: false, @@ -2461,7 +2598,7 @@ mod tests { last_token_time: None, total_itl_ms: 0.0, metric_tags, - concurrent_counter: None, + concurrent_slot: None, last_usage_stats: None, last_chat_id: None, stream_completed: false, @@ -2584,7 +2721,7 @@ mod tests { last_token_time: None, total_itl_ms: 0.0, metric_tags, - concurrent_counter: None, + concurrent_slot: None, last_usage_stats: None, last_chat_id: None, stream_completed: false, @@ -2791,7 +2928,7 @@ mod tests { last_token_time: None, total_itl_ms: 0.0, metric_tags: vec![], - concurrent_counter: Some(counter.clone()), + concurrent_slot: Some(ConcurrentSlot::Local(counter.clone())), last_usage_stats: None, last_chat_id: None, stream_completed: false, diff --git a/crates/services/src/completions/ports.rs b/crates/services/src/completions/ports.rs index 803b61f10..492f64241 100644 --- a/crates/services/src/completions/ports.rs +++ b/crates/services/src/completions/ports.rs @@ -3,6 +3,7 @@ use crate::UserId; use async_trait::async_trait; use inference_providers::StreamingResult; use serde::{Deserialize, Serialize}; +use std::time::Duration; use uuid::Uuid; /// Default concurrent request limit per organization per model @@ -129,6 +130,30 @@ pub trait OrganizationConcurrentLimitRepository: Send + Sync { async fn get_concurrent_limit(&self, org_id: Uuid) -> Result, anyhow::Error>; } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LeaseOutcome { + Admitted, + AtLimit { limit: u32, in_flight: i64 }, +} + +/// Repository trait for fleet-wide concurrency leases +/// Used by CompletionService so replicas admit against one shared count +#[async_trait] +pub trait ConcurrencyLeaseRepository: Send + Sync { + /// Record a lease if the organization is below its limit for the model + async fn try_acquire( + &self, + lease_id: Uuid, + organization_id: Uuid, + model_id: Uuid, + instance_id: &str, + default_limit: u32, + ttl: Duration, + ) -> Result; + + async fn release(&self, lease_ids: &[Uuid]) -> Result<(), anyhow::Error>; +} + #[async_trait] pub trait CompletionServiceTrait: Send + Sync { /// Create a streaming completion From 07a09c332fd929a396acd125bfff3da162757c94 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Wed, 12 Aug 2026 22:27:31 -0400 Subject: [PATCH 2/8] fix(concurrency): renew in place and bound the degraded path --- .gitignore | 3 + crates/config/src/types.rs | 4 +- .../src/repositories/concurrency_lease.rs | 99 ++++- crates/database/tests/concurrency_leases.rs | 146 +++++++- crates/services/src/completions/mod.rs | 340 +++++++++++++++++- crates/services/src/completions/ports.rs | 24 ++ 6 files changed, 598 insertions(+), 18 deletions(-) diff --git a/.gitignore b/.gitignore index 25ad8cc06..ecb03faf4 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,6 @@ oci.tar .DS_Store secret repro_*.sh +# rustc reads -C incremental= as a directory, so the reproducibility flag in +# .cargo/config.toml writes incremental artifacts to ./false +/false/ diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index 01fd55d53..3ff478057 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -52,7 +52,7 @@ impl Default for FleetConcurrencyConfig { fn default() -> Self { Self { mode: FleetConcurrencyMode::Off, - lease_ttl_seconds: 30, + lease_ttl_seconds: 60, instance_id: String::new(), } } @@ -75,7 +75,7 @@ impl FleetConcurrencyConfig { .ok() .and_then(|value| value.parse::().ok()) .filter(|seconds| *seconds > 0) - .unwrap_or(30), + .unwrap_or(60), instance_id: env::var("FLEET_CONCURRENCY_INSTANCE_ID") .ok() .filter(|value| !value.is_empty()) diff --git a/crates/database/src/repositories/concurrency_lease.rs b/crates/database/src/repositories/concurrency_lease.rs index 019203c86..aea946b4e 100644 --- a/crates/database/src/repositories/concurrency_lease.rs +++ b/crates/database/src/repositories/concurrency_lease.rs @@ -4,7 +4,7 @@ use crate::retry_db; use anyhow::{Context, Result}; use async_trait::async_trait; use services::common::RepositoryError; -use services::completions::ports::{ConcurrencyLeaseRepository, LeaseOutcome}; +use services::completions::ports::{ConcurrencyLeaseRepository, HeldLease, LeaseOutcome}; use std::time::Duration; use uuid::Uuid; @@ -112,6 +112,103 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { Ok(outcome) } + async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result<()> { + if lease_ids.is_empty() { + return Ok(()); + } + + let ttl_seconds = ttl.as_secs() as f64; + + retry_db!("renew_concurrency_leases", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + client + .execute( + r#" + UPDATE concurrency_leases + SET expires_at = NOW() + make_interval(secs => $2) + WHERE id = ANY($1) + "#, + &[&lease_ids, &ttl_seconds], + ) + .await + .map_err(map_db_error) + })?; + + Ok(()) + } + + async fn persist(&self, leases: &[HeldLease], instance_id: &str, ttl: Duration) -> Result<()> { + if leases.is_empty() { + return Ok(()); + } + + let ids: Vec = leases.iter().map(|lease| lease.id).collect(); + let organization_ids: Vec = + leases.iter().map(|lease| lease.organization_id).collect(); + let model_ids: Vec = leases.iter().map(|lease| lease.model_id).collect(); + let ttl_seconds = ttl.as_secs() as f64; + + retry_db!("persist_concurrency_leases", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + client + .execute( + r#" + INSERT INTO concurrency_leases + (id, organization_id, model_id, instance_id, expires_at) + SELECT held.id, held.organization_id, held.model_id, $4, + NOW() + make_interval(secs => $5) + FROM UNNEST($1::uuid[], $2::uuid[], $3::uuid[]) + AS held(id, organization_id, model_id) + ON CONFLICT (id) DO UPDATE SET expires_at = EXCLUDED.expires_at + "#, + &[ + &ids, + &organization_ids, + &model_ids, + &instance_id, + &ttl_seconds, + ], + ) + .await + .map_err(map_db_error) + })?; + + Ok(()) + } + + async fn sweep_expired(&self) -> Result { + let removed = retry_db!("sweep_expired_concurrency_leases", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + client + .execute( + "DELETE FROM concurrency_leases WHERE expires_at < NOW()", + &[], + ) + .await + .map_err(map_db_error) + })?; + + Ok(removed) + } + async fn release(&self, lease_ids: &[Uuid]) -> Result<()> { if lease_ids.is_empty() { return Ok(()); diff --git a/crates/database/tests/concurrency_leases.rs b/crates/database/tests/concurrency_leases.rs index 772a6d641..56339a5da 100644 --- a/crates/database/tests/concurrency_leases.rs +++ b/crates/database/tests/concurrency_leases.rs @@ -1,7 +1,7 @@ mod support; use database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository; -use services::completions::ports::{ConcurrencyLeaseRepository, LeaseOutcome}; +use services::completions::ports::{ConcurrencyLeaseRepository, HeldLease, LeaseOutcome}; use std::sync::Arc; use std::time::Duration; use uuid::Uuid; @@ -157,6 +157,150 @@ async fn released_leases_free_capacity_again() { ); } +/// Renewal is what lets a request outlive the TTL. Without it a request longer +/// than the TTL frees its own slot while still running. +#[tokio::test] +async fn renewal_keeps_a_lease_past_its_original_ttl() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = 1 WHERE id = $1", + &[&org.org_id], + ) + .await + .expect("rate limit update"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + let lease_id = Uuid::new_v4(); + let outcome = repository + .try_acquire( + lease_id, + org.org_id, + model.id, + "instance-a", + DEFAULT_LIMIT, + Duration::from_secs(1), + ) + .await + .expect("acquire"); + assert_eq!(outcome, LeaseOutcome::Admitted); + + repository.renew(&[lease_id], TTL).await.expect("renew"); + + tokio::time::sleep(Duration::from_millis(1200)).await; + + let blocked = repository + .try_acquire( + Uuid::new_v4(), + org.org_id, + model.id, + "instance-b", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + assert!( + matches!(blocked, LeaseOutcome::AtLimit { .. }), + "a renewed lease must still hold its slot after the original TTL" + ); +} + +/// Renewal must never recreate a row. A request that finished during the round +/// trip has already been deleted, and an insert here would strand a slot that +/// nothing holds, renews or releases. +#[tokio::test] +async fn renewal_cannot_recreate_a_released_lease() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + let lease_id = Uuid::new_v4(); + repository + .try_acquire( + lease_id, + org.org_id, + model.id, + "instance-a", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + + repository.release(&[lease_id]).await.expect("release"); + repository.renew(&[lease_id], TTL).await.expect("renew"); + + let live: i64 = pool + .get() + .await + .expect("connection") + .query_one( + "SELECT COUNT(*) FROM concurrency_leases WHERE id = $1", + &[&lease_id], + ) + .await + .expect("count") + .get(0); + assert_eq!( + live, 0, + "renewing a released lease must not bring it back to life" + ); +} + +/// Leases admitted while the store was unreachable are written back, so they +/// start counting against the fleet instead of only this replica. +#[tokio::test] +async fn pending_leases_are_written_back() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + let lease_id = Uuid::new_v4(); + let pending = [HeldLease { + id: lease_id, + organization_id: org.org_id, + model_id: model.id, + }]; + + repository + .persist(&pending, "instance-a", TTL) + .await + .expect("persist"); + + let live: i64 = pool + .get() + .await + .expect("connection") + .query_one( + "SELECT COUNT(*) FROM concurrency_leases WHERE id = $1 AND expires_at > NOW()", + &[&lease_id], + ) + .await + .expect("count") + .get(0); + assert_eq!(live, 1, "a degraded-path lease must reach the store"); +} + /// An expired lease must not keep counting, or a replica that died holding /// leases would lock the organization out permanently. #[tokio::test] diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index f4ace7f57..c797a7413 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -7,6 +7,7 @@ use crate::responses::models::ResponseId; use crate::usage::{RecordUsageServiceRequest, UsageServiceTrait}; use inference_providers::{ChatMessage, MessageRole, SSEEvent, StreamChunk, StreamingResult}; use moka::future::Cache; +use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; use tokio::sync::mpsc; @@ -522,6 +523,9 @@ pub(crate) enum ConcurrentSlot { Local(Arc), Lease { id: Uuid, + organization_id: Uuid, + model_id: Uuid, + held: Arc, release: mpsc::UnboundedSender, }, } @@ -532,8 +536,21 @@ impl ConcurrentSlot { Self::Local(counter) => { counter.fetch_sub(1, Ordering::Release); } - Self::Lease { id, release } => { - let _ = release.send(*id); + Self::Lease { + id, + organization_id, + model_id, + held, + release, + } => { + held.remove(*organization_id, *model_id, *id); + if release.send(*id).is_err() { + tracing::error!( + lease_id = %id, + "Concurrency lease release channel is closed; the lease will \ + hold capacity until it expires" + ); + } } } } @@ -577,6 +594,7 @@ pub struct CompletionServiceImpl { concurrent_limit: u32, /// Cache for per-organization concurrent limits (5-minute TTL) org_concurrent_limits: Cache, + last_known_limits: Cache, /// Repository for fetching organization concurrent limits organization_limit_repository: Arc, fleet_concurrency: Option, @@ -589,16 +607,143 @@ enum LeaseAdmission { Unavailable(anyhow::Error), } +/// Leases this replica currently holds. Renewal reads it to keep them alive, +/// release removes from it, and admission counts it when the lease store is +/// unreachable, so the degraded path can never admit on top of live leases. +#[derive(Default)] +pub(crate) struct HeldLeases { + by_model: std::sync::Mutex>>, +} + +/// Whether a held lease already has a row in the lease store. Renewal updates +/// stored leases and inserts pending ones, so an update can never recreate a +/// row that a release deleted while the round trip was outstanding. +#[derive(Clone, Copy, PartialEq, Eq)] +enum LeaseState { + Stored, + PendingWrite, +} + +impl HeldLeases { + fn insert(&self, organization_id: Uuid, model_id: Uuid, lease_id: Uuid, state: LeaseState) { + let mut guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + guard + .entry((organization_id, model_id)) + .or_default() + .insert(lease_id, state); + } + + /// Ids no longer held are skipped, so a lease released mid-write is not + /// resurrected in the registry either. + fn mark_stored(&self, leases: &[ports::HeldLease]) { + let mut guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + for lease in leases { + if let Some(ids) = guard.get_mut(&(lease.organization_id, lease.model_id)) { + if let Some(state) = ids.get_mut(&lease.id) { + *state = LeaseState::Stored; + } + } + } + } + + fn remove(&self, organization_id: Uuid, model_id: Uuid, lease_id: Uuid) { + let mut guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(ids) = guard.get_mut(&(organization_id, model_id)) { + ids.remove(&lease_id); + if ids.is_empty() { + guard.remove(&(organization_id, model_id)); + } + } + } + + fn count(&self, organization_id: Uuid, model_id: Uuid) -> usize { + let guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + guard + .get(&(organization_id, model_id)) + .map_or(0, |ids| ids.len()) + } + + /// Record a lease only if the model is below `limit`, deciding and + /// inserting under one lock so concurrent degraded admissions cannot both + /// observe the same count and both proceed. + fn insert_below( + &self, + organization_id: Uuid, + model_id: Uuid, + lease_id: Uuid, + limit: usize, + ) -> Result<(), usize> { + let mut guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + let ids = guard.entry((organization_id, model_id)).or_default(); + if ids.len() >= limit { + return Err(ids.len()); + } + ids.insert(lease_id, LeaseState::PendingWrite); + Ok(()) + } + + /// Held leases split by whether the store already has them. + fn snapshot(&self) -> (Vec, Vec) { + let guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + let mut stored = Vec::new(); + let mut pending = Vec::new(); + for ((organization_id, model_id), ids) in guard.iter() { + for (id, state) in ids { + let lease = ports::HeldLease { + id: *id, + organization_id: *organization_id, + model_id: *model_id, + }; + match state { + LeaseState::Stored => stored.push(lease), + LeaseState::PendingWrite => pending.push(lease), + } + } + } + (stored, pending) + } + + fn held_ids(&self) -> HashSet { + let guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); + guard.values().flat_map(|ids| ids.keys()).copied().collect() + } +} + struct FleetConcurrency { repository: Arc, + held: Arc, release: mpsc::UnboundedSender, instance_id: String, ttl: Duration, } +/// Renewals attempted within one lease TTL. Three gives two spare attempts +/// before a live request's lease could lapse. +const LEASE_RENEWALS_PER_TTL: u32 = 3; + +/// Each of these loops silently stops enforcing something if it exits, so a +/// death is reported rather than discarded. +fn supervise(task: &'static str, handle: tokio::task::JoinHandle<()>) { + tokio::spawn(async move { + match handle.await { + Ok(()) => { + tracing::warn!(task = task, "Fleet concurrency task stopped") + } + Err(error) => tracing::error!( + task = task, + error = %error, + "Fleet concurrency task died; limits are no longer maintained" + ), + } + }); +} + /// TTL for organization concurrent limit cache (5 minutes) const ORG_LIMIT_CACHE_TTL_SECS: u64 = 300; +/// How long a limit that was actually read stays usable as a fallback (24h). +const LAST_KNOWN_LIMIT_TTL_SECS: u64 = 86_400; + /// TTL for concurrent count cache entries (10 minutes). /// Safety net: if a counter gets stuck (e.g., due to a panic or proxy not propagating /// client disconnection), the entry expires and is replaced with a fresh zero counter. @@ -708,6 +853,10 @@ impl CompletionServiceImpl { concurrent_counts, concurrent_limit: DEFAULT_CONCURRENT_LIMIT, org_concurrent_limits, + last_known_limits: Cache::builder() + .time_to_live(Duration::from_secs(LAST_KNOWN_LIMIT_TTL_SECS)) + .max_capacity(10_000) + .build(), organization_limit_repository, fleet_concurrency: None, } @@ -724,7 +873,7 @@ impl CompletionServiceImpl { let (release, mut released) = mpsc::unbounded_channel::(); let releaser = repository.clone(); - tokio::spawn(async move { + let drain = tokio::spawn(async move { let mut batch = Vec::new(); while released.recv_many(&mut batch, 128).await > 0 { if let Err(error) = releaser.release(&batch).await { @@ -738,8 +887,90 @@ impl CompletionServiceImpl { } }); + supervise("release", drain); + + let held = Arc::new(HeldLeases::default()); + + let renewer = repository.clone(); + let renewing = held.clone(); + let renew_instance = instance_id.clone(); + let renew_every = ttl / LEASE_RENEWALS_PER_TTL; + let renew_task = tokio::spawn(async move { + let mut ticker = tokio::time::interval(renew_every); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + ticker.tick().await; + let (stored, pending) = renewing.snapshot(); + + let stored_ids: Vec = stored.iter().map(|lease| lease.id).collect(); + if let Err(error) = renewer.renew(&stored_ids, ttl).await { + tracing::warn!( + leases = stored_ids.len(), + error = %error, + "Failed to renew concurrency leases; in-flight requests continue" + ); + } + + if pending.is_empty() { + continue; + } + if let Err(error) = renewer.persist(&pending, &renew_instance, ttl).await { + tracing::warn!( + leases = pending.len(), + error = %error, + "Failed to store leases admitted while the lease store was down" + ); + continue; + } + renewing.mark_stored(&pending); + + // A pending lease released during the insert is written with no + // holder, so nothing would renew or release it again. + let held = renewing.held_ids(); + let orphaned: Vec = pending + .iter() + .map(|lease| lease.id) + .filter(|id| !held.contains(id)) + .collect(); + if !orphaned.is_empty() { + if let Err(error) = renewer.release(&orphaned).await { + tracing::warn!( + leases = orphaned.len(), + error = %error, + "Failed to drop leases released while being stored; they will expire" + ); + } + } + } + }); + + supervise("renew", renew_task); + + let sweeper = repository.clone(); + let sweep_task = tokio::spawn(async move { + let mut ticker = tokio::time::interval(ttl); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + ticker.tick().await; + match sweeper.sweep_expired().await { + Ok(0) => {} + Ok(removed) => tracing::info!( + removed = removed, + "Reclaimed concurrency leases whose holder stopped renewing" + ), + Err(error) => tracing::warn!( + error = %error, + "Failed to sweep expired concurrency leases" + ), + } + } + }); + + supervise("sweep", sweep_task); + self.fleet_concurrency = Some(FleetConcurrency { repository, + held, release, instance_id, ttl, @@ -904,22 +1135,43 @@ impl CompletionServiceImpl { let default_limit = self.concurrent_limit; let repo = self.organization_limit_repository.clone(); - self.org_concurrent_limits - .get_with(organization_id, async move { + // optionally_get_with does not cache a None, so a failed lookup falls back + // for this request only instead of pinning the default for the whole TTL. + let limit = self + .org_concurrent_limits + .optionally_get_with(organization_id, async move { match repo.get_concurrent_limit(organization_id).await { - Ok(Some(limit)) if limit > 0 => limit, - Ok(_) => default_limit, // Use default if NULL or 0 + Ok(Some(limit)) if limit > 0 => Some(limit), + Ok(_) => Some(default_limit), // Use default if NULL or 0 Err(e) => { tracing::warn!( organization_id = %organization_id, error = %e, "Failed to fetch org concurrent limit, using default" ); - default_limit + None } } }) - .await + .await; + + match limit { + Some(limit) => { + self.last_known_limits.insert(organization_id, limit).await; + limit + } + None => self + .last_known_org_concurrent_limit(organization_id) + .await + .unwrap_or(default_limit), + } + } + + /// Last limit actually read for this organization. Outliving the lookup + /// cache matters because an outage longer than that TTL would otherwise + /// leave the degraded path widening the cap to the global default. + async fn last_known_org_concurrent_limit(&self, organization_id: Uuid) -> Option { + self.last_known_limits.get(&organization_id).await } /// Create low-cardinality metric tags for a request @@ -1307,6 +1559,55 @@ impl CompletionServiceImpl { .collect() } + /// Admission while the lease store is unreachable. Counting this replica's + /// own leases degrades to the per-process behaviour that shipped before + /// fleet limits existed, and never admits on top of leases already live. + /// The lease is recorded locally only; the next renewal writes it back. + async fn admit_from_held_leases( + &self, + fleet: &FleetConcurrency, + organization_id: Uuid, + model_id: Uuid, + model_name: &str, + ) -> Result { + let limit = self + .last_known_org_concurrent_limit(organization_id) + .await + .unwrap_or(self.concurrent_limit); + let lease_id = Uuid::new_v4(); + + match fleet + .held + .insert_below(organization_id, model_id, lease_id, limit as usize) + { + Ok(()) => Ok(ConcurrentSlot::Lease { + id: lease_id, + organization_id, + model_id, + held: fleet.held.clone(), + release: fleet.release.clone(), + }), + Err(in_flight) => { + tracing::warn!( + organization_id = %organization_id, + model_id = %model_id, + model_name = %model_name, + current_count = in_flight, + limit = limit, + "Organization concurrent request limit exceeded for model on this replica" + ); + let message = format!( + "Concurrent request limit exceeded for model {model_name}. Organization limit: {limit} concurrent requests per model." + ); + self.record_error( + &ports::CompletionError::RateLimitExceeded(message.clone()), + Some(model_name), + ); + Err(ports::CompletionError::RateLimitExceeded(message)) + } + } + } + async fn try_acquire_lease( &self, fleet: &FleetConcurrency, @@ -1329,10 +1630,18 @@ impl CompletionServiceImpl { .map_err(LeaseAdmission::Unavailable)?; match outcome { - ports::LeaseOutcome::Admitted => Ok(ConcurrentSlot::Lease { - id: lease_id, - release: fleet.release.clone(), - }), + ports::LeaseOutcome::Admitted => { + fleet + .held + .insert(organization_id, model_id, lease_id, LeaseState::Stored); + Ok(ConcurrentSlot::Lease { + id: lease_id, + organization_id, + model_id, + held: fleet.held.clone(), + release: fleet.release.clone(), + }) + } ports::LeaseOutcome::AtLimit { limit, in_flight } => { tracing::warn!( organization_id = %organization_id, @@ -1374,8 +1683,11 @@ impl CompletionServiceImpl { organization_id = %organization_id, model_id = %model_id, error = %error, - "Fleet concurrency unavailable, falling back to per-process limit" + "Fleet concurrency unavailable, falling back to this replica's leases" ); + return self + .admit_from_held_leases(fleet, organization_id, model_id, model_name) + .await; } } } diff --git a/crates/services/src/completions/ports.rs b/crates/services/src/completions/ports.rs index 492f64241..f94c29785 100644 --- a/crates/services/src/completions/ports.rs +++ b/crates/services/src/completions/ports.rs @@ -136,6 +136,13 @@ pub enum LeaseOutcome { AtLimit { limit: u32, in_flight: i64 }, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct HeldLease { + pub id: Uuid, + pub organization_id: Uuid, + pub model_id: Uuid, +} + /// Repository trait for fleet-wide concurrency leases /// Used by CompletionService so replicas admit against one shared count #[async_trait] @@ -152,6 +159,23 @@ pub trait ConcurrencyLeaseRepository: Send + Sync { ) -> Result; async fn release(&self, lease_ids: &[Uuid]) -> Result<(), anyhow::Error>; + + /// Extend leases the store already has. Updating in place means a lease + /// released while this call was in flight stays deleted instead of being + /// recreated with no holder. + async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result<(), anyhow::Error>; + + /// Write leases the store does not have yet, which happens when a request + /// was admitted while it was unreachable. Returns the leases now stored. + async fn persist( + &self, + leases: &[HeldLease], + instance_id: &str, + ttl: Duration, + ) -> Result<(), anyhow::Error>; + + /// Remove leases whose holder stopped renewing them. + async fn sweep_expired(&self) -> Result; } #[async_trait] From 389ff684a647719006fa07b183480b1a471780fd Mon Sep 17 00:00:00 2001 From: neo-sky Date: Thu, 13 Aug 2026 23:55:50 -0400 Subject: [PATCH 3/8] fix(concurrency): zero limits, gauge metric, bounded sweep A rate limit of zero means unset everywhere else, so reading it as a real limit locked an org out of every model. in_flight went out as a counter when it is a depth reading, and the background loops sent metrics with no environment tag. Sweep deletes in batches now, and an unknown FLEET_CONCURRENCY_MODE warns instead of quietly leaving the feature off. --- .gitignore | 3 - crates/api/tests/common/mod.rs | 5 +- crates/config/src/types.rs | 19 +- .../sql/V0073__add_concurrency_leases.sql | 3 - .../src/repositories/concurrency_lease.rs | 27 +- crates/database/tests/concurrency_leases.rs | 135 +++++- .../completions/fleet_concurrency_tests.rs | 390 ++++++++++++++++++ crates/services/src/completions/mod.rs | 264 ++++++++++-- crates/services/src/completions/ports.rs | 15 +- crates/services/src/metrics/consts.rs | 18 + 10 files changed, 804 insertions(+), 75 deletions(-) create mode 100644 crates/services/src/completions/fleet_concurrency_tests.rs diff --git a/.gitignore b/.gitignore index ecb03faf4..25ad8cc06 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,3 @@ oci.tar .DS_Store secret repro_*.sh -# rustc reads -C incremental= as a directory, so the reproducibility flag in -# .cargo/config.toml writes incremental artifacts to ./false -/false/ diff --git a/crates/api/tests/common/mod.rs b/crates/api/tests/common/mod.rs index e85cd2214..0eb1fe7f3 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -434,9 +434,8 @@ where (server, database) } -/// Build several independent servers that share one database, standing in for -/// the replicas of a deployed fleet. Each gets its own services and provider -/// pool, so anything they agree on has to travel through the database. +/// Independent servers sharing one database, standing in for fleet replicas. +/// Each gets its own services, so agreement has to travel through the database. pub async fn setup_test_fleet( instances: usize, mutate: F, diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index 3ff478057..cf07b893a 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -60,13 +60,20 @@ impl Default for FleetConcurrencyConfig { impl FleetConcurrencyConfig { pub fn from_env() -> Self { - let mode = match env::var("FLEET_CONCURRENCY_MODE") - .unwrap_or_default() - .to_ascii_lowercase() - .as_str() - { + let requested = env::var("FLEET_CONCURRENCY_MODE").unwrap_or_default(); + let mode = match requested.to_ascii_lowercase().as_str() { "enforce" => FleetConcurrencyMode::Enforce, - _ => FleetConcurrencyMode::Off, + "" | "off" => FleetConcurrencyMode::Off, + _ => { + // eprintln for the reason given on the CHUTES_MODELS warning + // below. Silence here reads as enforcement being on when it is + // not. + eprintln!( + "WARN: unrecognised FLEET_CONCURRENCY_MODE '{requested}', \ + falling back to off" + ); + FleetConcurrencyMode::Off + } }; Self { diff --git a/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql b/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql index 3d55473d8..f558729df 100644 --- a/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql +++ b/crates/database/src/migrations/sql/V0073__add_concurrency_leases.sql @@ -12,6 +12,3 @@ CREATE INDEX idx_concurrency_leases_org_model CREATE INDEX idx_concurrency_leases_expires_at ON concurrency_leases (expires_at); - -CREATE INDEX idx_concurrency_leases_instance - ON concurrency_leases (instance_id, expires_at); diff --git a/crates/database/src/repositories/concurrency_lease.rs b/crates/database/src/repositories/concurrency_lease.rs index aea946b4e..3dd7642f1 100644 --- a/crates/database/src/repositories/concurrency_lease.rs +++ b/crates/database/src/repositories/concurrency_lease.rs @@ -18,8 +18,11 @@ impl PostgresConcurrencyLeaseRepository { Self { pool } } + /// Zero and negative mean unset, as on the per-process path. Reading zero + /// as a real limit would reject every request for the organization. fn effective_limit(stored: Option, default_limit: u32) -> u32 { stored + .filter(|value| *value > 0) .and_then(|value| u32::try_from(value).ok()) .unwrap_or(default_limit) } @@ -106,20 +109,20 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { .map_err(map_db_error)?; tx.commit().await.map_err(map_db_error)?; - Ok(LeaseOutcome::Admitted) + Ok(LeaseOutcome::Admitted { limit }) })?; Ok(outcome) } - async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result<()> { + async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result> { if lease_ids.is_empty() { - return Ok(()); + return Ok(Vec::new()); } let ttl_seconds = ttl.as_secs() as f64; - retry_db!("renew_concurrency_leases", { + let rows = retry_db!("renew_concurrency_leases", { let client = self .pool .get() @@ -128,11 +131,12 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { .map_err(RepositoryError::PoolError)?; client - .execute( + .query( r#" UPDATE concurrency_leases SET expires_at = NOW() + make_interval(secs => $2) WHERE id = ANY($1) + RETURNING id "#, &[&lease_ids, &ttl_seconds], ) @@ -140,7 +144,7 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { .map_err(map_db_error) })?; - Ok(()) + Ok(rows.iter().map(|row| row.get("id")).collect()) } async fn persist(&self, leases: &[HeldLease], instance_id: &str, ttl: Duration) -> Result<()> { @@ -197,9 +201,18 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { .context("Failed to get database connection") .map_err(RepositoryError::PoolError)?; + // Bounded so a backlog cannot become one long delete holding row + // locks against live admissions; the next tick takes the rest. client .execute( - "DELETE FROM concurrency_leases WHERE expires_at < NOW()", + r#" + DELETE FROM concurrency_leases + WHERE id = ANY( + SELECT id FROM concurrency_leases + WHERE expires_at < NOW() + LIMIT 1000 + ) + "#, &[], ) .await diff --git a/crates/database/tests/concurrency_leases.rs b/crates/database/tests/concurrency_leases.rs index 56339a5da..2d6f03448 100644 --- a/crates/database/tests/concurrency_leases.rs +++ b/crates/database/tests/concurrency_leases.rs @@ -1,3 +1,5 @@ +// Shared with other test binaries, which use a different subset of it. +#[allow(dead_code)] mod support; use database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository; @@ -60,7 +62,7 @@ async fn concurrent_acquires_never_exceed_the_limit() { let mut at_limit = 0usize; for attempt in attempts { match attempt.await.expect("task").expect("acquire") { - LeaseOutcome::Admitted => admitted += 1, + LeaseOutcome::Admitted { .. } => admitted += 1, LeaseOutcome::AtLimit { .. } => at_limit += 1, } } @@ -120,7 +122,7 @@ async fn released_leases_free_capacity_again() { ) .await .expect("acquire"); - assert_eq!(outcome, LeaseOutcome::Admitted); + assert!(matches!(outcome, LeaseOutcome::Admitted { .. })); held.push(lease_id); } @@ -150,9 +152,8 @@ async fn released_leases_free_capacity_again() { ) .await .expect("acquire"); - assert_eq!( - readmitted, - LeaseOutcome::Admitted, + assert!( + matches!(readmitted, LeaseOutcome::Admitted { .. }), "releasing a lease must free exactly one slot" ); } @@ -192,7 +193,7 @@ async fn renewal_keeps_a_lease_past_its_original_ttl() { ) .await .expect("acquire"); - assert_eq!(outcome, LeaseOutcome::Admitted); + assert!(matches!(outcome, LeaseOutcome::Admitted { .. })); repository.renew(&[lease_id], TTL).await.expect("renew"); @@ -243,7 +244,11 @@ async fn renewal_cannot_recreate_a_released_lease() { .expect("acquire"); repository.release(&[lease_id]).await.expect("release"); - repository.renew(&[lease_id], TTL).await.expect("renew"); + let renewed = repository.renew(&[lease_id], TTL).await.expect("renew"); + assert!( + renewed.is_empty(), + "a released lease must not report as renewed" + ); let live: i64 = pool .get() @@ -301,6 +306,115 @@ async fn pending_leases_are_written_back() { assert_eq!(live, 1, "a degraded-path lease must reach the store"); } +/// A zero or negative rate_limit means unset, not a limit of zero. Reading it +/// literally would reject every request for the organization. +#[tokio::test] +async fn a_zero_rate_limit_falls_back_to_the_default() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + + for stored in [0i32, -1i32] { + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = $1 WHERE id = $2", + &[&stored, &org.org_id], + ) + .await + .expect("rate limit update"); + + let outcome = repository + .try_acquire( + Uuid::new_v4(), + org.org_id, + model.id, + "instance-a", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + + assert!( + matches!(outcome, LeaseOutcome::Admitted { limit } if limit == DEFAULT_LIMIT), + "rate_limit {stored} must fall back to the default, got {outcome:?}" + ); + } +} + +/// cargo test --test concurrency_leases admission_throughput -- --ignored --nocapture +#[tokio::test] +#[ignore] +async fn admission_throughput() { + let concurrency: usize = std::env::var("BENCH_CONCURRENCY") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(16); + let rounds: usize = std::env::var("BENCH_ROUNDS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(200); + + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("bench-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = 100000 WHERE id = $1", + &[&org.org_id], + ) + .await + .expect("rate limit update"); + + let repository = Arc::new(PostgresConcurrencyLeaseRepository::new(pool.clone())); + let started = std::time::Instant::now(); + + for _ in 0..rounds { + let mut batch = Vec::with_capacity(concurrency); + for _ in 0..concurrency { + let repository = repository.clone(); + let org_id = org.org_id; + let model_id = model.id; + batch.push(tokio::spawn(async move { + let lease_id = Uuid::new_v4(); + repository + .try_acquire(lease_id, org_id, model_id, "bench", DEFAULT_LIMIT, TTL) + .await + .expect("acquire"); + lease_id + })); + } + let mut ids: Vec = Vec::with_capacity(concurrency); + for handle in batch { + ids.push(handle.await.expect("task")); + } + repository.release(&ids).await.expect("release"); + } + + let elapsed = started.elapsed(); + let total = rounds * concurrency; + println!( + "concurrency={concurrency} admissions: {total} in {elapsed:?} => {:.0}/s, {:.2}ms mean", + total as f64 / elapsed.as_secs_f64(), + elapsed.as_secs_f64() * 1000.0 / total as f64 + ); +} + /// An expired lease must not keep counting, or a replica that died holding /// leases would lock the organization out permanently. #[tokio::test] @@ -335,7 +449,7 @@ async fn expired_leases_stop_counting() { ) .await .expect("acquire"); - assert_eq!(outcome, LeaseOutcome::Admitted); + assert!(matches!(outcome, LeaseOutcome::Admitted { .. })); tokio::time::sleep(Duration::from_millis(1200)).await; @@ -350,9 +464,8 @@ async fn expired_leases_stop_counting() { ) .await .expect("acquire"); - assert_eq!( - after_expiry, - LeaseOutcome::Admitted, + assert!( + matches!(after_expiry, LeaseOutcome::Admitted { .. }), "a lease past its TTL must not hold capacity" ); } diff --git a/crates/services/src/completions/fleet_concurrency_tests.rs b/crates/services/src/completions/fleet_concurrency_tests.rs new file mode 100644 index 000000000..cd5efa190 --- /dev/null +++ b/crates/services/src/completions/fleet_concurrency_tests.rs @@ -0,0 +1,390 @@ +use super::*; +use crate::metrics::capturing::{CapturingMetricsService, MetricValue}; +use crate::test_utils::{MockAttestationService, MockUsageService}; +use std::collections::HashMap as StdHashMap; +use std::sync::Mutex as StdMutex; +use std::time::Instant; + +const LIMIT: u32 = 3; +const MODEL_NAME: &str = "test/model"; + +/// Stands in for the lease table shared by every replica. Admission is +/// serialised the way the advisory lock serialises it in Postgres. +#[derive(Default)] +struct SharedLeaseStore { + rows: StdMutex>, + limit: StdMutex>, + unavailable: std::sync::atomic::AtomicBool, +} + +impl SharedLeaseStore { + fn with_limit(limit: u32) -> Self { + Self { + limit: StdMutex::new(Some(limit)), + ..Default::default() + } + } + + fn set_unavailable(&self, unavailable: bool) { + self.unavailable + .store(unavailable, std::sync::atomic::Ordering::SeqCst); + } + + fn is_unavailable(&self) -> bool { + self.unavailable.load(std::sync::atomic::Ordering::SeqCst) + } + + fn live_count(&self, organization_id: Uuid, model_id: Uuid) -> usize { + let rows = self.rows.lock().unwrap(); + rows.values() + .filter(|(org, model, expires)| { + *org == organization_id && *model == model_id && *expires > Instant::now() + }) + .count() + } +} + +#[async_trait::async_trait] +impl ports::ConcurrencyLeaseRepository for SharedLeaseStore { + async fn try_acquire( + &self, + lease_id: Uuid, + organization_id: Uuid, + model_id: Uuid, + _instance_id: &str, + default_limit: u32, + ttl: Duration, + ) -> Result { + if self.is_unavailable() { + anyhow::bail!("lease store unavailable"); + } + + let limit = self.limit.lock().unwrap().unwrap_or(default_limit); + let mut rows = self.rows.lock().unwrap(); + let in_flight = rows + .values() + .filter(|(org, model, expires)| { + *org == organization_id && *model == model_id && *expires > Instant::now() + }) + .count(); + + if in_flight >= limit as usize { + return Ok(ports::LeaseOutcome::AtLimit { + limit, + in_flight: in_flight as i64, + }); + } + + rows.insert(lease_id, (organization_id, model_id, Instant::now() + ttl)); + Ok(ports::LeaseOutcome::Admitted { limit }) + } + + async fn release(&self, lease_ids: &[Uuid]) -> Result<(), anyhow::Error> { + let mut rows = self.rows.lock().unwrap(); + for id in lease_ids { + rows.remove(id); + } + Ok(()) + } + + async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result, anyhow::Error> { + if self.is_unavailable() { + anyhow::bail!("lease store unavailable"); + } + let mut rows = self.rows.lock().unwrap(); + let mut renewed = Vec::new(); + for id in lease_ids { + if let Some(row) = rows.get_mut(id) { + row.2 = Instant::now() + ttl; + renewed.push(*id); + } + } + Ok(renewed) + } + + async fn persist( + &self, + leases: &[ports::HeldLease], + _instance_id: &str, + ttl: Duration, + ) -> Result<(), anyhow::Error> { + if self.is_unavailable() { + anyhow::bail!("lease store unavailable"); + } + let mut rows = self.rows.lock().unwrap(); + for lease in leases { + rows.insert( + lease.id, + (lease.organization_id, lease.model_id, Instant::now() + ttl), + ); + } + Ok(()) + } + + async fn sweep_expired(&self) -> Result { + let mut rows = self.rows.lock().unwrap(); + let before = rows.len(); + rows.retain(|_, (_, _, expires)| *expires > Instant::now()); + Ok((before - rows.len()) as u64) + } +} + +/// Admission never consults the model catalogue, so this only has to exist. +struct EmptyModelsRepository; + +#[async_trait::async_trait] +impl crate::models::ModelsRepository for EmptyModelsRepository { + async fn get_all_active_models( + &self, + ) -> Result, anyhow::Error> { + Ok(Vec::new()) + } + + async fn get_model_by_name( + &self, + _name: &str, + ) -> Result, anyhow::Error> { + Ok(None) + } + + async fn resolve_and_get_model( + &self, + _identifier: &str, + ) -> Result, anyhow::Error> { + Ok(None) + } + + async fn get_configured_model_names(&self) -> Result, anyhow::Error> { + Ok(Vec::new()) + } +} + +struct StaticLimitRepository(Option); + +#[async_trait::async_trait] +impl ports::OrganizationConcurrentLimitRepository for StaticLimitRepository { + async fn get_concurrent_limit(&self, _org_id: Uuid) -> Result, anyhow::Error> { + Ok(self.0) + } +} + +struct FailingLimitRepository; + +#[async_trait::async_trait] +impl ports::OrganizationConcurrentLimitRepository for FailingLimitRepository { + async fn get_concurrent_limit(&self, _org_id: Uuid) -> Result, anyhow::Error> { + anyhow::bail!("limit lookup unavailable") + } +} + +fn replica( + store: Arc, + metrics: Arc, + limits: Arc, +) -> CompletionServiceImpl { + let pool = Arc::new(InferenceProviderPool::new( + None, + config::ExternalProvidersConfig::default(), + )); + CompletionServiceImpl::new( + pool, + Arc::new(MockAttestationService), + Arc::new(MockUsageService), + metrics, + Arc::new(EmptyModelsRepository), + limits, + ) + .with_fleet_concurrency(store, "test-instance".to_string(), Duration::from_secs(60)) +} + +fn counts(metrics: &CapturingMetricsService, name: &str) -> usize { + metrics + .get_metrics() + .into_iter() + .filter(|metric| metric.name == name && matches!(metric.value, MetricValue::Count(_))) + .count() +} + +/// The invariant the issue turns on: several replicas sharing one lease store +/// admit no more than the organization's limit between them. +#[tokio::test] +async fn replicas_share_one_limit() { + const REPLICAS: usize = 4; + const ATTEMPTS: usize = 20; + + let store = Arc::new(SharedLeaseStore::with_limit(LIMIT)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + + let replicas: Vec> = (0..REPLICAS) + .map(|_| { + Arc::new(replica( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(LIMIT))), + )) + }) + .collect(); + + let mut slots = Vec::new(); + for attempt in 0..ATTEMPTS { + let service = &replicas[attempt % REPLICAS]; + if let Ok(slot) = service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + { + slots.push(slot); + } + } + + assert_eq!( + slots.len(), + LIMIT as usize, + "{REPLICAS} replicas admitted {} against a shared limit of {LIMIT}", + slots.len() + ); + assert_eq!(store.live_count(organization_id, model_id), LIMIT as usize); + assert_eq!( + counts(&metrics, METRIC_CONCURRENCY_ADMITTED), + LIMIT as usize + ); + assert_eq!( + counts(&metrics, METRIC_CONCURRENCY_REJECTED), + ATTEMPTS - LIMIT as usize + ); +} + +/// Releasing on one replica must free capacity for a different replica. +#[tokio::test] +async fn capacity_freed_on_one_replica_is_visible_to_another() { + let store = Arc::new(SharedLeaseStore::with_limit(1)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + + let first = replica( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(1))), + ); + let second = replica( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(1))), + ); + + // The guard is what releases; a bare slot has no Drop of its own. + let guard = ConcurrentSlotGuard::new( + first + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .expect("first replica admitted"), + ); + assert!(second + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .is_err()); + + drop(guard); + tokio::time::sleep(Duration::from_millis(50)).await; + + assert!( + second + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .is_ok(), + "a slot released on one replica must become usable on another" + ); +} + +/// With the store unreachable a replica counts only its own leases, which is +/// the behaviour that shipped before fleet limits. It must never admit a fresh +/// limit on top of leases already live. +#[tokio::test] +async fn degraded_admission_does_not_stack_on_live_leases() { + let store = Arc::new(SharedLeaseStore::with_limit(LIMIT)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + + let service = replica( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(LIMIT))), + ); + + let mut slots = Vec::new(); + for _ in 0..LIMIT { + slots.push( + service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .expect("admitted while healthy"), + ); + } + + store.set_unavailable(true); + + assert!( + service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .is_err(), + "the degraded path must count leases this replica already holds" + ); + assert_eq!(counts(&metrics, METRIC_CONCURRENCY_DEGRADED), 1); +} + +/// A failed limit lookup must not widen the cap. The organization's real limit +/// was read once, so it is used rather than the global default. +#[tokio::test] +async fn a_failed_limit_lookup_keeps_the_known_limit() { + // The organization's real limit is 1, well under the global default of 64. + let store = Arc::new(SharedLeaseStore::with_limit(1)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + + let service = replica( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(1))), + ); + + let _slot = service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .expect("admitted while healthy"); + + store.set_unavailable(true); + + assert!( + service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .is_err(), + "an outage must not raise the cap to the global default" + ); +} + +#[tokio::test] +async fn an_unknown_limit_falls_back_to_the_default() { + let store = Arc::new(SharedLeaseStore::default()); + let metrics = Arc::new(CapturingMetricsService::new()); + let service = replica( + store.clone(), + metrics.clone(), + Arc::new(FailingLimitRepository), + ); + + store.set_unavailable(true); + let slot = service + .try_acquire_concurrent_slot(Uuid::new_v4(), Uuid::new_v4(), MODEL_NAME) + .await; + + assert!( + slot.is_ok(), + "with no limit ever read the default applies rather than rejecting" + ); +} diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index c797a7413..e2e051c66 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -518,7 +518,8 @@ where } } -#[derive(Clone)] +/// Deliberately not `Clone`: releasing two copies would decrement the local +/// counter twice, and an underflow there is what locks an organization out. pub(crate) enum ConcurrentSlot { Local(Arc), Lease { @@ -534,7 +535,11 @@ impl ConcurrentSlot { fn release(&self) { match self { Self::Local(counter) => { - counter.fetch_sub(1, Ordering::Release); + // Saturating: an underflow here would read as a permanently + // full counter and lock the organization out of the model. + let _ = counter.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + Some(current.saturating_sub(1)) + }); } Self::Lease { id, @@ -607,17 +612,15 @@ enum LeaseAdmission { Unavailable(anyhow::Error), } -/// Leases this replica currently holds. Renewal reads it to keep them alive, -/// release removes from it, and admission counts it when the lease store is -/// unreachable, so the degraded path can never admit on top of live leases. +/// Leases this replica holds. Read by renewal, release, and degraded +/// admission, which is what keeps the degraded path off live leases. #[derive(Default)] pub(crate) struct HeldLeases { by_model: std::sync::Mutex>>, } -/// Whether a held lease already has a row in the lease store. Renewal updates -/// stored leases and inserts pending ones, so an update can never recreate a -/// row that a release deleted while the round trip was outstanding. +/// Whether the store already has the row. Renewal updates stored leases and +/// inserts pending ones, so it can never recreate a released row. #[derive(Clone, Copy, PartialEq, Eq)] enum LeaseState { Stored, @@ -656,6 +659,7 @@ impl HeldLeases { } } + #[cfg(test)] fn count(&self, organization_id: Uuid, model_id: Uuid) -> usize { let guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); guard @@ -663,9 +667,8 @@ impl HeldLeases { .map_or(0, |ids| ids.len()) } - /// Record a lease only if the model is below `limit`, deciding and - /// inserting under one lock so concurrent degraded admissions cannot both - /// observe the same count and both proceed. + /// Decides and inserts under one lock, so two degraded admissions cannot + /// both observe the same count and both proceed. fn insert_below( &self, organization_id: Uuid, @@ -744,6 +747,12 @@ const ORG_LIMIT_CACHE_TTL_SECS: u64 = 300; /// How long a limit that was actually read stays usable as a fallback (24h). const LAST_KNOWN_LIMIT_TTL_SECS: u64 = 86_400; +/// The background loops have no model in scope, but leaving these untagged +/// would make prod and staging indistinguishable. +fn background_tags() -> Vec { + vec![format!("{}:{}", TAG_ENVIRONMENT, get_environment())] +} + /// TTL for concurrent count cache entries (10 minutes). /// Safety net: if a counter gets stuck (e.g., due to a panic or proxy not propagating /// client disconnection), the entry expires and is replaced with a fresh zero counter. @@ -895,7 +904,10 @@ impl CompletionServiceImpl { let renewing = held.clone(); let renew_instance = instance_id.clone(); let renew_every = ttl / LEASE_RENEWALS_PER_TTL; + let renew_metrics = self.metrics_service.clone(); let renew_task = tokio::spawn(async move { + let owned_tags = background_tags(); + let tags: Vec<&str> = owned_tags.iter().map(|tag| tag.as_str()).collect(); let mut ticker = tokio::time::interval(renew_every); ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); loop { @@ -903,13 +915,43 @@ impl CompletionServiceImpl { let (stored, pending) = renewing.snapshot(); let stored_ids: Vec = stored.iter().map(|lease| lease.id).collect(); - if let Err(error) = renewer.renew(&stored_ids, ttl).await { - tracing::warn!( - leases = stored_ids.len(), - error = %error, - "Failed to renew concurrency leases; in-flight requests continue" - ); + let mut pending = pending; + match renewer.renew(&stored_ids, ttl).await { + Ok(renewed) => { + // A missing row was swept after renewal stalled past the TTL. + // Re-storing only what is still held avoids reviving a release. + let renewed: HashSet = renewed.into_iter().collect(); + let held = renewing.held_ids(); + let lost: Vec = stored + .into_iter() + .filter(|lease| { + !renewed.contains(&lease.id) && held.contains(&lease.id) + }) + .collect(); + if !lost.is_empty() { + tracing::warn!( + leases = lost.len(), + "Re-storing held leases that expired before renewal" + ); + pending.extend(lost); + } + } + Err(error) => { + renew_metrics.record_count(METRIC_CONCURRENCY_RENEW_FAILED, 1, &tags); + tracing::warn!( + leases = stored_ids.len(), + error = %error, + "Failed to renew concurrency leases; in-flight requests continue" + ) + } } + // A histogram, not a counter: this is the current depth, and a + // counter would sum successive readings into a meaningless total. + renew_metrics.record_histogram( + METRIC_CONCURRENCY_IN_FLIGHT, + renewing.held_ids().len() as f64, + &tags, + ); if pending.is_empty() { continue; @@ -947,17 +989,27 @@ impl CompletionServiceImpl { supervise("renew", renew_task); let sweeper = repository.clone(); + let sweep_metrics = self.metrics_service.clone(); let sweep_task = tokio::spawn(async move { + let owned_tags = background_tags(); + let tags: Vec<&str> = owned_tags.iter().map(|tag| tag.as_str()).collect(); let mut ticker = tokio::time::interval(ttl); ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); loop { ticker.tick().await; match sweeper.sweep_expired().await { Ok(0) => {} - Ok(removed) => tracing::info!( - removed = removed, - "Reclaimed concurrency leases whose holder stopped renewing" - ), + Ok(removed) => { + sweep_metrics.record_count( + METRIC_CONCURRENCY_RECLAIMED, + removed as i64, + &tags, + ); + tracing::info!( + removed = removed, + "Reclaimed concurrency leases whose holder stopped renewing" + ) + } Err(error) => tracing::warn!( error = %error, "Failed to sweep expired concurrency leases" @@ -1167,8 +1219,7 @@ impl CompletionServiceImpl { } } - /// Last limit actually read for this organization. Outliving the lookup - /// cache matters because an outage longer than that TTL would otherwise + /// Outlives the lookup cache: an outage longer than that TTL would otherwise /// leave the degraded path widening the cap to the global default. async fn last_known_org_concurrent_limit(&self, organization_id: Uuid) -> Option { self.last_known_limits.get(&organization_id).await @@ -1521,6 +1572,13 @@ impl CompletionServiceImpl { .record_count(METRIC_REQUEST_ERRORS, 1, &tags_str); } + fn record_concurrency(&self, metric: &str, model_name: &str, scope: &str) { + let mut tags = Self::create_metric_tags(model_name); + tags.push(format!("{}:{}", TAG_SCOPE, scope)); + let tags: Vec<&str> = tags.iter().map(|tag| tag.as_str()).collect(); + self.metrics_service.record_count(metric, 1, &tags); + } + /// Convert completion messages to chat messages for inference providers fn prepare_chat_messages(messages: &[ports::CompletionMessage]) -> Vec { messages @@ -1559,10 +1617,8 @@ impl CompletionServiceImpl { .collect() } - /// Admission while the lease store is unreachable. Counting this replica's - /// own leases degrades to the per-process behaviour that shipped before - /// fleet limits existed, and never admits on top of leases already live. - /// The lease is recorded locally only; the next renewal writes it back. + /// Counts this replica's own leases, so a store outage degrades to the + /// per-process behaviour instead of admitting on top of live leases. async fn admit_from_held_leases( &self, fleet: &FleetConcurrency, @@ -1580,14 +1636,18 @@ impl CompletionServiceImpl { .held .insert_below(organization_id, model_id, lease_id, limit as usize) { - Ok(()) => Ok(ConcurrentSlot::Lease { - id: lease_id, - organization_id, - model_id, - held: fleet.held.clone(), - release: fleet.release.clone(), - }), + Ok(()) => { + self.record_concurrency(METRIC_CONCURRENCY_ADMITTED, model_name, SCOPE_REPLICA); + Ok(ConcurrentSlot::Lease { + id: lease_id, + organization_id, + model_id, + held: fleet.held.clone(), + release: fleet.release.clone(), + }) + } Err(in_flight) => { + self.record_concurrency(METRIC_CONCURRENCY_REJECTED, model_name, SCOPE_REPLICA); tracing::warn!( organization_id = %organization_id, model_id = %model_id, @@ -1616,6 +1676,7 @@ impl CompletionServiceImpl { model_name: &str, ) -> Result { let lease_id = Uuid::new_v4(); + let started = Instant::now(); let outcome = fleet .repository .try_acquire( @@ -1628,9 +1689,21 @@ impl CompletionServiceImpl { ) .await .map_err(LeaseAdmission::Unavailable)?; + self.metrics_service.record_latency( + METRIC_CONCURRENCY_ADMISSION_LATENCY, + started.elapsed(), + &Self::create_metric_tags(model_name) + .iter() + .map(|tag| tag.as_str()) + .collect::>(), + ); match outcome { - ports::LeaseOutcome::Admitted => { + ports::LeaseOutcome::Admitted { limit } => { + // The fleet path reads the limit inside the admission statement, + // so this is the only place the degraded path can learn it. + self.last_known_limits.insert(organization_id, limit).await; + self.record_concurrency(METRIC_CONCURRENCY_ADMITTED, model_name, SCOPE_FLEET); fleet .held .insert(organization_id, model_id, lease_id, LeaseState::Stored); @@ -1643,6 +1716,8 @@ impl CompletionServiceImpl { }) } ports::LeaseOutcome::AtLimit { limit, in_flight } => { + self.last_known_limits.insert(organization_id, limit).await; + self.record_concurrency(METRIC_CONCURRENCY_REJECTED, model_name, SCOPE_FLEET); tracing::warn!( organization_id = %organization_id, model_id = %model_id, @@ -1679,6 +1754,7 @@ impl CompletionServiceImpl { Ok(slot) => return Ok(slot), Err(LeaseAdmission::AtLimit(error)) => return Err(error), Err(LeaseAdmission::Unavailable(error)) => { + self.record_concurrency(METRIC_CONCURRENCY_DEGRADED, model_name, SCOPE_REPLICA); tracing::warn!( organization_id = %organization_id, model_id = %model_id, @@ -2544,9 +2620,127 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { pub use ports::*; +#[cfg(test)] +mod fleet_concurrency_tests; + #[cfg(test)] mod provider_attribution_tests; +#[cfg(test)] +mod held_lease_tests { + use super::*; + + fn ids() -> (Uuid, Uuid) { + (Uuid::new_v4(), Uuid::new_v4()) + } + + #[test] + fn insert_below_refuses_at_the_limit() { + let (org, model) = ids(); + let held = HeldLeases::default(); + + assert!(held.insert_below(org, model, Uuid::new_v4(), 2).is_ok()); + assert!(held.insert_below(org, model, Uuid::new_v4(), 2).is_ok()); + assert_eq!(held.insert_below(org, model, Uuid::new_v4(), 2), Err(2)); + assert_eq!(held.count(org, model), 2); + } + + #[test] + fn models_are_counted_separately() { + let (org, model) = ids(); + let other_model = Uuid::new_v4(); + let held = HeldLeases::default(); + + assert!(held.insert_below(org, model, Uuid::new_v4(), 1).is_ok()); + assert!(held + .insert_below(org, other_model, Uuid::new_v4(), 1) + .is_ok()); + assert_eq!(held.count(org, model), 1); + assert_eq!(held.count(org, other_model), 1); + } + + #[test] + fn releasing_frees_exactly_one_slot() { + let (org, model) = ids(); + let held = HeldLeases::default(); + let first = Uuid::new_v4(); + + assert!(held.insert_below(org, model, first, 1).is_ok()); + assert!(held.insert_below(org, model, Uuid::new_v4(), 1).is_err()); + + held.remove(org, model, first); + assert_eq!(held.count(org, model), 0); + assert!(held.insert_below(org, model, Uuid::new_v4(), 1).is_ok()); + } + + #[test] + fn removing_an_unknown_lease_does_not_free_a_slot() { + let (org, model) = ids(); + let held = HeldLeases::default(); + + assert!(held.insert_below(org, model, Uuid::new_v4(), 1).is_ok()); + held.remove(org, model, Uuid::new_v4()); + assert_eq!(held.count(org, model), 1); + } + + #[test] + fn new_leases_are_pending_until_marked_stored() { + let (org, model) = ids(); + let held = HeldLeases::default(); + let lease_id = Uuid::new_v4(); + + held.insert(org, model, lease_id, LeaseState::PendingWrite); + let (stored, pending) = held.snapshot(); + assert!(stored.is_empty()); + assert_eq!(pending.len(), 1); + + held.mark_stored(&pending); + let (stored, pending) = held.snapshot(); + assert_eq!(stored.len(), 1); + assert!(pending.is_empty()); + } + + #[test] + fn a_lease_released_before_being_marked_is_not_resurrected() { + let (org, model) = ids(); + let held = HeldLeases::default(); + let lease_id = Uuid::new_v4(); + + held.insert(org, model, lease_id, LeaseState::PendingWrite); + let (_, pending) = held.snapshot(); + + held.remove(org, model, lease_id); + held.mark_stored(&pending); + + assert_eq!(held.count(org, model), 0); + assert!(held.held_ids().is_empty()); + } + + #[test] + fn concurrent_admissions_cannot_exceed_the_limit() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let (org, model) = ids(); + let held = Arc::new(HeldLeases::default()); + let admitted = Arc::new(AtomicUsize::new(0)); + + std::thread::scope(|scope| { + for _ in 0..32 { + let held = held.clone(); + let admitted = admitted.clone(); + scope.spawn(move || { + if held.insert_below(org, model, Uuid::new_v4(), 4).is_ok() { + admitted.fetch_add(1, Ordering::Relaxed); + } + }); + } + }); + + assert_eq!(admitted.load(Ordering::Relaxed), 4); + assert_eq!(held.count(org, model), 4); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/services/src/completions/ports.rs b/crates/services/src/completions/ports.rs index f94c29785..5c8afdadd 100644 --- a/crates/services/src/completions/ports.rs +++ b/crates/services/src/completions/ports.rs @@ -130,9 +130,11 @@ pub trait OrganizationConcurrentLimitRepository: Send + Sync { async fn get_concurrent_limit(&self, org_id: Uuid) -> Result, anyhow::Error>; } +/// Both variants carry the limit that was applied, so the caller can remember +/// it for the degraded path without a second lookup. #[derive(Debug, Clone, PartialEq, Eq)] pub enum LeaseOutcome { - Admitted, + Admitted { limit: u32 }, AtLimit { limit: u32, in_flight: i64 }, } @@ -160,13 +162,12 @@ pub trait ConcurrencyLeaseRepository: Send + Sync { async fn release(&self, lease_ids: &[Uuid]) -> Result<(), anyhow::Error>; - /// Extend leases the store already has. Updating in place means a lease - /// released while this call was in flight stays deleted instead of being - /// recreated with no holder. - async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result<(), anyhow::Error>; + /// Extend leases the store already has, returning the ids still there. + /// Updates in place so a lease released mid-call stays deleted. + async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result, anyhow::Error>; - /// Write leases the store does not have yet, which happens when a request - /// was admitted while it was unreachable. Returns the leases now stored. + /// Write leases the store does not have yet, from requests admitted while + /// it was unreachable. async fn persist( &self, leases: &[HeldLease], diff --git a/crates/services/src/metrics/consts.rs b/crates/services/src/metrics/consts.rs index e6ef0d703..d6d8167f9 100644 --- a/crates/services/src/metrics/consts.rs +++ b/crates/services/src/metrics/consts.rs @@ -22,6 +22,24 @@ pub const METRIC_SIGNATURE_CREATION_DURATION: &str = "cloud_api.signature.creati pub const METRIC_ATTESTATION_REPORT_CACHE: &str = "cloud_api.attestation.report_cache"; pub const TAG_RESULT: &str = "result"; +// No organization tag: that breakdown comes from the concurrency_leases table, +// so it costs no metric cardinality. +pub const METRIC_CONCURRENCY_ADMITTED: &str = "cloud_api.concurrency.admitted"; +pub const METRIC_CONCURRENCY_REJECTED: &str = "cloud_api.concurrency.rejected"; +pub const METRIC_CONCURRENCY_DEGRADED: &str = "cloud_api.concurrency.degraded"; + +// What the lease store adds to the inference hot path, and the health of the +// background loops that keep leases accurate. +pub const METRIC_CONCURRENCY_ADMISSION_LATENCY: &str = "cloud_api.concurrency.admission_latency"; +pub const METRIC_CONCURRENCY_IN_FLIGHT: &str = "cloud_api.concurrency.in_flight"; +pub const METRIC_CONCURRENCY_RENEW_FAILED: &str = "cloud_api.concurrency.renew_failed"; +pub const METRIC_CONCURRENCY_RECLAIMED: &str = "cloud_api.concurrency.reclaimed"; + +// Whether admission was decided fleet-wide or by this replica alone. +pub const TAG_SCOPE: &str = "scope"; +pub const SCOPE_FLEET: &str = "fleet"; +pub const SCOPE_REPLICA: &str = "replica"; + // Usage/engagement metrics pub const METRIC_REQUEST_COUNT: &str = "cloud_api.request.count"; pub const METRIC_TOKENS_INPUT: &str = "cloud_api.tokens.input"; From f783811c20a269bcf582cfdfd4eb4ca09ef9debc Mon Sep 17 00:00:00 2001 From: neo-sky Date: Mon, 24 Aug 2026 20:54:59 -0400 Subject: [PATCH 4/8] Add shadow mode and cover the transport guard Shadow counts leases fleet-wide but leaves rejection to the per-replica limit, so the numbers can be measured before enforcement turns anyone away. Also covers the guard direct transport routes hold, which had no test. --- crates/api/src/lib.rs | 3 +- crates/config/src/types.rs | 3 + .../completions/fleet_concurrency_tests.rs | 110 +++++++++++++++++- crates/services/src/completions/mod.rs | 14 +++ crates/services/src/metrics/consts.rs | 1 + 5 files changed, 129 insertions(+), 2 deletions(-) diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 7d98b65a5..f8338cfb4 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -433,7 +433,7 @@ pub async fn init_domain_services_with_pool( org_limit_repository, ); - if config.fleet_concurrency.mode == config::FleetConcurrencyMode::Enforce { + if config.fleet_concurrency.mode != config::FleetConcurrencyMode::Off { let lease_repository = Arc::new( database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository::new( database.pool().clone(), @@ -443,6 +443,7 @@ pub async fn init_domain_services_with_pool( lease_repository, config.fleet_concurrency.instance_id.clone(), std::time::Duration::from_secs(config.fleet_concurrency.lease_ttl_seconds), + config.fleet_concurrency.mode == config::FleetConcurrencyMode::Enforce, ); } diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index 929f5a309..5126e32f2 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -37,6 +37,8 @@ pub enum FleetConcurrencyMode { /// Each replica counts only its own in-flight requests. #[default] Off, + /// Leases are counted fleet-wide but rejection stays per-replica. + Shadow, /// Admission is decided by the fleet-wide lease count. Enforce, } @@ -63,6 +65,7 @@ impl FleetConcurrencyConfig { let requested = env::var("FLEET_CONCURRENCY_MODE").unwrap_or_default(); let mode = match requested.to_ascii_lowercase().as_str() { "enforce" => FleetConcurrencyMode::Enforce, + "shadow" => FleetConcurrencyMode::Shadow, "" | "off" => FleetConcurrencyMode::Off, _ => { // eprintln for the reason given on the CHUTES_MODELS warning diff --git a/crates/services/src/completions/fleet_concurrency_tests.rs b/crates/services/src/completions/fleet_concurrency_tests.rs index cd5efa190..f2d32d9b2 100644 --- a/crates/services/src/completions/fleet_concurrency_tests.rs +++ b/crates/services/src/completions/fleet_concurrency_tests.rs @@ -181,6 +181,15 @@ fn replica( store: Arc, metrics: Arc, limits: Arc, +) -> CompletionServiceImpl { + replica_with(store, metrics, limits, true) +} + +fn replica_with( + store: Arc, + metrics: Arc, + limits: Arc, + enforcing: bool, ) -> CompletionServiceImpl { let pool = Arc::new(InferenceProviderPool::new( None, @@ -194,7 +203,12 @@ fn replica( Arc::new(EmptyModelsRepository), limits, ) - .with_fleet_concurrency(store, "test-instance".to_string(), Duration::from_secs(60)) + .with_fleet_concurrency( + store, + "test-instance".to_string(), + Duration::from_secs(60), + enforcing, + ) } fn counts(metrics: &CapturingMetricsService, name: &str) -> usize { @@ -388,3 +402,97 @@ async fn an_unknown_limit_falls_back_to_the_default() { "with no limit ever read the default applies rather than rejecting" ); } + +/// Shadow mode is only useful if it measures without rejecting: the fleet +/// count goes over the limit and the request still proceeds. +#[tokio::test] +async fn shadowing_records_the_verdict_without_rejecting() { + const REPLICAS: usize = 4; + const ATTEMPTS: usize = 20; + + let store = Arc::new(SharedLeaseStore::with_limit(LIMIT)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + + let replicas: Vec> = (0..REPLICAS) + .map(|_| { + Arc::new(replica_with( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(LIMIT))), + false, + )) + }) + .collect(); + + let mut slots = Vec::new(); + for attempt in 0..ATTEMPTS { + let service = &replicas[attempt % REPLICAS]; + if let Ok(slot) = service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + { + slots.push(slot); + } + } + + assert!( + slots.len() > LIMIT as usize, + "shadowing must not enforce the fleet limit, admitted {}", + slots.len() + ); + assert!( + counts(&metrics, METRIC_CONCURRENCY_WOULD_REJECT) > 0, + "the over-limit verdict is the whole signal shadow mode exists to give" + ); + assert_eq!( + counts(&metrics, METRIC_CONCURRENCY_REJECTED), + 0, + "nothing is rejected while shadowing" + ); +} + +/// The guard handed to direct transport routes must free fleet capacity when +/// dropped, or those routes hold leases until the TTL expires. +#[tokio::test] +async fn dropping_a_transport_guard_frees_fleet_capacity() { + use ports::CompletionServiceTrait; + + let store = Arc::new(SharedLeaseStore::with_limit(1)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + let service = replica( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(1))), + ); + + let guard = service + .acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .expect("the first request takes the only slot"); + assert_eq!(store.live_count(organization_id, model_id), 1); + + assert!( + service + .acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .is_err(), + "the limit is one, so a second transport request is refused" + ); + + drop(guard); + tokio::time::sleep(Duration::from_millis(50)).await; + + assert_eq!( + store.live_count(organization_id, model_id), + 0, + "the released lease must leave the store, not linger until it expires" + ); + service + .acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + .expect("the freed slot admits the next transport request"); +} diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 5ecb5527a..2c5a75373 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -667,6 +667,8 @@ pub struct CompletionServiceImpl { enum LeaseAdmission { AtLimit(ports::CompletionError), + /// Over the fleet limit while shadowing, so the per-process limit decides. + Shadowed, /// The lease store could not answer, so admission falls back to the /// per-process limit rather than rejecting the request. Unavailable(anyhow::Error), @@ -778,6 +780,7 @@ struct FleetConcurrency { release: mpsc::UnboundedSender, instance_id: String, ttl: Duration, + enforcing: bool, } /// Renewals attempted within one lease TTL. Three gives two spare attempts @@ -938,6 +941,7 @@ impl CompletionServiceImpl { repository: Arc, instance_id: String, ttl: Duration, + enforcing: bool, ) -> Self { let (release, mut released) = mpsc::unbounded_channel::(); let releaser = repository.clone(); @@ -1086,6 +1090,7 @@ impl CompletionServiceImpl { release, instance_id, ttl, + enforcing, }); self } @@ -1801,6 +1806,14 @@ impl CompletionServiceImpl { } ports::LeaseOutcome::AtLimit { limit, in_flight } => { self.last_known_limits.insert(organization_id, limit).await; + if !fleet.enforcing { + self.record_concurrency( + METRIC_CONCURRENCY_WOULD_REJECT, + model_name, + SCOPE_FLEET, + ); + return Err(LeaseAdmission::Shadowed); + } self.record_concurrency(METRIC_CONCURRENCY_REJECTED, model_name, SCOPE_FLEET); tracing::warn!( organization_id = %organization_id, @@ -1871,6 +1884,7 @@ impl CompletionServiceImpl { { Ok(slot) => return Ok(slot), Err(LeaseAdmission::AtLimit(error)) => return Err(error), + Err(LeaseAdmission::Shadowed) => {} Err(LeaseAdmission::Unavailable(error)) => { self.record_concurrency(METRIC_CONCURRENCY_DEGRADED, model_name, SCOPE_REPLICA); tracing::warn!( diff --git a/crates/services/src/metrics/consts.rs b/crates/services/src/metrics/consts.rs index 1466e3343..c3010e280 100644 --- a/crates/services/src/metrics/consts.rs +++ b/crates/services/src/metrics/consts.rs @@ -34,6 +34,7 @@ pub const TAG_RESULT: &str = "result"; pub const METRIC_CONCURRENCY_ADMITTED: &str = "cloud_api.concurrency.admitted"; pub const METRIC_CONCURRENCY_REJECTED: &str = "cloud_api.concurrency.rejected"; pub const METRIC_CONCURRENCY_DEGRADED: &str = "cloud_api.concurrency.degraded"; +pub const METRIC_CONCURRENCY_WOULD_REJECT: &str = "cloud_api.concurrency.would_reject"; // What the lease store adds to the inference hot path, and the health of the // background loops that keep leases accurate. From 6eca5fd9dfd1cd4e6edc9c390623e95aa8a38e67 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Tue, 25 Aug 2026 20:39:07 -0400 Subject: [PATCH 5/8] Bound the release queue and stop a retry rejecting its own lease A retry after a lost commit response counted the lease it had already written, so it could reject the request that was holding it. Releases now queue with a bound, and a dropped release or a failed sweep is counted rather than only logged. --- crates/config/Cargo.toml | 2 +- crates/config/src/types.rs | 21 ++++++--- .../sql/V0075__add_concurrency_leases.sql | 2 + .../src/repositories/concurrency_lease.rs | 3 +- crates/services/src/completions/mod.rs | 47 +++++++++++++++---- crates/services/src/completions/ports.rs | 7 +-- crates/services/src/metrics/consts.rs | 2 + 7 files changed, 63 insertions(+), 21 deletions(-) diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index 95949e797..08f9470ee 100644 --- a/crates/config/Cargo.toml +++ b/crates/config/Cargo.toml @@ -10,7 +10,7 @@ description = "Configuration management for cloud-api" dotenvy = "0.15.7" serde = { version = "1.0", features = ["derive"] } thiserror = "2.0" -uuid = { version = "1.23", features = ["serde"] } +uuid = { version = "1.23", features = ["serde", "v4"] } [dev-dependencies] tempfile = "3.27" diff --git a/crates/config/src/types.rs b/crates/config/src/types.rs index 5126e32f2..f99bc0388 100644 --- a/crates/config/src/types.rs +++ b/crates/config/src/types.rs @@ -1,5 +1,6 @@ use crate::ita::ItaAttestationConfig; use std::{collections::HashMap, env}; +use uuid::Uuid; #[derive(Debug, Clone)] pub struct ApiConfig { @@ -81,17 +82,25 @@ impl FleetConcurrencyConfig { Self { mode, - lease_ttl_seconds: env::var("FLEET_CONCURRENCY_LEASE_TTL_SECONDS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|seconds| *seconds > 0) - .unwrap_or(60), + lease_ttl_seconds: match env::var("FLEET_CONCURRENCY_LEASE_TTL_SECONDS") { + Ok(value) => match value.parse::() { + Ok(seconds) if seconds > 0 => seconds, + _ => { + eprintln!( + "WARN: invalid FLEET_CONCURRENCY_LEASE_TTL_SECONDS '{value}', \ + falling back to 60" + ); + 60 + } + }, + Err(_) => 60, + }, instance_id: env::var("FLEET_CONCURRENCY_INSTANCE_ID") .ok() .filter(|value| !value.is_empty()) .unwrap_or_else(|| { let host = env::var("HOSTNAME").unwrap_or_else(|_| "unknown".to_string()); - format!("{host}-{}", std::process::id()) + format!("{host}-{}", Uuid::new_v4().simple()) }), } } diff --git a/crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql b/crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql index f558729df..1fa53cd71 100644 --- a/crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql +++ b/crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql @@ -1,3 +1,5 @@ +-- One row per in-flight request, so replicas count against a shared limit +-- rather than one each. Rows outlive a dead replica and are swept on expiry. CREATE TABLE concurrency_leases ( id UUID PRIMARY KEY, organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE, diff --git a/crates/database/src/repositories/concurrency_lease.rs b/crates/database/src/repositories/concurrency_lease.rs index 3dd7642f1..3ef050d18 100644 --- a/crates/database/src/repositories/concurrency_lease.rs +++ b/crates/database/src/repositories/concurrency_lease.rs @@ -71,9 +71,10 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { WHERE organization_id = $1 AND model_id = $2 AND expires_at > NOW() + AND id <> $3 ) AS in_flight "#, - &[&organization_id, &model_id], + &[&organization_id, &model_id, &lease_id], ) .await .map_err(map_db_error)?; diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 2c5a75373..d9f2c799f 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -587,7 +587,7 @@ pub(crate) enum ConcurrentSlot { organization_id: Uuid, model_id: Uuid, held: Arc, - release: mpsc::UnboundedSender, + release: mpsc::Sender, }, } @@ -609,11 +609,12 @@ impl ConcurrentSlot { release, } => { held.remove(*organization_id, *model_id, *id); - if release.send(*id).is_err() { + if release.try_send(*id).is_err() { + held.record_dropped_release(); tracing::error!( lease_id = %id, - "Concurrency lease release channel is closed; the lease will \ - hold capacity until it expires" + "Could not queue a concurrency lease release; it will hold \ + capacity until it expires" ); } } @@ -679,6 +680,7 @@ enum LeaseAdmission { #[derive(Default)] pub(crate) struct HeldLeases { by_model: std::sync::Mutex>>, + dropped_releases: std::sync::atomic::AtomicU64, } /// Whether the store already has the row. Renewal updates stored leases and @@ -711,6 +713,16 @@ impl HeldLeases { } } + fn record_dropped_release(&self) { + self.dropped_releases + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + + fn take_dropped_releases(&self) -> u64 { + self.dropped_releases + .swap(0, std::sync::atomic::Ordering::Relaxed) + } + fn remove(&self, organization_id: Uuid, model_id: Uuid, lease_id: Uuid) { let mut guard = self.by_model.lock().unwrap_or_else(|e| e.into_inner()); if let Some(ids) = guard.get_mut(&(organization_id, model_id)) { @@ -777,7 +789,7 @@ impl HeldLeases { struct FleetConcurrency { repository: Arc, held: Arc, - release: mpsc::UnboundedSender, + release: mpsc::Sender, instance_id: String, ttl: Duration, enforcing: bool, @@ -787,6 +799,10 @@ struct FleetConcurrency { /// before a live request's lease could lapse. const LEASE_RENEWALS_PER_TTL: u32 = 3; +/// Bounded so a stalled lease store cannot let this queue grow without limit. +/// A dropped release costs one lease until its TTL, which the sweeper reclaims. +const RELEASE_QUEUE_DEPTH: usize = 8192; + /// Each of these loops silently stops enforcing something if it exits, so a /// death is reported rather than discarded. fn supervise(task: &'static str, handle: tokio::task::JoinHandle<()>) { @@ -943,7 +959,7 @@ impl CompletionServiceImpl { ttl: Duration, enforcing: bool, ) -> Self { - let (release, mut released) = mpsc::unbounded_channel::(); + let (release, mut released) = mpsc::channel::(RELEASE_QUEUE_DEPTH); let releaser = repository.clone(); let drain = tokio::spawn(async move { @@ -1016,6 +1032,14 @@ impl CompletionServiceImpl { renewing.held_ids().len() as f64, &tags, ); + let dropped = renewing.take_dropped_releases(); + if dropped > 0 { + renew_metrics.record_count( + METRIC_CONCURRENCY_RELEASE_DROPPED, + dropped as i64, + &tags, + ); + } if pending.is_empty() { continue; @@ -1074,10 +1098,13 @@ impl CompletionServiceImpl { "Reclaimed concurrency leases whose holder stopped renewing" ) } - Err(error) => tracing::warn!( - error = %error, - "Failed to sweep expired concurrency leases" - ), + Err(error) => { + sweep_metrics.record_count(METRIC_CONCURRENCY_SWEEP_FAILED, 1, &tags); + tracing::warn!( + error = %error, + "Failed to sweep expired concurrency leases" + ) + } } } }); diff --git a/crates/services/src/completions/ports.rs b/crates/services/src/completions/ports.rs index 207a3c650..f9d00c156 100644 --- a/crates/services/src/completions/ports.rs +++ b/crates/services/src/completions/ports.rs @@ -177,11 +177,11 @@ pub struct HeldLease { pub model_id: Uuid, } -/// Repository trait for fleet-wide concurrency leases -/// Used by CompletionService so replicas admit against one shared count +/// Holds the shared in-flight count so replicas admit against one limit +/// rather than one each. #[async_trait] pub trait ConcurrencyLeaseRepository: Send + Sync { - /// Record a lease if the organization is below its limit for the model + /// Record a lease if the organization is below its limit for the model. async fn try_acquire( &self, lease_id: Uuid, @@ -192,6 +192,7 @@ pub trait ConcurrencyLeaseRepository: Send + Sync { ttl: Duration, ) -> Result; + /// Give back leases whose requests have finished. async fn release(&self, lease_ids: &[Uuid]) -> Result<(), anyhow::Error>; /// Extend leases the store already has, returning the ids still there. diff --git a/crates/services/src/metrics/consts.rs b/crates/services/src/metrics/consts.rs index c3010e280..1f533c752 100644 --- a/crates/services/src/metrics/consts.rs +++ b/crates/services/src/metrics/consts.rs @@ -42,6 +42,8 @@ pub const METRIC_CONCURRENCY_ADMISSION_LATENCY: &str = "cloud_api.concurrency.ad pub const METRIC_CONCURRENCY_IN_FLIGHT: &str = "cloud_api.concurrency.in_flight"; pub const METRIC_CONCURRENCY_RENEW_FAILED: &str = "cloud_api.concurrency.renew_failed"; pub const METRIC_CONCURRENCY_RECLAIMED: &str = "cloud_api.concurrency.reclaimed"; +pub const METRIC_CONCURRENCY_RELEASE_DROPPED: &str = "cloud_api.concurrency.release_dropped"; +pub const METRIC_CONCURRENCY_SWEEP_FAILED: &str = "cloud_api.concurrency.sweep_failed"; // Whether admission was decided fleet-wide or by this replica alone. pub const TAG_SCOPE: &str = "scope"; From 272e2ea398c8264410d4aece79b9bc11a88320ed Mon Sep 17 00:00:00 2001 From: neo-sky Date: Tue, 25 Aug 2026 22:13:26 -0400 Subject: [PATCH 6/8] Merge main and move the lease migration to V0076 Main took V0075 for the orphaned-org-children migration, so refinery applied one and skipped the other. --- ...d_concurrency_leases.sql => V0076__add_concurrency_leases.sql} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename crates/database/src/migrations/sql/{V0075__add_concurrency_leases.sql => V0076__add_concurrency_leases.sql} (100%) diff --git a/crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql b/crates/database/src/migrations/sql/V0076__add_concurrency_leases.sql similarity index 100% rename from crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql rename to crates/database/src/migrations/sql/V0076__add_concurrency_leases.sql From a8d35b90b3a68b7c873c5fac86fea5bf98465393 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Thu, 27 Aug 2026 13:38:51 -0400 Subject: [PATCH 7/8] Keep the per-replica limit in charge while shadowing A fleet-admitted request returned before reaching the local counter, so a replica could run its lease allowance and its local allowance at once. While shadowing the fleet only observes now, including when the lease store is down. --- .../completions/fleet_concurrency_tests.rs | 75 +++++++++++++++++++ crates/services/src/completions/mod.rs | 46 +++++++++--- 2 files changed, 110 insertions(+), 11 deletions(-) diff --git a/crates/services/src/completions/fleet_concurrency_tests.rs b/crates/services/src/completions/fleet_concurrency_tests.rs index f2d32d9b2..9c4d7c323 100644 --- a/crates/services/src/completions/fleet_concurrency_tests.rs +++ b/crates/services/src/completions/fleet_concurrency_tests.rs @@ -496,3 +496,78 @@ async fn dropping_a_transport_guard_frees_fleet_capacity() { .await .expect("the freed slot admits the next transport request"); } + +/// Shadowing must not widen admission. The fleet lease is observation, so one +/// replica still admits no more than the organization limit. +#[tokio::test] +async fn shadowing_does_not_admit_past_the_replica_limit() { + const ATTEMPTS: usize = 20; + + let store = Arc::new(SharedLeaseStore::with_limit(LIMIT)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + let service = replica_with( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(LIMIT))), + false, + ); + + let mut slots = Vec::new(); + for _ in 0..ATTEMPTS { + if let Ok(slot) = service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + { + slots.push(slot); + } + } + + assert_eq!( + slots.len(), + LIMIT as usize, + "one shadowing replica admitted {} against a limit of {LIMIT}", + slots.len() + ); + assert_eq!( + store.live_count(organization_id, model_id), + LIMIT as usize, + "a lease is held for each admitted request and released for the rest" + ); +} + +/// A lease-store outage while shadowing must not switch the decision over to +/// the degraded limiter; the per-replica count still owns admission. +#[tokio::test] +async fn a_store_outage_while_shadowing_still_defers_to_the_replica_limit() { + let store = Arc::new(SharedLeaseStore::with_limit(LIMIT)); + let metrics = Arc::new(CapturingMetricsService::new()); + let organization_id = Uuid::new_v4(); + let model_id = Uuid::new_v4(); + let service = replica_with( + store.clone(), + metrics.clone(), + Arc::new(StaticLimitRepository(Some(LIMIT))), + false, + ); + + store.set_unavailable(true); + + let mut slots = Vec::new(); + for _ in 0..10 { + if let Ok(slot) = service + .try_acquire_concurrent_slot(organization_id, model_id, MODEL_NAME) + .await + { + slots.push(slot); + } + } + + assert_eq!( + slots.len(), + LIMIT as usize, + "shadowing through an outage admitted {} against a limit of {LIMIT}", + slots.len() + ); +} diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index d9f2c799f..cd8873bdb 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -589,17 +589,27 @@ pub(crate) enum ConcurrentSlot { held: Arc, release: mpsc::Sender, }, + Shadowed { + lease: Box, + local: Arc, + }, } impl ConcurrentSlot { + fn release_local(counter: &Arc) { + // Saturating: an underflow here would read as a permanently + // full counter and lock the organization out of the model. + let _ = counter.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + Some(current.saturating_sub(1)) + }); + } + fn release(&self) { match self { - Self::Local(counter) => { - // Saturating: an underflow here would read as a permanently - // full counter and lock the organization out of the model. - let _ = counter.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { - Some(current.saturating_sub(1)) - }); + Self::Local(counter) => Self::release_local(counter), + Self::Shadowed { lease, local } => { + lease.release(); + Self::release_local(local); } Self::Lease { id, @@ -1904,12 +1914,15 @@ impl CompletionServiceImpl { model_id: Uuid, model_name: &str, ) -> Result { + let mut shadow_lease: Option = None; + if let Some(fleet) = &self.fleet_concurrency { match self .try_acquire_lease(fleet, organization_id, model_id, model_name) .await { - Ok(slot) => return Ok(slot), + Ok(slot) if fleet.enforcing => return Ok(slot), + Ok(slot) => shadow_lease = Some(slot), Err(LeaseAdmission::AtLimit(error)) => return Err(error), Err(LeaseAdmission::Shadowed) => {} Err(LeaseAdmission::Unavailable(error)) => { @@ -1920,9 +1933,11 @@ impl CompletionServiceImpl { error = %error, "Fleet concurrency unavailable, falling back to this replica's leases" ); - return self - .admit_from_held_leases(fleet, organization_id, model_id, model_name) - .await; + if fleet.enforcing { + return self + .admit_from_held_leases(fleet, organization_id, model_id, model_name) + .await; + } } } } @@ -1956,13 +1971,22 @@ impl CompletionServiceImpl { &ports::CompletionError::RateLimitExceeded(msg.clone()), Some(model_name), ); + if let Some(lease) = shadow_lease { + lease.release(); + } return Err(ports::CompletionError::RateLimitExceeded(msg)); } if counter .compare_exchange_weak(current, current + 1, Ordering::AcqRel, Ordering::Acquire) .is_ok() { - return Ok(ConcurrentSlot::Local(counter)); + return Ok(match shadow_lease.take() { + Some(lease) => ConcurrentSlot::Shadowed { + lease: Box::new(lease), + local: counter, + }, + None => ConcurrentSlot::Local(counter), + }); } } } From 506dd419e801fb96bf310e640c8435d2d95163f5 Mon Sep 17 00:00:00 2001 From: neo-sky Date: Thu, 27 Aug 2026 14:26:25 -0400 Subject: [PATCH 8/8] Drop the lease a rejected retry already committed A lost response makes retry_db! re-run the acquire under the same id. If the free capacity went elsewhere in between, the rejection left the first attempt's row holding a slot that nothing renewed or released until its TTL. The reject path deletes it now. --- .../src/repositories/concurrency_lease.rs | 6 ++ crates/database/tests/concurrency_leases.rs | 76 +++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/crates/database/src/repositories/concurrency_lease.rs b/crates/database/src/repositories/concurrency_lease.rs index 3ef050d18..73efe4fe2 100644 --- a/crates/database/src/repositories/concurrency_lease.rs +++ b/crates/database/src/repositories/concurrency_lease.rs @@ -84,6 +84,12 @@ impl ConcurrencyLeaseRepository for PostgresConcurrencyLeaseRepository { let in_flight: i64 = row.get("in_flight"); if in_flight >= i64::from(limit) { + // A retry after a lost response finds the row its first attempt + // committed. Rejecting without dropping it holds a slot for a + // request nobody went on to track or release. + tx.execute("DELETE FROM concurrency_leases WHERE id = $1", &[&lease_id]) + .await + .map_err(map_db_error)?; tx.commit().await.map_err(map_db_error)?; return Ok::(LeaseOutcome::AtLimit { limit, diff --git a/crates/database/tests/concurrency_leases.rs b/crates/database/tests/concurrency_leases.rs index 2d6f03448..d699c39cc 100644 --- a/crates/database/tests/concurrency_leases.rs +++ b/crates/database/tests/concurrency_leases.rs @@ -306,6 +306,82 @@ async fn pending_leases_are_written_back() { assert_eq!(live, 1, "a degraded-path lease must reach the store"); } +/// A lost response makes `retry_db!` run the acquire again under the same id. +/// If the remaining capacity went to someone else meanwhile, the rejection has +/// to drop the row the first attempt committed; nothing else would, because the +/// caller takes the rejection and never registers the lease to release it. +#[tokio::test] +async fn a_rejected_retry_drops_the_lease_its_first_attempt_committed() { + let pool = support::test_pool().await.expect("test pool"); + let org = support::insert_org_fixture(&pool) + .await + .expect("org fixture"); + let model = support::insert_model(&pool, &format!("lease-{}", Uuid::new_v4())) + .await + .expect("model fixture"); + + pool.get() + .await + .expect("connection") + .execute( + "UPDATE organizations SET rate_limit = 1 WHERE id = $1", + &[&org.org_id], + ) + .await + .expect("rate limit update"); + + let repository = PostgresConcurrencyLeaseRepository::new(pool.clone()); + let lease_id = Uuid::new_v4(); + let already_committed = [ + HeldLease { + id: lease_id, + organization_id: org.org_id, + model_id: model.id, + }, + HeldLease { + id: Uuid::new_v4(), + organization_id: org.org_id, + model_id: model.id, + }, + ]; + repository + .persist(&already_committed, "instance-a", TTL) + .await + .expect("persist"); + + let outcome = repository + .try_acquire( + lease_id, + org.org_id, + model.id, + "instance-a", + DEFAULT_LIMIT, + TTL, + ) + .await + .expect("acquire"); + assert!( + matches!(outcome, LeaseOutcome::AtLimit { .. }), + "the other lease holds the only slot, got {outcome:?}" + ); + + let surviving: i64 = pool + .get() + .await + .expect("connection") + .query_one( + "SELECT COUNT(*) FROM concurrency_leases WHERE id = $1", + &[&lease_id], + ) + .await + .expect("count") + .get(0); + assert_eq!( + surviving, 0, + "a rejected acquire must not leave its lease holding capacity" + ); +} + /// A zero or negative rate_limit means unset, not a limit of zero. Reading it /// literally would reject every request for the organization. #[tokio::test]