diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 14c2e73f6..c2e74020b 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -440,14 +440,30 @@ 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::Off { + 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), + config.fleet_concurrency.mode == config::FleetConcurrencyMode::Enforce, + ); + } + + let completion_service = Arc::new(completion_service_impl); let brave_search_provider = Arc::new(services::responses::tools::brave::BraveWebSearchProvider::new()); @@ -2861,6 +2877,7 @@ mod tests { aml: config::AmlConfig::default(), usage_reporting: config::UsageReportingConfig::default(), ita: config::ItaAttestationConfig::default(), + fleet_concurrency: config::FleetConcurrencyConfig::default(), }; // Initialize services @@ -2975,6 +2992,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 dbeb1da70..12e6351b7 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -137,6 +137,7 @@ pub fn test_config() -> ApiConfig { ..config::UsageReportingConfig::default() }, ita: config::ItaAttestationConfig::default(), + fleet_concurrency: config::FleetConcurrencyConfig::default(), } } @@ -440,6 +441,36 @@ where (server, 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, +) -> ( + 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 fbb74d62a..e12c35873 100644 --- a/crates/api/tests/e2e_all/main.rs +++ b/crates/api/tests/e2e_all/main.rs @@ -43,6 +43,7 @@ mod external_providers; mod feature_requests; mod files; mod first_stream_event; +mod fleet_concurrency; mod function_tools; mod general; mod glm52_tier_routing; 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 60072e7d4..3895df2fa 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 { @@ -36,6 +37,81 @@ 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, + /// Leases are counted fleet-wide but rejection stays per-replica. + Shadow, + /// 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: 60, + instance_id: String::new(), + } + } +} + +impl FleetConcurrencyConfig { + pub fn from_env() -> Self { + 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 + // 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 { + mode, + 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}-{}", Uuid::new_v4().simple()) + }), + } + } } impl ApiConfig { @@ -78,6 +154,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/V0076__add_concurrency_leases.sql b/crates/database/src/migrations/sql/V0076__add_concurrency_leases.sql new file mode 100644 index 000000000..1fa53cd71 --- /dev/null +++ b/crates/database/src/migrations/sql/V0076__add_concurrency_leases.sql @@ -0,0 +1,16 @@ +-- 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, + 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); diff --git a/crates/database/src/repositories/concurrency_lease.rs b/crates/database/src/repositories/concurrency_lease.rs new file mode 100644 index 000000000..73efe4fe2 --- /dev/null +++ b/crates/database/src/repositories/concurrency_lease.rs @@ -0,0 +1,256 @@ +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, HeldLease, 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 } + } + + /// 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) + } +} + +#[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() + AND id <> $3 + ) AS in_flight + "#, + &[&organization_id, &model_id, &lease_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) { + // 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, + 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 { limit }) + })?; + + Ok(outcome) + } + + async fn renew(&self, lease_ids: &[Uuid], ttl: Duration) -> Result> { + if lease_ids.is_empty() { + return Ok(Vec::new()); + } + + let ttl_seconds = ttl.as_secs() as f64; + + let rows = retry_db!("renew_concurrency_leases", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + client + .query( + r#" + UPDATE concurrency_leases + SET expires_at = NOW() + make_interval(secs => $2) + WHERE id = ANY($1) + RETURNING id + "#, + &[&lease_ids, &ttl_seconds], + ) + .await + .map_err(map_db_error) + })?; + + Ok(rows.iter().map(|row| row.get("id")).collect()) + } + + 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)?; + + // Bounded so a backlog cannot become one long delete holding row + // locks against live admissions; the next tick takes the rest. + client + .execute( + r#" + DELETE FROM concurrency_leases + WHERE id = ANY( + SELECT id FROM concurrency_leases + WHERE expires_at < NOW() + LIMIT 1000 + ) + "#, + &[], + ) + .await + .map_err(map_db_error) + })?; + + Ok(removed) + } + + 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..d699c39cc --- /dev/null +++ b/crates/database/tests/concurrency_leases.rs @@ -0,0 +1,547 @@ +// Shared with other test binaries, which use a different subset of it. +#[allow(dead_code)] +mod support; + +use database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository; +use services::completions::ports::{ConcurrencyLeaseRepository, HeldLease, 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!(matches!(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!( + matches!(readmitted, LeaseOutcome::Admitted { .. }), + "releasing a lease must free exactly one slot" + ); +} + +/// 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!(matches!(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"); + 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() + .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"); +} + +/// 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] +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] +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!(matches!(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!( + matches!(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 03e8b3e4c..5be610384 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`] @@ -239,6 +239,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 { @@ -254,9 +255,16 @@ impl ResponseTemplate { cache_write_tokens: None, service_tier_override: 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 { @@ -1097,6 +1105,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 @@ -1224,6 +1236,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/fleet_concurrency_tests.rs b/crates/services/src/completions/fleet_concurrency_tests.rs new file mode 100644 index 000000000..9c4d7c323 --- /dev/null +++ b/crates/services/src/completions/fleet_concurrency_tests.rs @@ -0,0 +1,573 @@ +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 { + 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, + 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), + enforcing, + ) +} + +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" + ); +} + +/// 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"); +} + +/// 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 152796316..18140a4f7 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -10,8 +10,10 @@ use crate::usage::{ }; 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; use uuid::Uuid; // Create a new stream that intercepts messages, but passes the original ones through @@ -116,7 +118,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) @@ -690,9 +692,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(); } // Every successful finalization path reaches `StreamState::Done` only @@ -735,32 +736,84 @@ where } } +/// 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 { + id: Uuid, + organization_id: Uuid, + model_id: Uuid, + 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) => Self::release_local(counter), + Self::Shadowed { lease, local } => { + lease.release(); + Self::release_local(local); + } + Self::Lease { + id, + organization_id, + model_id, + held, + release, + } => { + held.remove(*organization_id, *model_id, *id); + if release.try_send(*id).is_err() { + held.record_dropped_release(); + tracing::error!( + lease_id = %id, + "Could not queue a concurrency lease release; it will hold \ + capacity until it expires" + ); + } + } + } + } +} + /// 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(); } } } @@ -775,13 +828,178 @@ 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, +} + +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), +} + +/// 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>>, + dropped_releases: std::sync::atomic::AtomicU64, +} + +/// 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, + 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 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)) { + ids.remove(&lease_id); + if ids.is_empty() { + guard.remove(&(organization_id, model_id)); + } + } + } + + #[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 + .get(&(organization_id, model_id)) + .map_or(0, |ids| ids.len()) + } + + /// 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, + 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::Sender, + instance_id: String, + ttl: Duration, + enforcing: bool, +} + +/// 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; + +/// 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<()>) { + 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; + +/// 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. @@ -891,10 +1109,187 @@ 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, } } + /// 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, + enforcing: bool, + ) -> Self { + let (release, mut released) = mpsc::channel::(RELEASE_QUEUE_DEPTH); + let releaser = repository.clone(); + + 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 { + tracing::warn!( + released = batch.len(), + error = %error, + "Failed to release concurrency leases; they will expire instead" + ); + } + batch.clear(); + } + }); + + 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_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 { + ticker.tick().await; + let (stored, pending) = renewing.snapshot(); + + let stored_ids: Vec = stored.iter().map(|lease| lease.id).collect(); + 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, + ); + 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; + } + 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_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) => { + sweep_metrics.record_count( + METRIC_CONCURRENCY_RECLAIMED, + removed as i64, + &tags, + ); + tracing::info!( + removed = removed, + "Reclaimed concurrency leases whose holder stopped renewing" + ) + } + Err(error) => { + sweep_metrics.record_count(METRIC_CONCURRENCY_SWEEP_FAILED, 1, &tags); + 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, + enforcing, + }); + self + } + /// Extract tools and tool_choice from the extra HashMap if present and /// parseable as the typed `ToolDefinition` / `ToolChoice` shapes. /// @@ -1052,22 +1447,42 @@ 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), + } + } + + /// 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 } /// Create low-cardinality metric tags for a request @@ -1441,6 +1856,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 @@ -1479,6 +1901,137 @@ impl CompletionServiceImpl { .collect() } + /// 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, + 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(()) => { + 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, + 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, + organization_id: Uuid, + model_id: Uuid, + model_name: &str, + ) -> Result { + let lease_id = Uuid::new_v4(); + let started = Instant::now(); + 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)?; + 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 { 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); + Ok(ConcurrentSlot::Lease { + id: lease_id, + organization_id, + model_id, + held: fleet.held.clone(), + release: fleet.release.clone(), + }) + } + 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, + 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), + )) + } + } + } + fn normalize_profiled_service_tier( model: &crate::models::ModelWithPricing, params: &mut inference_providers::ChatCompletionParams, @@ -1518,7 +2071,35 @@ impl CompletionServiceImpl { organization_id: Uuid, model_id: Uuid, model_name: &str, - ) -> Result, ports::CompletionError> { + ) -> 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) 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)) => { + self.record_concurrency(METRIC_CONCURRENCY_DEGRADED, model_name, SCOPE_REPLICA); + tracing::warn!( + organization_id = %organization_id, + model_id = %model_id, + error = %error, + "Fleet concurrency unavailable, falling back to this replica's leases" + ); + if fleet.enforcing { + return self + .admit_from_held_leases(fleet, organization_id, model_id, model_name) + .await; + } + } + } + } + // Get the dynamic limit for this organization (cached with 5-min TTL) let limit = self.get_org_concurrent_limit(organization_id).await; @@ -1548,13 +2129,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(counter); + return Ok(match shadow_lease.take() { + Some(lease) => ConcurrentSlot::Shadowed { + lease: Box::new(lease), + local: counter, + }, + None => ConcurrentSlot::Local(counter), + }); } } } @@ -1573,7 +2163,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, @@ -1614,7 +2204,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, @@ -1755,13 +2345,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, @@ -1949,12 +2539,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, @@ -2142,12 +2732,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 @@ -2206,12 +2796,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 @@ -2258,10 +2848,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) @@ -2306,10 +2896,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) @@ -2355,12 +2945,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 @@ -2422,9 +3012,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::*; @@ -2462,7 +3170,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, @@ -2679,7 +3387,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, @@ -2852,7 +3560,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, @@ -3006,7 +3714,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, @@ -3134,7 +3842,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, @@ -3345,7 +4053,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 00e930ac0..0e85c8bb7 100644 --- a/crates/services/src/completions/ports.rs +++ b/crates/services/src/completions/ports.rs @@ -3,10 +3,7 @@ use crate::UserId; use async_trait::async_trait; use inference_providers::StreamingResult; use serde::{Deserialize, Serialize}; -use std::sync::{ - atomic::{AtomicU32, Ordering}, - Arc, -}; +use std::time::Duration; use uuid::Uuid; /// Default concurrent request limit per organization per model @@ -16,20 +13,25 @@ pub const DEFAULT_CONCURRENT_LIMIT: u32 = 64; /// /// Routes that bypass the typed completion request path can hold this guard /// until their upstream response body completes or is dropped. -#[derive(Debug)] pub struct ConcurrentRequestGuard { - counter: Arc, + slot: super::ConcurrentSlot, +} + +impl std::fmt::Debug for ConcurrentRequestGuard { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("ConcurrentRequestGuard") + } } impl ConcurrentRequestGuard { - pub(crate) fn new(counter: Arc) -> Self { - Self { counter } + pub(crate) fn new(slot: super::ConcurrentSlot) -> Self { + Self { slot } } } impl Drop for ConcurrentRequestGuard { fn drop(&mut self) { - self.counter.fetch_sub(1, Ordering::Release); + self.slot.release(); } } @@ -162,6 +164,56 @@ 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 { limit: u32 }, + 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, +} + +/// 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. + async fn try_acquire( + &self, + lease_id: Uuid, + organization_id: Uuid, + model_id: Uuid, + instance_id: &str, + default_limit: u32, + 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. + /// 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, from requests admitted while + /// it was unreachable. + 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] pub trait CompletionServiceTrait: Send + Sync { /// Acquire one organization/model concurrency slot for a direct transport diff --git a/crates/services/src/metrics/consts.rs b/crates/services/src/metrics/consts.rs index 94bbcf15b..afd230552 100644 --- a/crates/services/src/metrics/consts.rs +++ b/crates/services/src/metrics/consts.rs @@ -29,6 +29,27 @@ 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"; +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. +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"; +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"; +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";