starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(core,store): merge first-attempt persistence and reuse hot-path state

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
zyycn
2026-07-26 20:20:15 +08:00
co-authored by Claude
parent 5c66dc198d
commit 7e7b88a59e
6 changed files with 277 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/engine/coordinator.rs
+39 -33
View File
@@ -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<S: ?Sized> {
client_api_key_ref: crate::policy::ClientApiKeyId,
timing_started_at: Instant,
deadline: SystemTime,
/// 会话级 deadline 计时器；deadline 固定，帧循环内复用而非逐事件新建。
deadline_timer: Fuse<Delay>,
pending_request: Option<NewModelRequest>,
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<Delay>,
) -> 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),
}
}
backend/crates/gateway-core/src/engine/execution.rs
+1 -1
View File
@@ -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 {
backend/crates/gateway-core/src/engine/mod.rs
+10
View File
@@ -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,
backend/crates/gateway-core/src/operation/mod.rs
+13 -1
View File
@@ -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<u64> {
@@ -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 {
backend/crates/gateway-store/src/postgres/execution.rs
+131 -42
View File
@@ -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::<Utc>::from(request.started_at),
deadline_at: DateTime::<Utc>::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::<Utc>::from(request.started_at),
deadline_at: DateTime::<Utc>::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<u64>, field: &'static str) -> StoreResult<Option<i64>> {
value.map(|value| to_i64(value, field)).transpose()
}
backend/crates/gateway-store/tests/postgres/execution.rs
+83 -1
View File
@@ -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 {