From 7e7b88a59e2c69d7eeaa03ed23602237490b14ef Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 26 Jul 2026 20:20:15 +0800 Subject: [PATCH] perf(core,store): merge first-attempt persistence and reuse hot-path state Combine the model-request insert with the first attempt into one ExecutionStore call (default implementation keeps the sequential behavior; the Postgres store overrides it with a single INSERT), removing one database round trip from the pre-first-token critical path. Reuse a session-scoped deadline timer in the event pull loop instead of registering a new global timer per event, and attach continuation session state in place so a uniquely held payload is no longer force-deep-copied by clone-before-make_mut. Co-authored-by: Claude --- .../gateway-core/src/engine/coordinator.rs | 72 ++++---- .../gateway-core/src/engine/execution.rs | 2 +- backend/crates/gateway-core/src/engine/mod.rs | 10 + .../crates/gateway-core/src/operation/mod.rs | 14 +- .../gateway-store/src/postgres/execution.rs | 173 +++++++++++++----- .../gateway-store/tests/postgres/execution.rs | 84 ++++++++- 6 files changed, 277 insertions(+), 78 deletions(-) diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 331836aa9..a1c8fe326 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -23,6 +23,7 @@ use crate::event::{ }; use crate::operation::{Operation, ProviderSessionState, RetrySafety}; use crate::routing::RoutingPlan; +use futures::future::Fuse; use futures::{FutureExt, StreamExt, pin_mut, select_biased}; use futures_timer::Delay; @@ -73,6 +74,12 @@ where client_api_key_ref, timing_started_at, deadline, + deadline_timer: Delay::new( + deadline + .duration_since(SystemTime::now()) + .unwrap_or(Duration::ZERO), + ) + .fuse(), pending_request: Some(request), request_persisted: false, operation, @@ -143,6 +150,8 @@ pub struct ResponseExecutionSession { client_api_key_ref: crate::policy::ClientApiKeyId, timing_started_at: Instant, deadline: SystemTime, + /// 会话级 deadline 计时器；deadline 固定，帧循环内复用而非逐事件新建。 + deadline_timer: Fuse, pending_request: Option, request_persisted: bool, operation: Operation, @@ -374,6 +383,7 @@ where &mut current.stream, self.cancellation.clone(), self.deadline, + &mut self.deadline_timer, ) .await }; @@ -614,21 +624,30 @@ where account.clone(), )); } - self.persist_request().await?; - self.engine - .store() - .record_attempt(AttemptRecord { - request_id: self.request_id.clone(), - attempt_count: next_attempt, - trigger, - provider_kind: metadata.provider().clone(), - provider_account_id: metadata.provider_account_id().cloned(), - provider_account_ref: metadata.provider_account_id().cloned(), - upstream_model_id: metadata.upstream_model().clone(), - upstream_transport: metadata.transport().as_str().to_owned(), - http_version: None, - }) - .await?; + let attempt_record = AttemptRecord { + request_id: self.request_id.clone(), + attempt_count: next_attempt, + trigger, + provider_kind: metadata.provider().clone(), + provider_account_id: metadata.provider_account_id().cloned(), + provider_account_ref: metadata.provider_account_id().cloned(), + upstream_model_id: metadata.upstream_model().clone(), + upstream_transport: metadata.transport().as_str().to_owned(), + http_version: None, + }; + if self.request_persisted { + self.engine.store().record_attempt(attempt_record).await?; + } else { + let request = self + .pending_request + .take() + .ok_or(EngineError::InvalidDeliveryState)?; + self.engine + .store() + .create_model_request_with_attempt(request, attempt_record) + .await?; + self.request_persisted = true; + } self.attempts = next_attempt.get(); self.current = Some(CurrentAttempt { stream, @@ -642,19 +661,6 @@ where Ok(()) } - async fn persist_request(&mut self) -> Result<(), EngineError> { - if self.request_persisted { - return Ok(()); - } - let request = self - .pending_request - .take() - .ok_or(EngineError::InvalidDeliveryState)?; - self.engine.store().create_model_request(request).await?; - self.request_persisted = true; - Ok(()) - } - async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> { let elapsed = elapsed_ms(self.timing_started_at); observe_event_timing(&mut self.timings, event, elapsed); @@ -1225,17 +1231,17 @@ async fn poll_stream_item( stream: &mut ProviderStream, cancellation: CancellationToken, deadline: SystemTime, + mut deadline_timer: &mut Fuse, ) -> PollBoundary { - let Ok(remaining) = deadline.duration_since(SystemTime::now()) else { + if SystemTime::now() >= deadline { return PollBoundary::Deadline; - }; + } let next = stream.next().fuse(); let cancelled = cancellation.cancelled().fuse(); - let timeout = Delay::new(remaining).fuse(); - pin_mut!(next, cancelled, timeout); + pin_mut!(next, cancelled); select_biased! { _ = cancelled => PollBoundary::Cancelled, - _ = timeout => PollBoundary::Deadline, + _ = deadline_timer => PollBoundary::Deadline, item = next => PollBoundary::Item(item), } } diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index bff85c85f..23183cbd9 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -868,7 +868,7 @@ fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeCont { return; } - *operation = operation.clone().with_provider_session_state(state.clone()); + operation.set_provider_session_state(state.clone()); } pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index f906a70d5..497fe53d0 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -596,6 +596,16 @@ pub struct RecoveryReport { pub trait ExecutionStore: Send + Sync { async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError>; async fn record_attempt(&self, attempt: AttemptRecord) -> Result<(), StoreError>; + /// 请求插入与首次 attempt 合并持久化；两者写同一行，支持合并写的 + /// store 可覆写为单次往返，缩短首 token 前的关键路径。 + async fn create_model_request_with_attempt( + &self, + request: NewModelRequest, + attempt: AttemptRecord, + ) -> Result<(), StoreError> { + self.create_model_request(request).await?; + self.record_attempt(attempt).await + } async fn mark_send_state( &self, request_id: &ModelRequestId, diff --git a/backend/crates/gateway-core/src/operation/mod.rs b/backend/crates/gateway-core/src/operation/mod.rs index 920ea0b86..24be45385 100644 --- a/backend/crates/gateway-core/src/operation/mod.rs +++ b/backend/crates/gateway-core/src/operation/mod.rs @@ -269,10 +269,15 @@ impl GenerateRequest { /// 附着同一客户端连接上一轮由 Provider 返回的不透明状态。 #[must_use] pub fn with_provider_session_state(mut self, state: ProviderSessionState) -> Self { - Arc::make_mut(&mut self.payload).provider_session_state = Some(state); + self.set_provider_session_state(state); self } + /// 原地附着会话状态；payload 独占时不复制正文。 + pub fn set_provider_session_state(&mut self, state: ProviderSessionState) { + Arc::make_mut(&mut self.payload).provider_session_state = Some(state); + } + /// 返回最大输出 token 数。 #[must_use] pub fn max_output_tokens(&self) -> Option { @@ -597,6 +602,13 @@ impl Operation { } } + /// 原地附着 Provider 私有状态；payload 独占时不复制正文。 + pub fn set_provider_session_state(&mut self, state: ProviderSessionState) { + if let Self::Generate(request) = self { + request.set_provider_session_state(state); + } + } + /// 返回稳定 operation 分类。 #[must_use] pub const fn kind(&self) -> OperationKind { diff --git a/backend/crates/gateway-store/src/postgres/execution.rs b/backend/crates/gateway-store/src/postgres/execution.rs index d8545891b..0b25d115a 100644 --- a/backend/crates/gateway-store/src/postgres/execution.rs +++ b/backend/crates/gateway-store/src/postgres/execution.rs @@ -295,6 +295,12 @@ pub struct ModelRequestRecoveryReport { #[async_trait] pub trait ModelRequestRepository: Send + Sync { async fn insert_model_request(&self, request: NewModelRequest) -> StoreResult<()>; + /// 请求行与首次 attempt 列一次插入；`attempt.attempt_count` 必须为 1。 + async fn insert_model_request_with_first_attempt( + &self, + request: NewModelRequest, + attempt: ModelRequestAttemptStart, + ) -> StoreResult<()>; async fn begin_model_request_attempt( &self, attempt: ModelRequestAttemptStart, @@ -383,6 +389,63 @@ impl ModelRequestRepository for PgExecutionStore { Ok(()) } + async fn insert_model_request_with_first_attempt( + &self, + request: NewModelRequest, + attempt: ModelRequestAttemptStart, + ) -> StoreResult<()> { + request.validate()?; + attempt.validate()?; + if attempt.attempt_count != 1 || attempt.model_request_id != request.id { + return Err(invalid("first attempt must target the inserted request")); + } + sqlx::query( + "insert into model_requests ( + id, client_api_key_id, client_api_key_ref, config_revision, protocol, + operation, endpoint, client_transport, requested_model_id, + client_ip, user_agent, reasoning_effort, + reasoning_preset, request_kind, subagent_kind, compact, + image_generation_requested, started_at, deadline_at, + provider_kind, provider_account_id, provider_account_ref, + upstream_model_id, upstream_transport, http_version, + attempt_count, upstream_send_state + ) values ( + $1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12, + $13, $14, $15, $16, $17, $18, $19, + $20, $21, $22, $23, $24, $25, 1, 'not_sent' + )", + ) + .bind(request.id) + .bind(request.client_api_key_id) + .bind(request.client_api_key_ref) + .bind(to_i64(request.config_revision, "config_revision")?) + .bind(request.protocol) + .bind(request.operation) + .bind(request.endpoint) + .bind(request.client_transport) + .bind(request.requested_model_id) + .bind(request.client_ip) + .bind(request.user_agent) + .bind(request.reasoning_effort) + .bind(request.reasoning_preset) + .bind(request.request_kind) + .bind(request.subagent_kind) + .bind(request.compact) + .bind(request.image_generation_requested) + .bind(request.started_at) + .bind(request.deadline_at) + .bind(attempt.provider_kind) + .bind(attempt.provider_account_id) + .bind(attempt.provider_account_ref) + .bind(attempt.upstream_model_id) + .bind(attempt.upstream_transport) + .bind(attempt.http_version) + .execute(&self.pool) + .await + .map_err(|_| postgres_unavailable("insert model request with first attempt"))?; + Ok(()) + } + async fn begin_model_request_attempt( &self, attempt: ModelRequestAttemptStart, @@ -632,53 +695,15 @@ impl ExecutionStore for PgExecutionStore { &self, request: CoreNewModelRequest, ) -> Result<(), CoreStoreError> { - self.insert_model_request(NewModelRequest { - id: request.id.as_str().to_owned(), - client_api_key_id: request - .client_api_key_id - .as_ref() - .map(|id| id.as_str().to_owned()), - client_api_key_ref: request.client_api_key_ref.as_str().to_owned(), - config_revision: request.config_revision.get(), - protocol: request.protocol, - operation: request.operation.as_str().to_owned(), - endpoint: request.endpoint, - client_transport: request.client_transport, - requested_model_id: request.requested_model.as_str().to_owned(), - client_ip: request.client_ip.map(|address| address.to_string()), - user_agent: request.user_agent, - reasoning_effort: request.reasoning_effort, - reasoning_preset: request.reasoning_preset, - request_kind: request.request_kind, - subagent_kind: request.subagent_kind, - compact: request.compact, - image_generation_requested: request.image_generation_requested, - started_at: DateTime::::from(request.started_at), - deadline_at: DateTime::::from(request.deadline_at), - }) - .await - .map_err(core_store_error) + self.insert_model_request(new_model_request_row(request)) + .await + .map_err(core_store_error) } async fn record_attempt(&self, attempt: CoreAttemptRecord) -> Result<(), CoreStoreError> { let expected_count = attempt.attempt_count.get(); let persisted = self - .begin_model_request_attempt(ModelRequestAttemptStart { - model_request_id: attempt.request_id.as_str().to_owned(), - attempt_count: expected_count, - provider_kind: attempt.provider_kind.as_str().to_owned(), - provider_account_id: attempt - .provider_account_id - .as_ref() - .map(|id| id.as_str().to_owned()), - provider_account_ref: attempt - .provider_account_ref - .as_ref() - .map(|id| id.as_str().to_owned()), - upstream_model_id: attempt.upstream_model_id.as_str().to_owned(), - upstream_transport: attempt.upstream_transport, - http_version: attempt.http_version, - }) + .begin_model_request_attempt(attempt_start_row(attempt)) .await .map_err(core_store_error)?; if persisted == expected_count { @@ -688,6 +713,24 @@ impl ExecutionStore for PgExecutionStore { } } + async fn create_model_request_with_attempt( + &self, + request: CoreNewModelRequest, + attempt: CoreAttemptRecord, + ) -> Result<(), CoreStoreError> { + // 合并写只对首次 attempt 成立（插入即带 attempt 列）；其余走基础两步。 + if attempt.attempt_count.get() != 1 { + self.create_model_request(request).await?; + return self.record_attempt(attempt).await; + } + self.insert_model_request_with_first_attempt( + new_model_request_row(request), + attempt_start_row(attempt), + ) + .await + .map_err(core_store_error) + } + async fn mark_send_state( &self, request_id: &ModelRequestId, @@ -937,6 +980,52 @@ fn core_store_error(error: StoreError) -> CoreStoreError { CoreStoreError::new(kind) } +fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest { + NewModelRequest { + id: request.id.as_str().to_owned(), + client_api_key_id: request + .client_api_key_id + .as_ref() + .map(|id| id.as_str().to_owned()), + client_api_key_ref: request.client_api_key_ref.as_str().to_owned(), + config_revision: request.config_revision.get(), + protocol: request.protocol, + operation: request.operation.as_str().to_owned(), + endpoint: request.endpoint, + client_transport: request.client_transport, + requested_model_id: request.requested_model.as_str().to_owned(), + client_ip: request.client_ip.map(|address| address.to_string()), + user_agent: request.user_agent, + reasoning_effort: request.reasoning_effort, + reasoning_preset: request.reasoning_preset, + request_kind: request.request_kind, + subagent_kind: request.subagent_kind, + compact: request.compact, + image_generation_requested: request.image_generation_requested, + started_at: DateTime::::from(request.started_at), + deadline_at: DateTime::::from(request.deadline_at), + } +} + +fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart { + ModelRequestAttemptStart { + model_request_id: attempt.request_id.as_str().to_owned(), + attempt_count: attempt.attempt_count.get(), + provider_kind: attempt.provider_kind.as_str().to_owned(), + provider_account_id: attempt + .provider_account_id + .as_ref() + .map(|id| id.as_str().to_owned()), + provider_account_ref: attempt + .provider_account_ref + .as_ref() + .map(|id| id.as_str().to_owned()), + upstream_model_id: attempt.upstream_model_id.as_str().to_owned(), + upstream_transport: attempt.upstream_transport, + http_version: attempt.http_version, + } +} + fn optional_i64(value: Option, field: &'static str) -> StoreResult> { value.map(|value| to_i64(value, field)).transpose() } diff --git a/backend/crates/gateway-store/tests/postgres/execution.rs b/backend/crates/gateway-store/tests/postgres/execution.rs index 3310b0ee9..cad4e52f3 100644 --- a/backend/crates/gateway-store/tests/postgres/execution.rs +++ b/backend/crates/gateway-store/tests/postgres/execution.rs @@ -4,7 +4,9 @@ use gateway_core::engine::{ ExecutionOutcome, ExecutionStore, ModelRequestFinalization as CoreModelRequestFinalization, ModelRequestId, ModelRequestTimings as CoreModelRequestTimings, UpstreamSendState, }; -use gateway_store::postgres::{ModelRequestRepository, NewModelRequest, PgExecutionStore}; +use gateway_store::postgres::{ + ModelRequestAttemptStart, ModelRequestRepository, NewModelRequest, PgExecutionStore, +}; use super::TestDatabase; @@ -41,6 +43,86 @@ fn model_request_rejects_mismatched_client_key_live_id() { assert!(request.validate().is_err()); } +#[tokio::test] +async fn merged_first_attempt_insert_should_match_sequential_semantics() { + let Some(database) = TestDatabase::create("execution_merged_insert").await else { + return; + }; + let repository = PgExecutionStore::new(database.pool.clone()); + let started_at = Utc::now(); + let request = NewModelRequest { + id: "req_merged".to_owned(), + client_api_key_id: None, + client_api_key_ref: "key_merged".to_owned(), + config_revision: 1, + protocol: "openai".to_owned(), + operation: "responses".to_owned(), + endpoint: "/v1/responses".to_owned(), + client_transport: "http_sse".to_owned(), + requested_model_id: "coding".to_owned(), + client_ip: None, + user_agent: None, + reasoning_effort: None, + reasoning_preset: None, + request_kind: None, + subagent_kind: None, + compact: false, + image_generation_requested: false, + started_at, + deadline_at: started_at + Duration::seconds(30), + }; + let attempt = ModelRequestAttemptStart { + model_request_id: "req_merged".to_owned(), + attempt_count: 1, + provider_kind: "openai".to_owned(), + provider_account_id: Some("acct_merged".to_owned()), + provider_account_ref: Some("acct_merged".to_owned()), + upstream_model_id: "gpt-test".to_owned(), + upstream_transport: "http_sse".to_owned(), + http_version: None, + }; + + repository + .insert_model_request_with_first_attempt(request, attempt) + .await + .expect("merged insert"); + + let (attempt_count, send_state, provider_kind, outcome): (i32, String, String, String) = + sqlx::query_as( + "select attempt_count, upstream_send_state, provider_kind, outcome + from model_requests where id = 'req_merged'", + ) + .fetch_one(&database.pool) + .await + .expect("load merged request"); + assert_eq!( + ( + attempt_count, + send_state.as_str(), + provider_kind.as_str(), + outcome.as_str() + ), + (1, "not_sent", "openai", "running") + ); + + // 后续 attempt 沿用常规 CAS 递增路径。 + let second = repository + .begin_model_request_attempt(ModelRequestAttemptStart { + model_request_id: "req_merged".to_owned(), + attempt_count: 2, + provider_kind: "openai".to_owned(), + provider_account_id: Some("acct_next".to_owned()), + provider_account_ref: Some("acct_next".to_owned()), + upstream_model_id: "gpt-test".to_owned(), + upstream_transport: "http_sse".to_owned(), + http_version: None, + }) + .await + .expect("second attempt"); + assert_eq!(second, 2); + database.close().await; +} + #[tokio::test] async fn downstream_commit_should_atomically_record_http_status_once() { let Some(database) = TestDatabase::create("execution_commit_status").await else {