From 63ca0f9419e97cc494162f2a3ecf92bf56622646 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:58:52 +0000 Subject: [PATCH 01/10] perf(attestation): write both chat signatures in one statement The provider (TEE) and gateway signature paths each stored the ecdsa and ed25519 signatures with two separate INSERT ... ON CONFLICT statements. On the streaming path this runs before the client receives [DONE], so on a replica far from the database every statement is a user-visible round trip (verification p50 is 38 ms on cpu01 and 377 ms on cpu02). AttestationRepository gains add_chat_signatures with a default one-by-one implementation, and the Postgres repository overrides it with a single multi-row upsert using the same conflict clause. Both callers collect the signatures first and write once. --- .../database/src/repositories/attestation.rs | 61 +++++++++++++++++++ .../src/attestation/chat_signatures.rs | 54 ++++++++-------- .../src/attestation/gateway_signatures.rs | 48 +++++++-------- crates/services/src/attestation/ports.rs | 14 +++++ 4 files changed, 124 insertions(+), 53 deletions(-) diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index e25197334..d3919fe0c 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -72,6 +72,67 @@ impl AttestationRepository for PgAttestationRepository { Ok(()) } + /// One round trip for all signatures of a chat: the streaming path stores + /// an ecdsa and an ed25519 signature before the client sees `[DONE]`, and + /// on a replica far from the database each statement is a full network + /// round trip. + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + if signatures.is_empty() { + return Ok(()); + } + let client = self + .pool + .get() + .await + .map_err(|e| AttestationError::RepositoryError(e.to_string()))?; + + let signature_kinds: Vec> = signatures + .iter() + .map(|signature| signature.signature_kind.map(|kind| kind.as_str())) + .collect(); + let mut sql = String::from( + "INSERT INTO chat_signatures (chat_id, text, signature, signing_address, signing_algo, signature_kind) VALUES ", + ); + let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); + for (index, (signature, signature_kind)) in + signatures.iter().zip(signature_kinds.iter()).enumerate() + { + if index > 0 { + sql.push_str(", "); + } + let base = index * 6; + sql.push_str(&format!( + "(${}, ${}, ${}, ${}, ${}, ${})", + base + 1, + base + 2, + base + 3, + base + 4, + base + 5, + base + 6 + )); + params.push(&chat_id); + params.push(&signature.text); + params.push(&signature.signature); + params.push(&signature.signing_address); + params.push(&signature.signing_algo); + params.push(signature_kind); + } + sql.push_str( + " ON CONFLICT (chat_id, signing_algo) DO UPDATE SET text = EXCLUDED.text, signature = EXCLUDED.signature, signing_address = EXCLUDED.signing_address, signature_kind = EXCLUDED.signature_kind, updated_at = NOW()", + ); + + client + .execute(sql.as_str(), ¶ms) + .await + .map_err(|e| AttestationError::RepositoryError(e.to_string()))?; + + Ok(()) + } + async fn get_chat_signature( &self, chat_id: &str, diff --git a/crates/services/src/attestation/chat_signatures.rs b/crates/services/src/attestation/chat_signatures.rs index 6fa936809..971f12961 100644 --- a/crates/services/src/attestation/chat_signatures.rs +++ b/crates/services/src/attestation/chat_signatures.rs @@ -76,6 +76,7 @@ impl AttestationService { let environment = get_environment(); let env_tag = format!("{TAG_ENVIRONMENT}:{environment}"); let result: Result<(), AttestationError> = async { + let mut signatures = Vec::with_capacity(2); for algo in ["ecdsa", "ed25519"] { let provider_signature = provider .get_signature(chat_id, Some(algo.to_string())) @@ -103,38 +104,39 @@ impl AttestationService { ); AttestationError::ProviderError(e.to_string()) })?; - let signature = ChatSignature { + signatures.push(ChatSignature { text: provider_signature.text, signature: provider_signature.signature, signing_address: provider_signature.signing_address, signing_algo: provider_signature.signing_algo, signature_kind: Some(SignatureKind::ProviderTee), - }; - - self.repository - .add_chat_signature(chat_id, signature) - .await - .map_err(|e| { - tracing::error!( - %chat_id, - error = %e, - "Failed to store chat signature in repository for algorithm: {}", - algo - ); - let duration = start_time.elapsed(); - self.metrics_service.record_count( - METRIC_VERIFICATION_FAILURE, - 1, - &[&format!("{TAG_REASON}:{REASON_REPOSITORY_ERROR}"), &env_tag], - ); - self.metrics_service.record_latency( - METRIC_VERIFICATION_DURATION, - duration, - &[&env_tag], - ); - AttestationError::RepositoryError(e.to_string()) - })?; + }); } + + // Both algorithms land in one statement: this runs before the + // client sees `[DONE]`, so every saved round trip is user-visible. + self.repository + .add_chat_signatures(chat_id, signatures) + .await + .map_err(|e| { + tracing::error!( + %chat_id, + error = %e, + "Failed to store chat signatures in repository" + ); + let duration = start_time.elapsed(); + self.metrics_service.record_count( + METRIC_VERIFICATION_FAILURE, + 1, + &[&format!("{TAG_REASON}:{REASON_REPOSITORY_ERROR}"), &env_tag], + ); + self.metrics_service.record_latency( + METRIC_VERIFICATION_DURATION, + duration, + &[&env_tag], + ); + AttestationError::RepositoryError(e.to_string()) + })?; Ok(()) } .await; diff --git a/crates/services/src/attestation/gateway_signatures.rs b/crates/services/src/attestation/gateway_signatures.rs index 630d7cd2e..c16479ceb 100644 --- a/crates/services/src/attestation/gateway_signatures.rs +++ b/crates/services/src/attestation/gateway_signatures.rs @@ -18,6 +18,7 @@ impl AttestationService { let env_tag = format!("{TAG_ENVIRONMENT}:{environment}"); let signature_text = format!("{request_hash}:{response_hash}"); + let mut signatures = Vec::with_capacity(2); for algo in ["ecdsa", "ed25519"] { let (signature_hex, signing_address) = match algo { "ed25519" => self.sign_ed25519_gateway_signature(&signature_text), @@ -26,35 +27,28 @@ impl AttestationService { "Unknown signing algorithm: {algo}" ))), }?; - - self.repository - .add_chat_signature( - signature_id, - ChatSignature { - text: signature_text.clone(), - signature: signature_hex, - signing_address, - signing_algo: algo.to_string(), - signature_kind: Some(SignatureKind::Gateway), - }, - ) - .await - .map_err(|e| { - tracing::error!( - "Failed to store {} signature in repository for algorithm: {}", - id_label, - algo - ); - AttestationError::RepositoryError(e.to_string()) - })?; - tracing::info!( - signature_kind = id_label, - signature_id = signature_id, - signing_algo = algo, - "Stored gateway signature" - ); + signatures.push(ChatSignature { + text: signature_text.clone(), + signature: signature_hex, + signing_address, + signing_algo: algo.to_string(), + signature_kind: Some(SignatureKind::Gateway), + }); } + self.repository + .add_chat_signatures(signature_id, signatures) + .await + .map_err(|e| { + tracing::error!("Failed to store {} signatures in repository", id_label); + AttestationError::RepositoryError(e.to_string()) + })?; + tracing::info!( + signature_kind = id_label, + signature_id = signature_id, + "Stored gateway signatures" + ); + let duration = start_time.elapsed(); self.metrics_service .record_count(METRIC_SIGNATURE_CREATION_SUCCESS, 1, &[&env_tag]); diff --git a/crates/services/src/attestation/ports.rs b/crates/services/src/attestation/ports.rs index a8fdac25a..e1a1599e3 100644 --- a/crates/services/src/attestation/ports.rs +++ b/crates/services/src/attestation/ports.rs @@ -101,6 +101,20 @@ pub trait AttestationRepository: Send + Sync { chat_id: &str, signature: ChatSignature, ) -> Result<(), AttestationError>; + + /// Store several signatures for one chat id. Backends that can write them + /// in a single statement should override this; the default stores them + /// one by one and stops at the first failure. + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + for signature in signatures { + self.add_chat_signature(chat_id, signature).await?; + } + Ok(()) + } async fn get_chat_signature( &self, chat_id: &str, From e06400008e40cdb83b7c6571260133b646ee9a47 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:04:20 +0000 Subject: [PATCH 02/10] perf(database): reuse prepared statements on the inference hot path tokio-postgres prepares a raw SQL string on every call, so each query on the chat-completions path was two network round trips: Parse/Describe, then Bind/Execute. The eleven statements that run for every inference request (API-key auth and workspace/organization lookup, per-key spend, staking source, balance, limits, model resolution, model-by-id, the billing transaction and the chat-signature upsert) now go through deadpool's per-connection statement cache, which makes each of them one round trip. On the cpu02 replicas, ~61 ms from the Postgres leader, that is roughly 0.6 s less database waiting per request. A cached statement cannot survive a change to its result shape: after a rolling deploy migrates the schema, a replica still on the old release gets SQLSTATE 0A000 "cached plan must not change result type" on its SELECT * statements. The CachedStatements helpers drop the stale statement from the connection cache on that error and map_db_error classifies it as retryable, so retry_db! re-prepares and re-runs it. An e2e test adds a column to workspaces under warm caches and checks that requests keep succeeding; it fails with a 500 when the recovery is disabled. --- .config/nextest.toml | 2 +- crates/api/tests/e2e_all/main.rs | 1 + .../tests/e2e_all/prepared_statement_cache.rs | 84 ++++++++++++ crates/database/src/repositories/api_key.rs | 5 +- .../database/src/repositories/attestation.rs | 3 +- crates/database/src/repositories/mod.rs | 1 + crates/database/src/repositories/model.rs | 5 +- .../src/repositories/organization_limits.rs | 3 +- .../organization_staking_farm_sources.rs | 3 +- .../src/repositories/organization_usage.rs | 11 +- .../src/repositories/statement_cache.rs | 127 ++++++++++++++++++ crates/database/src/repositories/utils.rs | 10 ++ crates/database/src/repositories/workspace.rs | 3 +- 13 files changed, 244 insertions(+), 14 deletions(-) create mode 100644 crates/api/tests/e2e_all/prepared_statement_cache.rs create mode 100644 crates/database/src/repositories/statement_cache.rs diff --git a/.config/nextest.toml b/.config/nextest.toml index 41628550e..5051225a4 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -32,7 +32,7 @@ max-threads = 4 # These tests intentionally touch database-wide state (DDL, broad encryption # scans, or ACCESS EXCLUSIVE locks), so they must not overlap any other test. [[profile.default.overrides]] -filter = 'package(api) & binary(e2e_all) & (test(/^database_encryption::/) | test(/^reporting_usage::usage_export::usage_export_database_timeout_returns_504_and_stops_query$/) | test(/^reporting_usage::usage_summary::usage_summary_database_timeout_returns_504$/))' +filter = 'package(api) & binary(e2e_all) & (test(/^database_encryption::/) | test(/^prepared_statement_cache::/) | test(/^reporting_usage::usage_export::usage_export_database_timeout_returns_504_and_stops_query$/) | test(/^reporting_usage::usage_summary::usage_summary_database_timeout_returns_504$/))' test-group = "e2e-db" threads-required = "num-test-threads" diff --git a/crates/api/tests/e2e_all/main.rs b/crates/api/tests/e2e_all/main.rs index e0684b5c3..e107533ba 100644 --- a/crates/api/tests/e2e_all/main.rs +++ b/crates/api/tests/e2e_all/main.rs @@ -66,6 +66,7 @@ mod org_system_prompt; mod organization_deletion; mod pagination_validation; mod patroni_failover; +mod prepared_statement_cache; mod privacy_classify; mod privacy_redact; mod provider_errors; diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs new file mode 100644 index 000000000..254d8638a --- /dev/null +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -0,0 +1,84 @@ +//! Hot-path queries run through the connection's prepared-statement cache. +//! During a rolling deploy a replica on the previous release keeps serving +//! after the new release has run its migrations; if one of those migrations +//! adds a column to a table read with `SELECT *`, Postgres refuses the old +//! plan with "cached plan must not change result type". The repositories must +//! recover from that transparently, not surface a 500. + +use crate::common::*; + +#[tokio::test] +async fn cached_statements_recover_from_a_column_added_under_them() { + let (server, database) = setup_test_server_with_database().await; + let (session_id, _email) = setup_unique_test_session(&database).await; + let org = create_org_with_session(&server, &session_id).await; + let api_key = get_api_key_for_org_with_session(&server, org.id.clone(), &session_id).await; + + let probe = || async { + server + .get("/v1/files?limit=1") + .add_header("Authorization", format!("Bearer {api_key}")) + .add_header("User-Agent", MOCK_USER_AGENT) + .await + }; + + // Warm the cache: API-key auth resolves the workspace and organization + // with a `SELECT w.*` on every request, through whichever pooled + // connection serves it. Several requests seed several connections. + for _ in 0..8 { + let response = probe().await; + assert_eq!(response.status_code(), 200, "{}", response.text()); + } + + // A migration on the next release adds a column to `workspaces` while + // this process still holds statements prepared against the old shape. + let column = format!( + "rolling_deploy_{}", + uuid::Uuid::new_v4().simple().to_string()[..8].to_string() + ); + { + let client = database + .pool() + .get() + .await + .expect("failed to get database connection"); + client + .execute( + &format!("ALTER TABLE workspaces ADD COLUMN {column} TEXT"), + &[], + ) + .await + .expect("failed to add column"); + } + + // Every request must still succeed: the stale plan is dropped from the + // cache and the statement re-prepared inside the repository retry. + // Collect instead of asserting so the column is dropped even on failure; + // a leftover column would fail the database-encryption classification + // scans that share this database. + let mut outcomes = Vec::with_capacity(8); + for _ in 0..8 { + let response = probe().await; + outcomes.push((response.status_code(), response.text())); + } + + let client = database + .pool() + .get() + .await + .expect("failed to get database connection"); + client + .execute( + &format!("ALTER TABLE workspaces DROP COLUMN IF EXISTS {column}"), + &[], + ) + .await + .expect("failed to drop column"); + + for (status, body) in outcomes { + assert_eq!( + status, 200, + "request after a schema change must recover, got: {body}" + ); + } +} diff --git a/crates/database/src/repositories/api_key.rs b/crates/database/src/repositories/api_key.rs index d587eede9..e6b885087 100644 --- a/crates/database/src/repositories/api_key.rs +++ b/crates/database/src/repositories/api_key.rs @@ -1,5 +1,6 @@ use crate::models::ApiKey; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -149,7 +150,7 @@ impl ApiKeyRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT ak.* FROM api_keys ak @@ -192,7 +193,7 @@ impl ApiKeyRepository { .map_err(RepositoryError::PoolError)?; client - .execute( + .cached_execute( "UPDATE api_keys SET last_used_at = NOW() WHERE id = $1", &[&id], ) diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index e25197334..c38c7389a 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -1,3 +1,4 @@ +use crate::repositories::statement_cache::CachedStatements; use async_trait::async_trait; use services::attestation::{ ports::AttestationRepository, AttestationError, ChatSignature, SignatureKind, @@ -62,7 +63,7 @@ impl AttestationRepository for PgAttestationRepository { .map_err(|e| AttestationError::RepositoryError(e.to_string()))?; let signature_kind = signature.signature_kind.map(|kind| kind.as_str()); client - .execute( + .cached_execute( "INSERT INTO chat_signatures (chat_id, text, signature, signing_address, signing_algo, signature_kind) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (chat_id, signing_algo) DO UPDATE SET text = EXCLUDED.text, signature = EXCLUDED.signature, signing_address = EXCLUDED.signing_address, signature_kind = EXCLUDED.signature_kind, updated_at = NOW()", &[&chat_id, &signature.text, &signature.signature, &signature.signing_address, &signature.signing_algo, &signature_kind], ) diff --git a/crates/database/src/repositories/mod.rs b/crates/database/src/repositories/mod.rs index f5411510a..37db640ad 100644 --- a/crates/database/src/repositories/mod.rs +++ b/crates/database/src/repositories/mod.rs @@ -33,6 +33,7 @@ pub mod retry; pub mod service; pub mod service_usage_repository_impl; pub mod session; +pub mod statement_cache; pub mod usage_repository_impl; pub mod user; pub mod utils; diff --git a/crates/database/src/repositories/model.rs b/crates/database/src/repositories/model.rs index 3b41eb40d..be6428cae 100644 --- a/crates/database/src/repositories/model.rs +++ b/crates/database/src/repositories/model.rs @@ -1,6 +1,7 @@ use crate::constants::DEFAULT_MODEL_OWNED_BY; use crate::models::{Model, ModelHistory, UpdateModelPricingRequest}; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -269,7 +270,7 @@ impl ModelRepository { .map_err(RepositoryError::PoolError)?; client - .query( + .cached_query( r#" SELECT id, model_name, model_display_name, model_description, model_icon, @@ -1301,7 +1302,7 @@ impl ModelRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT m.id, diff --git a/crates/database/src/repositories/organization_limits.rs b/crates/database/src/repositories/organization_limits.rs index f56c52e94..3bf05ed60 100644 --- a/crates/database/src/repositories/organization_limits.rs +++ b/crates/database/src/repositories/organization_limits.rs @@ -1,5 +1,6 @@ use crate::models::{OrganizationLimitsHistory, UpdateOrganizationLimitsDbRequest}; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -122,7 +123,7 @@ impl OrganizationLimitsRepository { .map_err(RepositoryError::PoolError)?; client - .query( + .cached_query( r#" SELECT id, organization_id, spend_limit, credit_type, source, currency, effective_from, effective_until, diff --git a/crates/database/src/repositories/organization_staking_farm_sources.rs b/crates/database/src/repositories/organization_staking_farm_sources.rs index edace0482..817d7e333 100644 --- a/crates/database/src/repositories/organization_staking_farm_sources.rs +++ b/crates/database/src/repositories/organization_staking_farm_sources.rs @@ -1,4 +1,5 @@ use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -165,7 +166,7 @@ impl StakingFarmRepository for OrganizationStakingFarmSourcesRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT id, organization_id, near_account_id, network_id, contract_id, farm_product_id, farm_price_id, credit_nano_usd_per_reward_unit, diff --git a/crates/database/src/repositories/organization_usage.rs b/crates/database/src/repositories/organization_usage.rs index ae3cba21f..a14e8bc54 100644 --- a/crates/database/src/repositories/organization_usage.rs +++ b/crates/database/src/repositories/organization_usage.rs @@ -3,6 +3,7 @@ use crate::models::{ ServedProviderType, StopReason, }; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -47,7 +48,7 @@ impl OrganizationUsageRepository { .map_err(RepositoryError::PoolError)?; client - .query_one( + .cached_query_one( r#" SELECT COALESCE(SUM(total_cost), 0)::BIGINT as total_spend FROM organization_usage_log @@ -93,7 +94,7 @@ impl OrganizationUsageRepository { .served_provider_type .map(|provider| provider.as_str()); let maybe_row = transaction - .query_opt( + .cached_query_opt( r#" INSERT INTO organization_usage_log ( id, organization_id, workspace_id, api_key_id, @@ -146,7 +147,7 @@ impl OrganizationUsageRepository { Some(row) => { // New insert succeeded — update organization balance transaction - .execute( + .cached_execute( r#" INSERT INTO organization_balance ( organization_id, @@ -188,7 +189,7 @@ impl OrganizationUsageRepository { ); let existing = client - .query_one( + .cached_query_one( r#" SELECT * FROM organization_usage_log @@ -220,7 +221,7 @@ impl OrganizationUsageRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT organization_id, total_spent, last_usage_at, total_requests, total_tokens, updated_at diff --git a/crates/database/src/repositories/statement_cache.rs b/crates/database/src/repositories/statement_cache.rs new file mode 100644 index 000000000..750ddc1cf --- /dev/null +++ b/crates/database/src/repositories/statement_cache.rs @@ -0,0 +1,127 @@ +//! Per-connection prepared-statement reuse for hot-path queries. +//! +//! `tokio_postgres` prepares a raw SQL string on every call, so each query is +//! two network round trips (Parse/Describe, then Bind/Execute). deadpool keeps +//! a per-connection statement cache; going through it makes a repeated query +//! one round trip. On a replica far from the database this is the difference +//! between ~60 ms and ~120 ms per query. +//! +//! The one thing a cached statement cannot survive is a change to its result +//! shape: after a migration adds a column to a table read with `SELECT *`, +//! Postgres rejects the old plan with SQLSTATE 0A000 "cached plan must not +//! change result type". That happens on replicas still running the previous +//! release during a rolling deploy. The helpers below drop the stale statement +//! from the connection's cache when they see that error, and +//! [`map_db_error`](super::utils::map_db_error) classifies it as retryable so +//! the surrounding `retry_db!` block re-prepares and re-runs the statement. + +use async_trait::async_trait; +use deadpool_postgres::{Client, Transaction}; +use tokio_postgres::{error::SqlState, types::ToSql, Error, Row}; + +const STALE_PLAN_MESSAGE: &str = "cached plan must not change result type"; + +/// True when `err` is Postgres refusing a prepared statement whose result +/// columns changed since it was prepared. +pub fn is_stale_cached_plan(err: &Error) -> bool { + err.as_db_error().is_some_and(|db_err| { + db_err.code() == &SqlState::FEATURE_NOT_SUPPORTED + && db_err.message().contains(STALE_PLAN_MESSAGE) + }) +} + +/// Run queries through the connection's prepared-statement cache. +#[async_trait] +pub trait CachedStatements { + async fn cached_query( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error>; + + async fn cached_query_opt( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error>; + + async fn cached_query_one( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result; + + async fn cached_execute(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) + -> Result; +} + +macro_rules! impl_cached_statements { + ($ty:ty) => { + #[async_trait] + impl CachedStatements for $ty { + async fn cached_query( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error> { + let statement = self.prepare_cached(sql).await?; + let result = self.query(&statement, params).await; + self.forget_if_stale(sql, &result); + result + } + + async fn cached_query_opt( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error> { + let statement = self.prepare_cached(sql).await?; + let result = self.query_opt(&statement, params).await; + self.forget_if_stale(sql, &result); + result + } + + async fn cached_query_one( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result { + let statement = self.prepare_cached(sql).await?; + let result = self.query_one(&statement, params).await; + self.forget_if_stale(sql, &result); + result + } + + async fn cached_execute( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result { + let statement = self.prepare_cached(sql).await?; + let result = self.execute(&statement, params).await; + self.forget_if_stale(sql, &result); + result + } + } + + impl StaleStatementCleanup for $ty { + fn forget_if_stale(&self, sql: &str, result: &Result) { + if let Err(err) = result { + if is_stale_cached_plan(err) { + tracing::warn!( + "Dropping prepared statement whose result shape changed; retrying" + ); + self.statement_cache.remove(sql, &[]); + } + } + } + } + }; +} + +trait StaleStatementCleanup { + fn forget_if_stale(&self, sql: &str, result: &Result); +} + +impl_cached_statements!(Client); +impl_cached_statements!(Transaction<'_>); diff --git a/crates/database/src/repositories/utils.rs b/crates/database/src/repositories/utils.rs index 6bf50cd0b..a3ab42edc 100644 --- a/crates/database/src/repositories/utils.rs +++ b/crates/database/src/repositories/utils.rs @@ -41,6 +41,16 @@ pub fn map_db_error(err: tokio_postgres::Error) -> RepositoryError { &SqlState::QUERY_CANCELED => RepositoryError::QueryTimeout, + // A cached prepared statement outlived a schema change (rolling + // deploy). The statement has already been dropped from the + // connection's cache (see `statement_cache`); classify as + // retryable so `retry_db!` re-prepares and re-runs it. + _ if super::statement_cache::is_stale_cached_plan(&err) => { + RepositoryError::ConnectionFailed( + "prepared statement invalidated by a schema change".to_string(), + ) + } + // Default case - wrap in generic database error _ => RepositoryError::DatabaseError(anyhow::anyhow!( "Database error ({}): {}", diff --git a/crates/database/src/repositories/workspace.rs b/crates/database/src/repositories/workspace.rs index 92f24b1ab..6302b2e76 100644 --- a/crates/database/src/repositories/workspace.rs +++ b/crates/database/src/repositories/workspace.rs @@ -1,3 +1,4 @@ +use crate::repositories::statement_cache::CachedStatements; use crate::{ models::{CreateWorkspaceRequest, UpdateWorkspaceRequest, Workspace}, pool::DbPool, @@ -380,7 +381,7 @@ impl WorkspaceRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT w.*, From 975296066e15aeeb059e83128be24827fea2eb14 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:16:35 +0000 Subject: [PATCH 03/10] review: tolerate duplicate algorithms and keep partial signature writes - A multi-row ON CONFLICT DO UPDATE raises 21000 when two rows of the same statement share the arbiter key; signing_algo comes from the backend, so the repository now collapses duplicates and keeps the last one, matching the previous one-at-a-time overwrite. Unit-tested. - The provider path stores whatever was fetched before propagating a later fetch error, so a backend that serves ecdsa but fails ed25519 still leaves the chat verifiable by ecdsa, as it did before. --- .../database/src/repositories/attestation.rs | 63 +++++++++++++++++++ .../src/attestation/chat_signatures.rs | 26 +++++--- 2 files changed, 81 insertions(+), 8 deletions(-) diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index d3919fe0c..f27fa35bc 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -48,6 +48,27 @@ impl PgAttestationRepository { } } +/// Collapse signatures that share a `signing_algo`, keeping the last one. +/// +/// A multi-row `INSERT ... ON CONFLICT DO UPDATE` raises SQLSTATE 21000 +/// ("cannot affect row a second time") when two rows of the same statement +/// hit the same arbiter key, whereas writing them one at a time simply let +/// the later row overwrite the earlier. `signing_algo` is supplied by the +/// inference backend, so the batch must tolerate duplicates the same way. +fn last_signature_per_algo(signatures: Vec) -> Vec { + let mut deduped: Vec = Vec::with_capacity(signatures.len()); + for signature in signatures { + match deduped + .iter_mut() + .find(|existing| existing.signing_algo == signature.signing_algo) + { + Some(existing) => *existing = signature, + None => deduped.push(signature), + } + } + deduped +} + #[async_trait] impl AttestationRepository for PgAttestationRepository { async fn add_chat_signature( @@ -81,6 +102,7 @@ impl AttestationRepository for PgAttestationRepository { chat_id: &str, signatures: Vec, ) -> Result<(), AttestationError> { + let signatures = last_signature_per_algo(signatures); if signatures.is_empty() { return Ok(()); } @@ -164,3 +186,44 @@ impl AttestationRepository for PgAttestationRepository { self.row_to_chat_signature(row) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn signature(algo: &str, text: &str) -> ChatSignature { + ChatSignature { + text: text.to_string(), + signature: format!("sig-{text}"), + signing_address: "addr".to_string(), + signing_algo: algo.to_string(), + signature_kind: None, + } + } + + #[test] + fn distinct_algorithms_are_kept_in_order() { + let out = last_signature_per_algo(vec![signature("ecdsa", "a"), signature("ed25519", "b")]); + assert_eq!(out.len(), 2); + assert_eq!(out[0].signing_algo, "ecdsa"); + assert_eq!(out[1].signing_algo, "ed25519"); + } + + #[test] + fn duplicate_algorithm_keeps_the_last_signature() { + let out = last_signature_per_algo(vec![ + signature("ecdsa", "first"), + signature("ed25519", "b"), + signature("ecdsa", "second"), + ]); + assert_eq!(out.len(), 2); + assert_eq!(out[0].signing_algo, "ecdsa"); + assert_eq!(out[0].text, "second"); + assert_eq!(out[1].signing_algo, "ed25519"); + } + + #[test] + fn empty_input_stays_empty() { + assert!(last_signature_per_algo(Vec::new()).is_empty()); + } +} diff --git a/crates/services/src/attestation/chat_signatures.rs b/crates/services/src/attestation/chat_signatures.rs index 971f12961..0cfa7e28a 100644 --- a/crates/services/src/attestation/chat_signatures.rs +++ b/crates/services/src/attestation/chat_signatures.rs @@ -75,8 +75,8 @@ impl AttestationService { let environment = get_environment(); let env_tag = format!("{TAG_ENVIRONMENT}:{environment}"); - let result: Result<(), AttestationError> = async { - let mut signatures = Vec::with_capacity(2); + let mut signatures = Vec::with_capacity(2); + let fetched: Result<(), AttestationError> = async { for algo in ["ecdsa", "ed25519"] { let provider_signature = provider .get_signature(chat_id, Some(algo.to_string())) @@ -113,8 +113,19 @@ impl AttestationService { }); } - // Both algorithms land in one statement: this runs before the - // client sees `[DONE]`, so every saved round trip is user-visible. + Ok(()) + } + .await; + + // Store whatever was fetched even if a later algorithm failed: the + // one-at-a-time implementation persisted each signature as it came, + // so a backend that serves ecdsa but fails ed25519 still leaves the + // chat verifiable by ecdsa. Both algorithms land in one statement: + // this runs before the client sees `[DONE]`, so every saved round + // trip is user-visible. + let stored: Result<(), AttestationError> = if signatures.is_empty() { + Ok(()) + } else { self.repository .add_chat_signatures(chat_id, signatures) .await @@ -136,10 +147,9 @@ impl AttestationService { &[&env_tag], ); AttestationError::RepositoryError(e.to_string()) - })?; - Ok(()) - } - .await; + }) + }; + let result = fetched.and(stored); provider.unpin_chat_connection(chat_id); result?; From 63618a601a1ccbb6a9f6c55c93dd61d94fe6efe4 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:21:35 +0000 Subject: [PATCH 04/10] review: evict stale statements pool-wide and retry on the same connection - A stale statement is evicted from every pooled connection's cache (deadpool StatementCaches::clear via the object's pool handle), so a retry on a different connection cannot trip over another stale copy. - Outside a transaction the statement is re-prepared and re-run once on the same connection, so the recovery does not consume retry_db! attempts. Inside a transaction the block is restarted by retry_db! after a pool-wide eviction (record_usage). - Detection keys on the PostgreSQL routine (RevalidateCachedQuery) with the English message as a fallback, so a non-English lc_messages does not disable the recovery; SQLSTATE 26000 (missing prepared statement) is handled the same way as a defensive measure. - The e2e test warms and probes concurrently so every connection of the test pool holds the stale plan; lint fix for the column-name format. --- .../tests/e2e_all/prepared_statement_cache.rs | 23 +-- .../src/repositories/organization_usage.rs | 18 ++- .../src/repositories/statement_cache.rs | 140 +++++++++++++----- crates/database/src/repositories/utils.rs | 13 +- 4 files changed, 135 insertions(+), 59 deletions(-) diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs index 254d8638a..1b892827f 100644 --- a/crates/api/tests/e2e_all/prepared_statement_cache.rs +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -6,6 +6,7 @@ //! recover from that transparently, not surface a 500. use crate::common::*; +use futures::future::join_all; #[tokio::test] async fn cached_statements_recover_from_a_column_added_under_them() { @@ -23,10 +24,10 @@ async fn cached_statements_recover_from_a_column_added_under_them() { }; // Warm the cache: API-key auth resolves the workspace and organization - // with a `SELECT w.*` on every request, through whichever pooled - // connection serves it. Several requests seed several connections. - for _ in 0..8 { - let response = probe().await; + // with a `SELECT w.*` on every request. Probes run concurrently so the + // statement is cached on every connection of the test pool, not just + // one; that is the shape of a warm production replica. + for response in join_all((0..8).map(|_| probe())).await { assert_eq!(response.status_code(), 200, "{}", response.text()); } @@ -34,7 +35,7 @@ async fn cached_statements_recover_from_a_column_added_under_them() { // this process still holds statements prepared against the old shape. let column = format!( "rolling_deploy_{}", - uuid::Uuid::new_v4().simple().to_string()[..8].to_string() + &uuid::Uuid::new_v4().simple().to_string()[..8] ); { let client = database @@ -56,11 +57,13 @@ async fn cached_statements_recover_from_a_column_added_under_them() { // Collect instead of asserting so the column is dropped even on failure; // a leftover column would fail the database-encryption classification // scans that share this database. - let mut outcomes = Vec::with_capacity(8); - for _ in 0..8 { - let response = probe().await; - outcomes.push((response.status_code(), response.text())); - } + // Concurrent again: every warm connection holds a stale plan, and each + // request must recover on whichever connection it lands on. + let outcomes: Vec<_> = join_all((0..8).map(|_| probe())) + .await + .into_iter() + .map(|response| (response.status_code(), response.text())) + .collect(); let client = database .pool() diff --git a/crates/database/src/repositories/organization_usage.rs b/crates/database/src/repositories/organization_usage.rs index a14e8bc54..1e0001297 100644 --- a/crates/database/src/repositories/organization_usage.rs +++ b/crates/database/src/repositories/organization_usage.rs @@ -3,7 +3,9 @@ use crate::models::{ ServedProviderType, StopReason, }; use crate::pool::DbPool; -use crate::repositories::statement_cache::CachedStatements; +use crate::repositories::statement_cache::{ + invalidate_pool_statement_caches, is_stale_statement, CachedStatements, +}; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -78,6 +80,16 @@ impl OrganizationUsageRepository { .context("Failed to get database connection") .map_err(RepositoryError::PoolError)?; + // A stale statement inside the transaction aborts it; evict the + // statement from every pooled connection so the retry does not + // trip over another stale copy, then let `retry_db!` restart. + let pool = deadpool_postgres::Client::pool(&client); + let map_tx_error = |err: tokio_postgres::Error| { + if is_stale_statement(&err) { + invalidate_pool_statement_caches(pool.as_ref()); + } + map_db_error(err) + }; let transaction = client.transaction().await.map_err(map_db_error)?; let id = Uuid::new_v4(); @@ -141,7 +153,7 @@ impl OrganizationUsageRepository { ], ) .await - .map_err(map_db_error)?; + .map_err(map_tx_error)?; let (row, was_inserted) = match maybe_row { Some(row) => { @@ -173,7 +185,7 @@ impl OrganizationUsageRepository { ], ) .await - .map_err(map_db_error)?; + .map_err(map_tx_error)?; transaction.commit().await.map_err(map_db_error)?; (row, true) diff --git a/crates/database/src/repositories/statement_cache.rs b/crates/database/src/repositories/statement_cache.rs index 750ddc1cf..4612e63d5 100644 --- a/crates/database/src/repositories/statement_cache.rs +++ b/crates/database/src/repositories/statement_cache.rs @@ -6,19 +6,29 @@ //! one round trip. On a replica far from the database this is the difference //! between ~60 ms and ~120 ms per query. //! -//! The one thing a cached statement cannot survive is a change to its result -//! shape: after a migration adds a column to a table read with `SELECT *`, -//! Postgres rejects the old plan with SQLSTATE 0A000 "cached plan must not -//! change result type". That happens on replicas still running the previous -//! release during a rolling deploy. The helpers below drop the stale statement -//! from the connection's cache when they see that error, and -//! [`map_db_error`](super::utils::map_db_error) classifies it as retryable so -//! the surrounding `retry_db!` block re-prepares and re-runs the statement. +//! A cached statement can outlive the server-side object it points at: +//! +//! - after a migration adds a column to a table read with `SELECT *`, Postgres +//! refuses the old plan with SQLSTATE 0A000 "cached plan must not change +//! result type" (a replica still on the previous release during a rolling +//! deploy hits this on every warm connection); +//! - a statement can be missing on the server altogether, SQLSTATE 26000, +//! if the connection was reset underneath the cache. +//! +//! Both are handled the same way: the stale entry is evicted from every +//! connection's cache in the pool (a schema change invalidates all of them +//! at once), the statement is re-prepared and re-run on the same connection +//! when no transaction is involved, and [`map_db_error`](super::utils::map_db_error) +//! classifies the error as retryable so `retry_db!` re-runs transactional +//! blocks with a fresh statement. use async_trait::async_trait; -use deadpool_postgres::{Client, Transaction}; +use deadpool_postgres::{Client, Pool, Transaction}; use tokio_postgres::{error::SqlState, types::ToSql, Error, Row}; +/// PostgreSQL routine that raises the stale-plan error; stable across server +/// locales, unlike the message text. +const STALE_PLAN_ROUTINE: &str = "RevalidateCachedQuery"; const STALE_PLAN_MESSAGE: &str = "cached plan must not change result type"; /// True when `err` is Postgres refusing a prepared statement whose result @@ -26,10 +36,33 @@ const STALE_PLAN_MESSAGE: &str = "cached plan must not change result type"; pub fn is_stale_cached_plan(err: &Error) -> bool { err.as_db_error().is_some_and(|db_err| { db_err.code() == &SqlState::FEATURE_NOT_SUPPORTED - && db_err.message().contains(STALE_PLAN_MESSAGE) + && (db_err.routine() == Some(STALE_PLAN_ROUTINE) + || db_err.message().contains(STALE_PLAN_MESSAGE)) }) } +/// True when the server no longer has the prepared statement the cache +/// handed out. +pub fn is_missing_prepared_statement(err: &Error) -> bool { + err.as_db_error() + .is_some_and(|db_err| db_err.code() == &SqlState::INVALID_SQL_STATEMENT_NAME) +} + +/// True for either way a cached statement handle can go bad. +pub fn is_stale_statement(err: &Error) -> bool { + is_stale_cached_plan(err) || is_missing_prepared_statement(err) +} + +/// Evict every cached statement on every connection of `pool`. A schema +/// change invalidates the same statement on all warm connections, so a +/// per-connection eviction would leave the next retry to trip over another +/// stale copy. +pub fn invalidate_pool_statement_caches(pool: Option<&Pool>) { + if let Some(pool) = pool { + pool.manager().statement_caches.clear(); + } +} + /// Run queries through the connection's prepared-statement cache. #[async_trait] pub trait CachedStatements { @@ -55,6 +88,25 @@ pub trait CachedStatements { -> Result; } +/// Runs a cached statement, and on a stale handle evicts it pool-wide and +/// re-runs once on the same connection (outside a transaction only). +macro_rules! run_cached { + ($self:ident, $method:ident, $sql:expr, $params:expr) => {{ + let statement = $self.prepare_cached($sql).await?; + match $self.$method(&statement, $params).await { + Err(err) if is_stale_statement(&err) => { + $self.evict_stale($sql); + if !$self.can_retry_in_place() { + return Err(err); + } + let statement = $self.prepare_cached($sql).await?; + $self.$method(&statement, $params).await + } + result => result, + } + }}; +} + macro_rules! impl_cached_statements { ($ty:ty) => { #[async_trait] @@ -64,10 +116,7 @@ macro_rules! impl_cached_statements { sql: &str, params: &[&(dyn ToSql + Sync)], ) -> Result, Error> { - let statement = self.prepare_cached(sql).await?; - let result = self.query(&statement, params).await; - self.forget_if_stale(sql, &result); - result + run_cached!(self, query, sql, params) } async fn cached_query_opt( @@ -75,10 +124,7 @@ macro_rules! impl_cached_statements { sql: &str, params: &[&(dyn ToSql + Sync)], ) -> Result, Error> { - let statement = self.prepare_cached(sql).await?; - let result = self.query_opt(&statement, params).await; - self.forget_if_stale(sql, &result); - result + run_cached!(self, query_opt, sql, params) } async fn cached_query_one( @@ -86,10 +132,7 @@ macro_rules! impl_cached_statements { sql: &str, params: &[&(dyn ToSql + Sync)], ) -> Result { - let statement = self.prepare_cached(sql).await?; - let result = self.query_one(&statement, params).await; - self.forget_if_stale(sql, &result); - result + run_cached!(self, query_one, sql, params) } async fn cached_execute( @@ -97,30 +140,47 @@ macro_rules! impl_cached_statements { sql: &str, params: &[&(dyn ToSql + Sync)], ) -> Result { - let statement = self.prepare_cached(sql).await?; - let result = self.execute(&statement, params).await; - self.forget_if_stale(sql, &result); - result - } - } - - impl StaleStatementCleanup for $ty { - fn forget_if_stale(&self, sql: &str, result: &Result) { - if let Err(err) = result { - if is_stale_cached_plan(err) { - tracing::warn!( - "Dropping prepared statement whose result shape changed; retrying" - ); - self.statement_cache.remove(sql, &[]); - } - } + run_cached!(self, execute, sql, params) } } }; } trait StaleStatementCleanup { - fn forget_if_stale(&self, sql: &str, result: &Result); + /// Drop stale statements after the server rejected `sql`. + fn evict_stale(&self, sql: &str); + /// Whether the statement can be re-prepared and re-run on this + /// connection right away. + fn can_retry_in_place(&self) -> bool; +} + +impl StaleStatementCleanup for Client { + fn evict_stale(&self, _sql: &str) { + tracing::warn!("Dropping cached prepared statements after a schema change; retrying"); + match Client::pool(self) { + Some(pool) => invalidate_pool_statement_caches(Some(&pool)), + None => self.statement_cache.clear(), + } + } + + fn can_retry_in_place(&self) -> bool { + true + } +} + +impl StaleStatementCleanup for Transaction<'_> { + fn evict_stale(&self, _sql: &str) { + tracing::warn!("Dropping cached prepared statements after a schema change; retrying"); + // The pool is not reachable from a transaction; callers that hold + // the client evict pool-wide through `invalidate_pool_statement_caches`. + self.statement_cache.clear(); + } + + fn can_retry_in_place(&self) -> bool { + // The transaction is aborted after the error; the surrounding + // `retry_db!` restarts it with a fresh statement. + false + } } impl_cached_statements!(Client); diff --git a/crates/database/src/repositories/utils.rs b/crates/database/src/repositories/utils.rs index a3ab42edc..20763849f 100644 --- a/crates/database/src/repositories/utils.rs +++ b/crates/database/src/repositories/utils.rs @@ -41,13 +41,14 @@ pub fn map_db_error(err: tokio_postgres::Error) -> RepositoryError { &SqlState::QUERY_CANCELED => RepositoryError::QueryTimeout, - // A cached prepared statement outlived a schema change (rolling - // deploy). The statement has already been dropped from the - // connection's cache (see `statement_cache`); classify as - // retryable so `retry_db!` re-prepares and re-runs it. - _ if super::statement_cache::is_stale_cached_plan(&err) => { + // A cached prepared statement outlived the server-side object it + // points at (schema change during a rolling deploy, or a reset + // connection). The caches have already been evicted (see + // `statement_cache`); classify as retryable so `retry_db!` + // re-prepares and re-runs it. + _ if super::statement_cache::is_stale_statement(&err) => { RepositoryError::ConnectionFailed( - "prepared statement invalidated by a schema change".to_string(), + "cached prepared statement is no longer valid".to_string(), ) } From d8ad797bef61ca35a61483c10ec069502218bcd9 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:34:00 +0000 Subject: [PATCH 05/10] review: log signature id and error when a gateway signature write fails --- crates/services/src/attestation/gateway_signatures.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/services/src/attestation/gateway_signatures.rs b/crates/services/src/attestation/gateway_signatures.rs index c16479ceb..07c4e18bd 100644 --- a/crates/services/src/attestation/gateway_signatures.rs +++ b/crates/services/src/attestation/gateway_signatures.rs @@ -40,7 +40,12 @@ impl AttestationService { .add_chat_signatures(signature_id, signatures) .await .map_err(|e| { - tracing::error!("Failed to store {} signatures in repository", id_label); + tracing::error!( + signature_kind = id_label, + signature_id = signature_id, + error = %e, + "Failed to store gateway signatures in repository" + ); AttestationError::RepositoryError(e.to_string()) })?; tracing::info!( From 58e7f4ccb2fb047cba5643130c33399900a5ab12 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:34:21 +0000 Subject: [PATCH 06/10] review: make the e2e column cleanup best-effort and assert it last --- .../tests/e2e_all/prepared_statement_cache.rs | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs index 1b892827f..78e7227dc 100644 --- a/crates/api/tests/e2e_all/prepared_statement_cache.rs +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -65,18 +65,21 @@ async fn cached_statements_recover_from_a_column_added_under_them() { .map(|response| (response.status_code(), response.text())) .collect(); - let client = database - .pool() - .get() - .await - .expect("failed to get database connection"); - client - .execute( - &format!("ALTER TABLE workspaces DROP COLUMN IF EXISTS {column}"), - &[], - ) - .await - .expect("failed to drop column"); + // Best-effort cleanup that never panics: a leftover column breaks the + // database-encryption classification scans that share this database, so + // the drop must run even when the pool or the statement misbehaves, and + // the assertions below must still report the real outcome. + let cleanup = async { + let client = database.pool().get().await?; + client + .execute( + &format!("ALTER TABLE workspaces DROP COLUMN IF EXISTS {column}"), + &[], + ) + .await?; + Ok::<(), Box>(()) + }; + let cleanup = cleanup.await; for (status, body) in outcomes { assert_eq!( @@ -84,4 +87,6 @@ async fn cached_statements_recover_from_a_column_added_under_them() { "request after a schema change must recover, got: {body}" ); } + + cleanup.expect("failed to drop the temporary column"); } From af537a8aef1b798bd1336fb8b9d5de37985a23bd Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:52:52 +0000 Subject: [PATCH 07/10] test(database): verify cached billing transaction recovery --- .../tests/e2e_all/prepared_statement_cache.rs | 220 ++++++++++++++++++ 1 file changed, 220 insertions(+) diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs index 78e7227dc..a442a06df 100644 --- a/crates/api/tests/e2e_all/prepared_statement_cache.rs +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -6,7 +6,10 @@ //! recover from that transparently, not surface a 500. use crate::common::*; +use anyhow::{ensure, Context, Result}; +use database::{models::RecordUsageRequest, repositories::OrganizationUsageRepository}; use futures::future::join_all; +use uuid::Uuid; #[tokio::test] async fn cached_statements_recover_from_a_column_added_under_them() { @@ -90,3 +93,220 @@ async fn cached_statements_recover_from_a_column_added_under_them() { cleanup.expect("failed to drop the temporary column"); } + +enum UsageStatementInvalidation { + AddedColumn, + Deallocated, +} + +async fn usage_fixture() -> (database::pool::DbPool, RecordUsageRequest) { + let pool = db_setup::create_test_pool().await; + let client = pool.get().await.expect("fixture connection"); + let suffix = Uuid::new_v4().to_string(); + let user_id = Uuid::parse_str(MOCK_USER_ID).unwrap(); + let row = client + .query_one( + "WITH org AS ( + INSERT INTO organizations (name) VALUES ($1) RETURNING id + ), workspace AS ( + INSERT INTO workspaces (name, organization_id, created_by_user_id) + SELECT $1, org.id, $2 FROM org RETURNING id + ), api_key AS ( + INSERT INTO api_keys (key_hash, key_prefix, name, workspace_id, created_by_user_id) + SELECT lpad(replace(workspace.id::text, '-', ''), 64, '0'), 'sk-test', $1, workspace.id, $2 + FROM workspace RETURNING id + ) + SELECT org.id AS organization_id, workspace.id AS workspace_id, + api_key.id AS api_key_id, models.id AS model_id + FROM org, workspace, api_key, models WHERE models.model_name = $3", + &[&suffix, &user_id, &E2E_QWEN_MODEL_NAME], + ) + .await + .expect("test-owned usage fixture"); + let request = RecordUsageRequest { + organization_id: row.get("organization_id"), + workspace_id: row.get("workspace_id"), + api_key_id: row.get("api_key_id"), + model_id: row.get("model_id"), + model_name: E2E_QWEN_MODEL_NAME.to_string(), + input_tokens: 12, + output_tokens: 8, + input_cost: 12_000_000, + output_cost: 16_000_000, + total_cost: 28_000_000, + inference_type: services::usage::InferenceType::ChatCompletion + .as_str() + .to_string(), + ttft_ms: None, + avg_itl_ms: None, + inference_id: Some(Uuid::new_v4()), + provider_request_id: None, + stop_reason: None, + response_id: None, + image_count: None, + cache_read_tokens: 0, + cache_write_tokens: 0, + billing_details: None, + service_tier: None, + context_band: None, + served_provider_tier: None, + served_provider_type: None, + served_via_fallback: false, + }; + drop(client); + (pool, request) +} + +async fn assert_usage_recovers_without_double_charging(invalidation: UsageStatementInvalidation) { + let (pool, request) = usage_fixture().await; + let repository = OrganizationUsageRepository::new(pool.clone()); + let pool_size = pool.status().expect("active pool").max_size; + // More warm connections than retry_db!'s three attempts proves that + // recovery cannot rely on eventually acquiring a cold connection. + assert!(pool_size > 3); + let mut connections = Vec::new(); + for _ in 0..pool_size { + connections.push(pool.get().await.expect("reserve pool connection")); + } + for _ in 0..pool_size { + // Only this connection is available to the real repository. Retain + // it after warming, then release the next unvisited connection. + drop(connections.remove(0)); + let mut warm_request = request.clone(); + warm_request.inference_id = Some(Uuid::new_v4()); + assert!( + repository + .record_usage(warm_request) + .await + .unwrap() + .was_inserted + ); + let client = pool.get().await.expect("reacquire warmed connection"); + let cached_insert: bool = client + .query_one( + "SELECT EXISTS ( + SELECT 1 FROM pg_prepared_statements + WHERE statement LIKE '%INSERT INTO organization_usage_log%' + AND statement NOT LIKE '%pg_prepared_statements%' + )", + &[], + ) + .await + .expect("inspect warmed session") + .get(0); + assert!( + cached_insert, + "every session must hold the real usage INSERT" + ); + connections.push(client); + } + + let column = format!("usage_cache_{}", Uuid::new_v4().simple()); + match invalidation { + UsageStatementInvalidation::AddedColumn => { + connections[0] + .batch_execute(&format!( + "ALTER TABLE organization_usage_log ADD COLUMN {column} TEXT" + )) + .await + .expect("invalidate the cached INSERT RETURNING result type"); + } + UsageStatementInvalidation::Deallocated => { + for client in &connections { + // Invalidate server handles while retaining deadpool's cache. + client.batch_execute("DEALLOCATE ALL").await.unwrap(); + } + } + } + drop(connections); + + // Return errors instead of panicking until after DDL cleanup, so a failed + // regression does not leave an unclassified column for encryption tests. + let outcome = async { + let recovered = repository + .record_usage(request.clone()) + .await + .context("record usage after invalidating every warm connection")?; + ensure!(recovered.was_inserted, "recovered usage must be inserted"); + let mut reserved = Vec::new(); + for _ in 1..pool_size { + reserved.push(pool.get().await?); + } + let mut duplicate_request = request.clone(); + duplicate_request.total_cost *= 2; + let duplicate = repository.record_usage(duplicate_request).await?; + ensure!(duplicate.id == recovered.id && !duplicate.was_inserted); + ensure!(duplicate.total_cost == request.total_cost); + + // Only one connection is free: the new write must reuse the session + // (and cached INSERT) whose duplicate transaction just rolled back. + let mut after_rollback = request.clone(); + after_rollback.inference_id = Some(Uuid::new_v4()); + ensure!(repository.record_usage(after_rollback).await?.was_inserted); + + let expected_requests = pool_size as i64 + 2; + let expected_spend = expected_requests * request.total_cost; + let expected_tokens = + expected_requests * i64::from(request.input_tokens + request.output_tokens); + let client = pool.get().await?; + let usage = client + .query_one( + "SELECT COUNT(*)::BIGINT, SUM(total_cost)::BIGINT, SUM(total_tokens)::BIGINT, + COUNT(*) FILTER (WHERE inference_id = $2)::BIGINT + FROM organization_usage_log WHERE organization_id = $1", + &[&request.organization_id, &request.inference_id], + ) + .await?; + ensure!( + usage.get::<_, i64>(0) == expected_requests, + "usage row count" + ); + ensure!(usage.get::<_, i64>(1) == expected_spend, "usage cost"); + ensure!(usage.get::<_, i64>(2) == expected_tokens, "usage tokens"); + ensure!( + usage.get::<_, i64>(3) == 1, + "recovered inference appears once" + ); + drop(client); + let balance = repository + .get_balance(request.organization_id) + .await? + .context("organization balance must exist")?; + ensure!( + balance.total_requests == expected_requests, + "balance requests" + ); + ensure!( + balance.total_spent == expected_spend, + "balance charged exactly once" + ); + ensure!(balance.total_tokens == expected_tokens, "balance tokens"); + Ok::<(), anyhow::Error>(()) + } + .await; + + let cleanup: Result<()> = async { + if matches!(invalidation, UsageStatementInvalidation::AddedColumn) { + pool.get() + .await? + .batch_execute(&format!( + "ALTER TABLE organization_usage_log DROP COLUMN IF EXISTS {column}" + )) + .await?; + } + Ok(()) + } + .await; + outcome.expect("usage transaction must recover without duplicate billing"); + cleanup.expect("remove temporary usage column"); +} + +#[tokio::test] +async fn usage_transaction_recovers_from_changed_result_type_without_double_charging() { + assert_usage_recovers_without_double_charging(UsageStatementInvalidation::AddedColumn).await; +} + +#[tokio::test] +async fn usage_transaction_recovers_from_missing_statements_without_double_charging() { + assert_usage_recovers_without_double_charging(UsageStatementInvalidation::Deallocated).await; +} From e60a6e62fe7668efde6743036a9a757e5354a666 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:56:48 +0000 Subject: [PATCH 08/10] test(database): exercise recovery after a partial billing transaction --- .../tests/e2e_all/prepared_statement_cache.rs | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs index a442a06df..df58e7d51 100644 --- a/crates/api/tests/e2e_all/prepared_statement_cache.rs +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -96,7 +96,7 @@ async fn cached_statements_recover_from_a_column_added_under_them() { enum UsageStatementInvalidation { AddedColumn, - Deallocated, + DeallocatedBalance, } async fn usage_fixture() -> (database::pool::DbPool, RecordUsageRequest) { @@ -211,10 +211,24 @@ async fn assert_usage_recovers_without_double_charging(invalidation: UsageStatem .await .expect("invalidate the cached INSERT RETURNING result type"); } - UsageStatementInvalidation::Deallocated => { + UsageStatementInvalidation::DeallocatedBalance => { for client in &connections { - // Invalidate server handles while retaining deadpool's cache. - client.batch_execute("DEALLOCATE ALL").await.unwrap(); + // Leave the usage INSERT valid, so it succeeds before the + // balance UPSERT fails. Recovery must roll back that usage + // row and restart the entire transaction to charge once. + let balance_statement: String = client + .query_one( + "SELECT name FROM pg_prepared_statements WHERE statement LIKE $1", + &[&"%INSERT INTO organization_balance%"], + ) + .await + .expect("every warmed session must have the balance UPSERT") + .get(0); + let quoted_name = balance_statement.replace('"', "\"\""); + client + .batch_execute(&format!("DEALLOCATE \"{quoted_name}\"")) + .await + .expect("invalidate only the server's balance statement"); } } } @@ -238,8 +252,8 @@ async fn assert_usage_recovers_without_double_charging(invalidation: UsageStatem ensure!(duplicate.id == recovered.id && !duplicate.was_inserted); ensure!(duplicate.total_cost == request.total_cost); - // Only one connection is free: the new write must reuse the session - // (and cached INSERT) whose duplicate transaction just rolled back. + // Only one connection is free: a new write must succeed on the + // session whose duplicate transaction just rolled back. let mut after_rollback = request.clone(); after_rollback.inference_id = Some(Uuid::new_v4()); ensure!(repository.record_usage(after_rollback).await?.was_inserted); @@ -307,6 +321,7 @@ async fn usage_transaction_recovers_from_changed_result_type_without_double_char } #[tokio::test] -async fn usage_transaction_recovers_from_missing_statements_without_double_charging() { - assert_usage_recovers_without_double_charging(UsageStatementInvalidation::Deallocated).await; +async fn usage_transaction_recovers_from_missing_balance_statement_without_double_charging() { + assert_usage_recovers_without_double_charging(UsageStatementInvalidation::DeallocatedBalance) + .await; } From 37b67d642918c5bdf277af9705dd9e3cb769cfa6 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:57:13 +0000 Subject: [PATCH 09/10] fix(attestation): preserve partial signatures when stream fetches time out --- .../chat_signature_lifecycle_tests.rs | 25 +- .../src/attestation/chat_signatures.rs | 77 ++-- crates/services/src/attestation/mod.rs | 3 + crates/services/src/attestation/ports.rs | 10 + .../attestation/provider_signature_tests.rs | 413 ++++++++++++++++++ .../services/src/attestation/service_trait.rs | 16 +- crates/services/src/completions/mod.rs | 8 +- 7 files changed, 514 insertions(+), 38 deletions(-) create mode 100644 crates/services/src/attestation/provider_signature_tests.rs diff --git a/crates/services/src/attestation/chat_signature_lifecycle_tests.rs b/crates/services/src/attestation/chat_signature_lifecycle_tests.rs index 27bce8d96..fd026083a 100644 --- a/crates/services/src/attestation/chat_signature_lifecycle_tests.rs +++ b/crates/services/src/attestation/chat_signature_lifecycle_tests.rs @@ -36,18 +36,35 @@ use crate::{ /// Repository that records stored signatures and succeeds. #[derive(Clone, Default)] -struct RecordingRepository { +pub(super) struct RecordingRepository { stored: Arc>>, + batch_sizes: Arc>>, } impl RecordingRepository { - fn stored(&self) -> Vec<(String, ChatSignature)> { + pub(super) fn stored(&self) -> Vec<(String, ChatSignature)> { self.stored.lock().map(|s| s.clone()).unwrap_or_default() } + + pub(super) fn batch_sizes(&self) -> Vec { + self.batch_sizes.lock().unwrap().clone() + } } #[async_trait] impl AttestationRepository for RecordingRepository { + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + self.batch_sizes.lock().unwrap().push(signatures.len()); + for signature in signatures { + self.add_chat_signature(chat_id, signature).await?; + } + Ok(()) + } + async fn add_chat_signature( &self, chat_id: &str, @@ -72,7 +89,7 @@ impl AttestationRepository for RecordingRepository { } /// Repository whose store always fails. -struct FailingRepository; +pub(super) struct FailingRepository; #[async_trait] impl AttestationRepository for FailingRepository { @@ -251,7 +268,7 @@ async fn pool_with_pinned_chat(chat_id: &str) -> (Arc, Ar (pool, provider) } -fn lifecycle_service( +pub(super) fn lifecycle_service( repository: Arc, pool: Arc, ) -> AttestationService { diff --git a/crates/services/src/attestation/chat_signatures.rs b/crates/services/src/attestation/chat_signatures.rs index 0cfa7e28a..41ed3ee05 100644 --- a/crates/services/src/attestation/chat_signatures.rs +++ b/crates/services/src/attestation/chat_signatures.rs @@ -7,11 +7,17 @@ use super::{ }; use crate::{metrics::consts::*, usage::StopReason}; -/// Upper bound on the gateway signature store at end-of-stream. The client is +/// Upper bound on signature finalization at end-of-stream. The client is /// still waiting for the held-back `[DONE]` while this runs, so it must be /// bounded; on timeout the routing pin is still released and the stream ends /// without a stored signature (logged by the caller). -pub(in crate::attestation) const STREAM_SIGNATURE_STORE_TIMEOUT: Duration = Duration::from_secs(5); +pub(crate) const STREAM_SIGNATURE_STORE_TIMEOUT: Duration = Duration::from_secs(5); + +/// Reserve the final second of the streaming deadline for persisting fetched +/// signatures. An outer timeout would otherwise drop an already-fetched ECDSA +/// signature if the ED25519 fetch stalls before the batch can be written. +pub(in crate::attestation) const PROVIDER_SIGNATURE_FETCH_TIMEOUT: Duration = + STREAM_SIGNATURE_STORE_TIMEOUT.saturating_sub(Duration::from_secs(1)); impl AttestationService { pub(in crate::attestation) async fn get_chat_signature_impl( @@ -54,6 +60,7 @@ impl AttestationService { pub(in crate::attestation) async fn store_chat_signature_from_provider_impl( &self, chat_id: &str, + fetch_deadline: Option, ) -> Result<(), AttestationError> { let start_time = std::time::Instant::now(); let provider = self @@ -78,32 +85,43 @@ impl AttestationService { let mut signatures = Vec::with_capacity(2); let fetched: Result<(), AttestationError> = async { for algo in ["ecdsa", "ed25519"] { - let provider_signature = provider - .get_signature(chat_id, Some(algo.to_string())) - .await - .map_err(|e| { - // The error string embeds the backend URL on connection - // failures — without it (and chat_id) these events are - // impossible to attribute to a model/backend. - tracing::error!( - %chat_id, - error = %e, - "Failed to get chat signature from provider for algorithm: {}", - algo - ); - let duration = start_time.elapsed(); - self.metrics_service.record_count( - METRIC_VERIFICATION_FAILURE, - 1, - &[&format!("{TAG_REASON}:{REASON_INFERENCE_ERROR}"), &env_tag], - ); - self.metrics_service.record_latency( - METRIC_VERIFICATION_DURATION, - duration, - &[&env_tag], - ); - AttestationError::ProviderError(e.to_string()) - })?; + let fetch = async { + provider + .get_signature(chat_id, Some(algo.to_string())) + .await + .map_err(|e| AttestationError::ProviderError(e.to_string())) + }; + let provider_signature = match fetch_deadline { + Some(deadline) => match tokio::time::timeout_at(deadline, fetch).await { + Ok(result) => result, + Err(_) => Err(AttestationError::ProviderError(format!( + "Timed out fetching chat signature for algorithm: {algo}" + ))), + }, + None => fetch.await, + } + .inspect_err(|e| { + // The error string embeds the backend URL on connection + // failures — without it (and chat_id) these events are + // impossible to attribute to a model/backend. + tracing::error!( + %chat_id, + error = %e, + "Failed to get chat signature from provider for algorithm: {}", + algo + ); + let duration = start_time.elapsed(); + self.metrics_service.record_count( + METRIC_VERIFICATION_FAILURE, + 1, + &[&format!("{TAG_REASON}:{REASON_INFERENCE_ERROR}"), &env_tag], + ); + self.metrics_service.record_latency( + METRIC_VERIFICATION_DURATION, + duration, + &[&env_tag], + ); + })?; signatures.push(ChatSignature { text: provider_signature.text, signature: provider_signature.signature, @@ -117,7 +135,8 @@ impl AttestationService { } .await; - // Store whatever was fetched even if a later algorithm failed: the + // Store whatever was fetched even if a later algorithm failed or hit + // the fetch deadline, leaving time before the outer stream timeout. The // one-at-a-time implementation persisted each signature as it came, // so a backend that serves ecdsa but fails ed25519 still leaves the // chat verifiable by ecdsa. Both algorithms land in one statement: diff --git a/crates/services/src/attestation/mod.rs b/crates/services/src/attestation/mod.rs index e54dfea42..0a2ef0812 100644 --- a/crates/services/src/attestation/mod.rs +++ b/crates/services/src/attestation/mod.rs @@ -11,6 +11,8 @@ mod lifecycle; pub mod measurement; pub mod models; pub mod ports; +#[cfg(test)] +mod provider_signature_tests; mod report; pub mod report_data; mod service_trait; @@ -22,6 +24,7 @@ use config::ItaAttestationConfig; use ed25519_dalek::{SigningKey, VerifyingKey}; use k256::ecdsa::{SigningKey as EcdsaSigningKey, VerifyingKey as EcdsaVerifyingKey}; +pub(crate) use chat_signatures::STREAM_SIGNATURE_STORE_TIMEOUT; pub use environment::{ compute_spki_hash, load_tls_cert_fingerprint, load_vpc_info, load_vpc_shared_secret, }; diff --git a/crates/services/src/attestation/ports.rs b/crates/services/src/attestation/ports.rs index e1a1599e3..85e8e46b9 100644 --- a/crates/services/src/attestation/ports.rs +++ b/crates/services/src/attestation/ports.rs @@ -23,6 +23,16 @@ pub trait AttestationServiceTrait: Send + Sync { chat_id: &str, ) -> Result<(), AttestationError>; + /// Fetch and store provider signatures before a stream sends `[DONE]`. + /// Implementations may reserve part of the caller's finalization deadline + /// to persist partial fetches; background callers use the method above. + async fn store_stream_chat_signature_from_provider( + &self, + chat_id: &str, + ) -> Result<(), AttestationError> { + self.store_chat_signature_from_provider(chat_id).await + } + /// Store a chat signature directly over gateway-emitted bytes. /// Creates a signature with text format "request_hash:response_hash" /// and stores both ECDSA and ED25519 signatures. diff --git a/crates/services/src/attestation/provider_signature_tests.rs b/crates/services/src/attestation/provider_signature_tests.rs new file mode 100644 index 000000000..a087cbd65 --- /dev/null +++ b/crates/services/src/attestation/provider_signature_tests.rs @@ -0,0 +1,413 @@ +//! Provider signature batching must retain successful fetches when the other +//! algorithm fails or times out before the streaming finalization deadline. +use std::{sync::Arc, time::Duration}; + +use async_trait::async_trait; +use config::ExternalProvidersConfig; +use inference_providers::{self as ip, mock::MockProvider, InferenceProvider}; + +use super::{ + chat_signature_lifecycle_tests::{lifecycle_service, FailingRepository, RecordingRepository}, + ports::{AttestationRepository, AttestationServiceTrait}, + AttestationError, AttestationService, ChatSignature, SignatureKind, + STREAM_SIGNATURE_STORE_TIMEOUT, +}; +use crate::{ + inference_provider_pool::InferenceProviderPool, + metrics::{ + capturing::{CapturingMetricsService, MetricValue}, + consts::*, + }, +}; + +#[derive(Clone, Copy)] +enum FetchBehavior { + Success, + FailSecond, + StallSecond, + StallFirst, + SlowSecond, +} + +struct ControlledProvider { + inner: MockProvider, + behavior: FetchBehavior, +} + +#[async_trait] +impl InferenceProvider for ControlledProvider { + async fn get_signature( + &self, + chat_id: &str, + signing_algo: Option, + ) -> Result { + match (signing_algo.as_deref(), self.behavior) { + (Some("ecdsa"), FetchBehavior::StallFirst) + | (Some("ed25519"), FetchBehavior::StallSecond) => std::future::pending().await, + (Some("ed25519"), FetchBehavior::FailSecond) => Err( + ip::CompletionError::CompletionError("synthetic fetch failure".to_string()), + ), + (Some("ed25519"), FetchBehavior::SlowSecond) => { + tokio::time::sleep(Duration::from_secs(6)).await; + self.inner.get_signature(chat_id, signing_algo).await + } + _ => { + // Spending part of the budget on ECDSA catches a timeout that + // is accidentally restarted for each algorithm. + if matches!(self.behavior, FetchBehavior::StallSecond) { + tokio::time::sleep(Duration::from_secs(2)).await; + } + self.inner.get_signature(chat_id, signing_algo).await + } + } + } + + fn unpin_chat_connection(&self, chat_id: &str) { + self.inner.unpin_chat_connection(chat_id); + } + async fn models(&self) -> Result { + self.inner.models().await + } + + async fn chat_completion_stream( + &self, + params: ip::ChatCompletionParams, + request_hash: String, + ) -> Result { + self.inner + .chat_completion_stream(params, request_hash) + .await + } + + async fn chat_completion( + &self, + params: ip::ChatCompletionParams, + request_hash: String, + ) -> Result { + self.inner.chat_completion(params, request_hash).await + } + + async fn text_completion_stream( + &self, + params: ip::CompletionParams, + ) -> Result { + self.inner.text_completion_stream(params).await + } + + async fn image_generation( + &self, + params: ip::ImageGenerationParams, + request_hash: String, + ) -> Result { + self.inner.image_generation(params, request_hash).await + } + + async fn image_edit( + &self, + params: Arc, + request_hash: String, + ) -> Result { + self.inner.image_edit(params, request_hash).await + } + + async fn score( + &self, + params: ip::ScoreParams, + request_hash: String, + ) -> Result { + self.inner.score(params, request_hash).await + } + + async fn rerank( + &self, + params: ip::RerankParams, + ) -> Result { + self.inner.rerank(params).await + } + + async fn embeddings_raw( + &self, + body: bytes::Bytes, + extra: std::collections::HashMap, + ) -> Result { + self.inner.embeddings_raw(body, extra).await + } + + async fn privacy_classify_raw( + &self, + body: bytes::Bytes, + extra: std::collections::HashMap, + ) -> Result { + self.inner.privacy_classify_raw(body, extra).await + } + + async fn get_attestation_report( + &self, + model: String, + signing_algo: Option, + nonce: Option, + signing_address: Option, + include_tls_fingerprint: bool, + ) -> Result, ip::models::AttestationError> { + self.inner + .get_attestation_report( + model, + signing_algo, + nonce, + signing_address, + include_tls_fingerprint, + ) + .await + } + + async fn audio_transcription( + &self, + params: ip::AudioTranscriptionParams, + request_hash: String, + ) -> Result { + self.inner.audio_transcription(params, request_hash).await + } +} + +async fn provider_service( + chat_id: &str, + behavior: FetchBehavior, + repository: Arc, +) -> ( + AttestationService, + Arc, + Arc, +) { + let pool = Arc::new(InferenceProviderPool::new( + None, + ExternalProvidersConfig::default(), + )); + let provider = Arc::new(ControlledProvider { + inner: MockProvider::new(), + behavior, + }); + pool.store_chat_id_mapping(chat_id.to_string(), provider.clone()) + .await; + let mut service = lifecycle_service(repository, pool); + let metrics = Arc::new(CapturingMetricsService::new()); + service.metrics_service = metrics.clone(); + (service, provider, metrics) +} + +fn assert_failure_metric(metrics: &CapturingMetricsService, reason: &str) { + let failures: Vec<_> = metrics + .get_metrics() + .into_iter() + .filter(|metric| metric.name == METRIC_VERIFICATION_FAILURE) + .collect(); + assert_eq!(failures.len(), 1); + assert!(matches!(failures[0].value, MetricValue::Count(1))); + assert!(failures[0].tags.contains(&format!("{TAG_REASON}:{reason}"))); + assert!(!metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_SUCCESS)); +} + +#[tokio::test] +async fn provider_success_stores_both_signatures_in_one_batch_and_unpins() { + let chat_id = "chatcmpl-provider-batch-success"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::Success, + Arc::new(repository.clone()), + ) + .await; + service + .store_stream_chat_signature_from_provider(chat_id) + .await + .unwrap(); + assert_eq!(repository.batch_sizes(), vec![2]); + let stored = repository.stored(); + assert_eq!( + stored + .iter() + .map(|(_, sig)| sig.signing_algo.as_str()) + .collect::>(), + vec!["ecdsa", "ed25519"] + ); + assert!(stored + .iter() + .all(|(id, sig)| id == chat_id && sig.signature_kind == Some(SignatureKind::ProviderTee))); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert!(metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_SUCCESS + && matches!(metric.value, MetricValue::Count(1)))); + assert!(!metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_FAILURE)); +} + +#[tokio::test] +async fn provider_second_fetch_error_preserves_first_signature_and_unpins() { + let chat_id = "chatcmpl-provider-partial-error"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::FailSecond, + Arc::new(repository.clone()), + ) + .await; + let result = service + .store_stream_chat_signature_from_provider(chat_id) + .await; + assert!(matches!(result, Err(AttestationError::ProviderError(_)))); + assert_eq!(repository.batch_sizes(), vec![1]); + assert_eq!(repository.stored()[0].1.signing_algo, "ecdsa"); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_INFERENCE_ERROR); +} + +/// Make the flush consume part of its reserved budget, so a passing test +/// requires the write to finish before the actual five-second outer timeout. +struct DelayedRepository(RecordingRepository); + +#[async_trait] +impl AttestationRepository for DelayedRepository { + async fn add_chat_signature( + &self, + chat_id: &str, + signature: ChatSignature, + ) -> Result<(), AttestationError> { + self.0.add_chat_signature(chat_id, signature).await + } + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + tokio::time::sleep(Duration::from_millis(500)).await; + self.0.add_chat_signatures(chat_id, signatures).await + } + async fn get_chat_signature( + &self, + chat_id: &str, + signing_algo: &str, + ) -> Result { + self.0.get_chat_signature(chat_id, signing_algo).await + } +} + +#[tokio::test(start_paused = true)] +async fn provider_stalled_second_fetch_flushes_first_before_stream_timeout() { + let chat_id = "chatcmpl-provider-partial-timeout"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::StallSecond, + Arc::new(DelayedRepository(repository.clone())), + ) + .await; + let started = tokio::time::Instant::now(); + // The same hard deadline used by InterceptStream must not cancel the + // partial write, even after ECDSA has used two seconds of the fetch budget. + let result = tokio::time::timeout( + STREAM_SIGNATURE_STORE_TIMEOUT, + service.store_stream_chat_signature_from_provider(chat_id), + ) + .await + .expect("partial signature must be flushed before stream finalization cancels it"); + assert!( + matches!(result, Err(AttestationError::ProviderError(message)) if message.contains("Timed out") && message.contains("ed25519")) + ); + assert_eq!(repository.batch_sizes(), vec![1]); + assert_eq!(repository.stored()[0].1.signing_algo, "ecdsa"); + assert!(started.elapsed() >= Duration::from_millis(4500)); + assert!(started.elapsed() < STREAM_SIGNATURE_STORE_TIMEOUT); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_INFERENCE_ERROR); +} + +#[tokio::test(start_paused = true)] +async fn provider_stalled_first_fetch_unpins_without_an_empty_write() { + let chat_id = "chatcmpl-provider-first-timeout"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::StallFirst, + Arc::new(repository.clone()), + ) + .await; + let result = tokio::time::timeout( + STREAM_SIGNATURE_STORE_TIMEOUT, + service.store_stream_chat_signature_from_provider(chat_id), + ) + .await + .unwrap(); + assert!( + matches!(result, Err(AttestationError::ProviderError(message)) if message.contains("Timed out")) + ); + assert!(repository.batch_sizes().is_empty()); + assert!(repository.stored().is_empty()); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_INFERENCE_ERROR); +} + +#[tokio::test] +async fn provider_repository_error_is_metriced_and_unpins() { + let chat_id = "chatcmpl-provider-repository-error"; + let (service, provider, metrics) = + provider_service(chat_id, FetchBehavior::Success, Arc::new(FailingRepository)).await; + let result = service + .store_stream_chat_signature_from_provider(chat_id) + .await; + assert!(matches!(result, Err(AttestationError::RepositoryError(_)))); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_REPOSITORY_ERROR); +} + +#[tokio::test(start_paused = true)] +async fn background_provider_fetch_keeps_its_existing_timeout() { + let chat_id = "chatcmpl-provider-background-slow"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::SlowSecond, + Arc::new(repository.clone()), + ) + .await; + let started = tokio::time::Instant::now(); + service + .store_chat_signature_from_provider(chat_id) + .await + .unwrap(); + assert!(started.elapsed() >= Duration::from_secs(6)); + assert_eq!(repository.batch_sizes(), vec![2]); + assert_eq!(repository.stored().len(), 2); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert!(metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_SUCCESS)); + assert!(!metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_FAILURE)); +} diff --git a/crates/services/src/attestation/service_trait.rs b/crates/services/src/attestation/service_trait.rs index e85855d7a..a0dedd58a 100644 --- a/crates/services/src/attestation/service_trait.rs +++ b/crates/services/src/attestation/service_trait.rs @@ -1,6 +1,7 @@ use async_trait::async_trait; use super::{ + chat_signatures::PROVIDER_SIGNATURE_FETCH_TIMEOUT, ita::{ItaTokenQuery, ItaTokenResponse}, models::AttestationReport, ports, AttestationError, AttestationService, SignatureLookupResult, @@ -21,7 +22,20 @@ impl ports::AttestationServiceTrait for AttestationService { &self, chat_id: &str, ) -> Result<(), AttestationError> { - self.store_chat_signature_from_provider_impl(chat_id).await + self.store_chat_signature_from_provider_impl(chat_id, None) + .await + } + + async fn store_stream_chat_signature_from_provider( + &self, + chat_id: &str, + ) -> Result<(), AttestationError> { + // Start the budget before provider lookup and share it across both + // algorithms. Non-streaming/background callers retain their existing + // provider timeouts, without this streaming-specific deadline. + let deadline = tokio::time::Instant::now() + PROVIDER_SIGNATURE_FETCH_TIMEOUT; + self.store_chat_signature_from_provider_impl(chat_id, Some(deadline)) + .await } async fn store_chat_signature( diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index bb7ddef87..f63bc8d8e 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -1,6 +1,7 @@ pub mod ports; use crate::attestation::ports::AttestationServiceTrait; +use crate::attestation::STREAM_SIGNATURE_STORE_TIMEOUT; use crate::inference_provider_pool::InferenceProviderPool; use crate::models::ModelsRepository; use crate::responses::models::ResponseId; @@ -22,7 +23,6 @@ use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use tracing::Instrument; -const FINALIZE_TIMEOUT_SECS: u64 = 5; /// A raw provider terminal marker is only expected to contain `data: [DONE]` /// plus its SSE line ending. Bound it before accepting it as a signed terminal /// event so malformed padding cannot be retained or signed. @@ -213,8 +213,8 @@ where Box::pin(async move { match tokio::time::timeout( - Duration::from_secs(FINALIZE_TIMEOUT_SECS), - attestation_service.store_chat_signature_from_provider(&chat_id), + STREAM_SIGNATURE_STORE_TIMEOUT, + attestation_service.store_stream_chat_signature_from_provider(&chat_id), ) .await { @@ -227,7 +227,7 @@ where %organization_id, %model_id, "Timeout storing chat signature after {}s", - FINALIZE_TIMEOUT_SECS + STREAM_SIGNATURE_STORE_TIMEOUT.as_secs() ); // The provider-store implementation normally unpins after // it completes. A timeout cancels that future before its From 95eb797db4beeea1985fd83ae70465c346a16f2a Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 01:02:48 +0000 Subject: [PATCH 10/10] perf(attestation): cache the batched signature upsert --- .../tests/e2e_all/prepared_statement_cache.rs | 70 +++++++++++++++++++ .../database/src/repositories/attestation.rs | 2 +- 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs index df58e7d51..ab8757f49 100644 --- a/crates/api/tests/e2e_all/prepared_statement_cache.rs +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -325,3 +325,73 @@ async fn usage_transaction_recovers_from_missing_balance_statement_without_doubl assert_usage_recovers_without_double_charging(UsageStatementInvalidation::DeallocatedBalance) .await; } + +#[tokio::test] +async fn batched_signatures_reuse_a_prepared_statement() { + use services::attestation::{ports::AttestationRepository, ChatSignature, SignatureKind}; + + let (_server, database) = setup_test_server_with_database().await; + let pool = database.pool(); + // Keep one connection available so both repository calls and the session + // inspection below use the same PostgreSQL prepared-statement cache. + let mut held_connections = Vec::new(); + for _ in 1..pool.status().expect("database pool").max_size { + held_connections.push(pool.get().await.expect("hold another connection")); + } + + let chat_id = format!("batch-cache-{}", uuid::Uuid::new_v4()); + let signatures = |text: &str| { + ["ecdsa", "ed25519"] + .into_iter() + .map(|algo| ChatSignature { + text: text.to_string(), + signature: format!("synthetic-{algo}"), + signing_address: "synthetic-address".to_string(), + signing_algo: algo.to_string(), + signature_kind: Some(SignatureKind::ProviderTee), + }) + .collect() + }; + let cached_names = || async { + pool.get() + .await + .expect("inspect the warmed connection") + .query( + "SELECT name FROM pg_prepared_statements WHERE statement LIKE $1 ORDER BY name", + &[&"INSERT INTO chat_signatures%"], + ) + .await + .expect("inspect prepared statements") + .into_iter() + .map(|row| row.get::<_, String>("name")) + .collect::>() + }; + + database + .attestation + .add_chat_signatures(&chat_id, signatures("first")) + .await + .expect("store both signatures"); + let first_names = cached_names().await; + assert_eq!(first_names.len(), 1, "the batch must be cached"); + + database + .attestation + .add_chat_signatures(&chat_id, signatures("updated")) + .await + .expect("update both signatures"); + assert_eq!( + cached_names().await, + first_names, + "reuse the same statement" + ); + for algo in ["ecdsa", "ed25519"] { + let stored = database + .attestation + .get_chat_signature(&chat_id, algo) + .await + .expect("read the updated signature"); + assert_eq!(stored.text, "updated"); + assert_eq!(stored.signature_kind, Some(SignatureKind::ProviderTee)); + } +} diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index d1d51d7ab..f975154d0 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -149,7 +149,7 @@ impl AttestationRepository for PgAttestationRepository { ); client - .execute(sql.as_str(), ¶ms) + .cached_execute(sql.as_str(), ¶ms) .await .map_err(|e| AttestationError::RepositoryError(e.to_string()))?;