From cc42e78334d0036344d35fa29010c6fd80e06486 Mon Sep 17 00:00:00 2001 From: zyycn Date: Wed, 29 Jul 2026 00:11:26 +0800 Subject: [PATCH] refactor(store): buffer data-plane side effects Move PostgreSQL execution observations and Redis release, circuit, and continuation writes onto bounded non-blocking queues. Supervise independent flush workers so database or coordination failures can drop recoverable side effects without delaying protocol delivery. Co-authored-by: Codex --- backend/crates/gateway-store/src/lib.rs | 110 +++++- .../src/postgres/execution_buffer.rs | 274 ++++++++++++++ .../crates/gateway-store/src/postgres/mod.rs | 2 + .../src/redis/coordination_buffer.rs | 345 ++++++++++++++++++ backend/crates/gateway-store/src/redis/mod.rs | 2 + .../tests/coordination_buffer.rs | 247 +++++++++++++ .../gateway-store/tests/execution_buffer.rs | 156 ++++++++ backend/crates/gateway-store/tests/main.rs | 2 + 8 files changed, 1118 insertions(+), 20 deletions(-) create mode 100644 backend/crates/gateway-store/src/postgres/execution_buffer.rs create mode 100644 backend/crates/gateway-store/src/redis/coordination_buffer.rs create mode 100644 backend/crates/gateway-store/tests/coordination_buffer.rs create mode 100644 backend/crates/gateway-store/tests/execution_buffer.rs diff --git a/backend/crates/gateway-store/src/lib.rs b/backend/crates/gateway-store/src/lib.rs index cc2351981..36a6cb1c1 100644 --- a/backend/crates/gateway-store/src/lib.rs +++ b/backend/crates/gateway-store/src/lib.rs @@ -21,7 +21,7 @@ use gateway_core::CoreStorePorts; use gateway_core::health::{HealthProbe, HealthState}; use gateway_core::provider_ports::ProviderStorePorts; use gateway_core::task::{ - ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDisabledReason, WorkerId, + DaemonRestartPolicy, ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerId, WorkerKind, WorkerLeaderLeasePort, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError, }; @@ -318,28 +318,40 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { }), ); - let execution = Arc::new(postgres::PgExecutionStore::new(pool.clone())); + let execution_repository = Arc::new(postgres::PgExecutionStore::new(pool.clone())); + let (execution, execution_writer) = + postgres::BufferedExecutionStore::new(Arc::clone(&execution_repository)); + let execution = Arc::new(execution); let retention = Arc::new(postgres::PgRetentionRepository::new(pool.clone())); - let core_ports = CoreStorePorts::new( - execution.clone(), - ( - Arc::new(redis::RedisClientAdmissionRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - )?), - Arc::new(postgres::PgClientAdmissionRecoveryRepository::new( - pool.clone(), - )), - ), + let admissions: Arc = Arc::new( + redis::RedisClientAdmissionRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?, + ); + let circuits: Arc = Arc::new(redis::RedisProviderCircuitRepository::new( redis_connection.clone(), REDIS_NAMESPACE, redis::ProviderCircuitPolicy::default(), - )?), + )?); + let continuation: Arc = Arc::new(redis::RedisNativeContinuationRepository::new( redis_connection.clone(), REDIS_NAMESPACE, - )?), + )?); + let (admissions, admission_release_writer) = + redis::BufferedClientAdmissionPort::new(admissions); + let (circuits, circuit_feedback_writer) = redis::BufferedProviderCircuitPort::new(circuits); + let (continuation, continuation_writer) = + redis::BufferedNativeContinuationPort::new(continuation); + let core_ports = CoreStorePorts::new( + execution, + ( + Arc::new(admissions), + Arc::new(postgres::PgClientAdmissionRecoveryRepository::new( + pool.clone(), + )), + ), + Arc::new(circuits), + Arc::new(continuation), ( Arc::new(postgres::PgRuntimeSnapshotRepository::new(pool.clone())), Arc::new(redis::RedisRuntimeChangeRepository::new( @@ -373,7 +385,14 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { connection: redis_connection, }), ]; - let worker_contributions = store_worker_contributions(execution, retention)?; + let worker_contributions = store_worker_contributions( + execution_repository, + execution_writer, + admission_release_writer, + circuit_feedback_writer, + continuation_writer, + retention, + )?; Ok(StoreBundle { admin_ports, core_ports, @@ -386,12 +405,27 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { fn store_worker_contributions( execution: Arc, + execution_writer: postgres::ExecutionObservationWriter, + admission_release_writer: redis::ClientAdmissionReleaseWriter, + circuit_feedback_writer: redis::ProviderCircuitFeedbackWriter, + continuation_writer: redis::NativeContinuationWriter, retention: Arc, ) -> StoreResult> { let stale_id = WorkerId::try_new(WorkerKind::StaleModelRequestRecovery, "postgres") .map_err(worker_definition_error)?; let retention_id = WorkerId::try_new(WorkerKind::Retention, "postgres").map_err(worker_definition_error)?; + let ops_flush_id = + WorkerId::try_new(WorkerKind::OpsFlush, "postgres").map_err(worker_definition_error)?; + let admission_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_admission") + .map_err(worker_definition_error)?; + let circuit_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_circuit") + .map_err(worker_definition_error)?; + let continuation_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_continuation") + .map_err(worker_definition_error)?; + let ops_flush_restart = + DaemonRestartPolicy::try_new(Duration::from_secs(1), Duration::from_secs(60)) + .map_err(worker_definition_error)?; Ok(vec![ WorkerContribution::Registration(scheduled_worker( stale_id, @@ -403,10 +437,46 @@ fn store_worker_contributions( Duration::from_secs(60 * 60), Box::new(RetentionTask { retention }), )?), - WorkerContribution::Disabled { - kind: WorkerKind::OpsFlush, - reason: WorkerDisabledReason::NoBufferedOpsEvents, - }, + WorkerContribution::Registration( + WorkerRegistration::try_new( + ops_flush_id, + WorkerRunnable::Daemon { + restart: ops_flush_restart, + task: Box::new(execution_writer), + }, + ) + .map_err(worker_definition_error)?, + ), + WorkerContribution::Registration( + WorkerRegistration::try_new( + admission_flush_id, + WorkerRunnable::Daemon { + restart: ops_flush_restart, + task: Box::new(admission_release_writer), + }, + ) + .map_err(worker_definition_error)?, + ), + WorkerContribution::Registration( + WorkerRegistration::try_new( + circuit_flush_id, + WorkerRunnable::Daemon { + restart: ops_flush_restart, + task: Box::new(circuit_feedback_writer), + }, + ) + .map_err(worker_definition_error)?, + ), + WorkerContribution::Registration( + WorkerRegistration::try_new( + continuation_flush_id, + WorkerRunnable::Daemon { + restart: ops_flush_restart, + task: Box::new(continuation_writer), + }, + ) + .map_err(worker_definition_error)?, + ), ]) } diff --git a/backend/crates/gateway-store/src/postgres/execution_buffer.rs b/backend/crates/gateway-store/src/postgres/execution_buffer.rs new file mode 100644 index 000000000..0987bd39b --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/execution_buffer.rs @@ -0,0 +1,274 @@ +//! 数据面执行观测的非阻塞 PostgreSQL 写入队列。 + +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::time::SystemTime; + +use async_trait::async_trait; +use gateway_core::engine::{ + AttemptRecord, CancellationToken, ExecutionStore, IntermediateFailure, + ModelRequestFinalization, ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport, + UpstreamSendState, +}; +use gateway_core::error::StoreError; +use gateway_core::task::{DaemonTask, WorkerTaskError}; +use tokio::sync::{Mutex, mpsc}; + +const DEFAULT_QUEUE_CAPACITY: usize = 4_096; + +/// 将数据面观测写入转换为有界、非阻塞的进程内命令。 +/// +/// 队列满、worker 尚未启动或已经退出时只丢弃观测并记录告警；协议数据面不会 +/// 等待 PostgreSQL，也不会看到 Store 错误。启动恢复仍直接访问底层 Store。 +pub struct BufferedExecutionStore { + inner: Arc, + sender: mpsc::Sender, +} + +impl BufferedExecutionStore { + #[must_use] + pub fn new(inner: Arc) -> (Self, ExecutionObservationWriter) { + Self::with_capacity( + inner, + NonZeroUsize::new(DEFAULT_QUEUE_CAPACITY).expect("queue capacity is non-zero"), + ) + } + + #[must_use] + pub fn with_capacity( + inner: Arc, + capacity: NonZeroUsize, + ) -> (Self, ExecutionObservationWriter) { + let (sender, receiver) = mpsc::channel(capacity.get()); + ( + Self { + inner: Arc::clone(&inner), + sender, + }, + ExecutionObservationWriter { + inner, + receiver: Mutex::new(receiver), + }, + ) + } + + fn enqueue(&self, write: ExecutionObservationWrite) { + if let Err(error) = self.sender.try_send(write) { + let (reason, write) = match error { + mpsc::error::TrySendError::Full(write) => ("full", write), + mpsc::error::TrySendError::Closed(write) => ("closed", write), + }; + tracing::warn!( + operation = write.operation(), + reason, + "执行观测队列不可用，已丢弃本次写入" + ); + } + } +} + +#[async_trait] +impl ExecutionStore for BufferedExecutionStore +where + S: ExecutionStore + ?Sized, +{ + async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::Create(Box::new(request))); + Ok(()) + } + + async fn record_attempt(&self, attempt: AttemptRecord) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::Attempt(Box::new(attempt))); + Ok(()) + } + + async fn create_model_request_with_attempt( + &self, + request: NewModelRequest, + attempt: AttemptRecord, + ) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::CreateWithAttempt(Box::new(( + request, attempt, + )))); + Ok(()) + } + + async fn mark_send_state( + &self, + request_id: &ModelRequestId, + state: UpstreamSendState, + ) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::MarkSendState { + request_id: request_id.clone(), + state, + }); + Ok(()) + } + + async fn mark_downstream_committed( + &self, + request_id: &ModelRequestId, + committed_at: SystemTime, + client_status_code: Option, + ) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::MarkDownstreamCommitted { + request_id: request_id.clone(), + committed_at, + client_status_code, + }); + Ok(()) + } + + async fn record_client_status( + &self, + request_id: &ModelRequestId, + client_status_code: u16, + ) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::RecordClientStatus { + request_id: request_id.clone(), + client_status_code, + }); + Ok(()) + } + + async fn record_intermediate_failure( + &self, + failure: IntermediateFailure, + ) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::IntermediateFailure(Box::new( + failure, + ))); + Ok(()) + } + + async fn record_probe_failure(&self, failure: ProbeFailure) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::ProbeFailure(Box::new(failure))); + Ok(()) + } + + async fn finalize_model_request( + &self, + finalization: ModelRequestFinalization, + ) -> Result<(), StoreError> { + self.enqueue(ExecutionObservationWrite::Finalize(Box::new(finalization))); + Ok(()) + } + + async fn recover_expired(&self, now: SystemTime) -> Result { + self.inner.recover_expired(now).await + } +} + +/// 由 Host 监督的单消费者，保证同一请求的观测命令按入队顺序落库。 +pub struct ExecutionObservationWriter { + inner: Arc, + receiver: Mutex>, +} + +impl DaemonTask for ExecutionObservationWriter +where + S: ExecutionStore + ?Sized, +{ + fn run( + &self, + cancellation: CancellationToken, + ) -> futures::future::BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut receiver = self.receiver.lock().await; + loop { + let write = tokio::select! { + () = cancellation.cancelled() => return Ok(()), + write = receiver.recv() => write, + }; + let Some(write) = write else { + return Err(WorkerTaskError::safe("execution observation queue closed")); + }; + let operation = write.operation(); + if let Err(error) = write.persist(self.inner.as_ref()).await { + tracing::warn!( + operation, + error_kind = ?error.kind(), + "执行观测写入失败，数据面不受影响" + ); + } + } + }) + } +} + +enum ExecutionObservationWrite { + Create(Box), + Attempt(Box), + CreateWithAttempt(Box<(NewModelRequest, AttemptRecord)>), + MarkSendState { + request_id: ModelRequestId, + state: UpstreamSendState, + }, + MarkDownstreamCommitted { + request_id: ModelRequestId, + committed_at: SystemTime, + client_status_code: Option, + }, + RecordClientStatus { + request_id: ModelRequestId, + client_status_code: u16, + }, + IntermediateFailure(Box), + ProbeFailure(Box), + Finalize(Box), +} + +impl ExecutionObservationWrite { + const fn operation(&self) -> &'static str { + match self { + Self::Create(_) => "create_model_request", + Self::Attempt(_) => "record_attempt", + Self::CreateWithAttempt(_) => "create_model_request_with_attempt", + Self::MarkSendState { .. } => "mark_send_state", + Self::MarkDownstreamCommitted { .. } => "mark_downstream_committed", + Self::RecordClientStatus { .. } => "record_client_status", + Self::IntermediateFailure(_) => "record_intermediate_failure", + Self::ProbeFailure(_) => "record_probe_failure", + Self::Finalize(_) => "finalize_model_request", + } + } + + async fn persist(self, store: &S) -> Result<(), StoreError> + where + S: ExecutionStore + ?Sized, + { + match self { + Self::Create(request) => store.create_model_request(*request).await, + Self::Attempt(attempt) => store.record_attempt(*attempt).await, + Self::CreateWithAttempt(write) => { + let (request, attempt) = *write; + store + .create_model_request_with_attempt(request, attempt) + .await + } + Self::MarkSendState { request_id, state } => { + store.mark_send_state(&request_id, state).await + } + Self::MarkDownstreamCommitted { + request_id, + committed_at, + client_status_code, + } => { + store + .mark_downstream_committed(&request_id, committed_at, client_status_code) + .await + } + Self::RecordClientStatus { + request_id, + client_status_code, + } => { + store + .record_client_status(&request_id, client_status_code) + .await + } + Self::IntermediateFailure(failure) => store.record_intermediate_failure(*failure).await, + Self::ProbeFailure(failure) => store.record_probe_failure(*failure).await, + Self::Finalize(finalization) => store.finalize_model_request(*finalization).await, + } + } +} diff --git a/backend/crates/gateway-store/src/postgres/mod.rs b/backend/crates/gateway-store/src/postgres/mod.rs index a4dbe78c1..4ce275b0f 100644 --- a/backend/crates/gateway-store/src/postgres/mod.rs +++ b/backend/crates/gateway-store/src/postgres/mod.rs @@ -11,6 +11,7 @@ mod admin_security_audit; mod admission_recovery; mod client_keys; mod execution; +mod execution_buffer; mod observability; mod ops_events; mod provider_accounts; @@ -22,6 +23,7 @@ pub use admin_security_audit::*; pub use admission_recovery::*; pub use client_keys::*; pub use execution::*; +pub use execution_buffer::*; pub use observability::*; pub use ops_events::*; pub use provider_accounts::*; diff --git a/backend/crates/gateway-store/src/redis/coordination_buffer.rs b/backend/crates/gateway-store/src/redis/coordination_buffer.rs new file mode 100644 index 000000000..3bc9574b2 --- /dev/null +++ b/backend/crates/gateway-store/src/redis/coordination_buffer.rs @@ -0,0 +1,345 @@ +//! 数据面 Redis 写副作用的有界非阻塞队列。 + +use std::num::NonZeroUsize; +use std::sync::Arc; + +use futures::future::{BoxFuture, ready}; +use gateway_core::engine::admission::{ + ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery, + ClientAdmissionRequest, ClientAdmissionRestoreResult, +}; +use gateway_core::engine::continuation::{ + NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId, +}; +use gateway_core::engine::execution::{ + ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort, +}; +use gateway_core::engine::{CancellationToken, ModelRequestId}; +use gateway_core::policy::ClientApiKeyId; +use gateway_core::routing::ProviderKind; +use gateway_core::task::{DaemonTask, WorkerTaskError}; +use tokio::sync::{Mutex, mpsc}; + +const DEFAULT_QUEUE_CAPACITY: usize = 4_096; + +/// 准入读取保持强一致；终态释放只做有界入队，Redis 失败由 worker 吸收。 +#[derive(Clone)] +pub struct BufferedClientAdmissionPort { + inner: Arc, + sender: mpsc::Sender, +} + +impl BufferedClientAdmissionPort { + #[must_use] + pub fn new(inner: Arc) -> (Self, ClientAdmissionReleaseWriter) { + Self::with_capacity( + inner, + NonZeroUsize::new(DEFAULT_QUEUE_CAPACITY).expect("queue capacity is non-zero"), + ) + } + + #[must_use] + pub fn with_capacity( + inner: Arc, + capacity: NonZeroUsize, + ) -> (Self, ClientAdmissionReleaseWriter) { + let (sender, receiver) = mpsc::channel(capacity.get()); + ( + Self { + inner: Arc::clone(&inner), + sender, + }, + ClientAdmissionReleaseWriter { + inner, + receiver: Mutex::new(receiver), + }, + ) + } + + fn enqueue(&self, release: AdmissionRelease) -> bool { + match self.sender.try_send(release) { + Ok(()) => true, + Err(error) => { + let reason = match error { + mpsc::error::TrySendError::Full(_) => "full", + mpsc::error::TrySendError::Closed(_) => "closed", + }; + tracing::warn!( + operation = "release_client_admission", + reason, + "Redis 协调队列不可用，已丢弃本次准入释放；租约将由 TTL 收敛" + ); + false + } + } + } +} + +impl ClientAdmissionPort for BufferedClientAdmissionPort { + fn admit( + &self, + request: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result> { + self.inner.admit(request) + } + + fn release<'a>( + &'a self, + client_api_key_id: &'a ClientApiKeyId, + model_request_id: &'a ModelRequestId, + ) -> BoxFuture<'a, Result> { + let enqueued = self.enqueue(AdmissionRelease { + client_api_key_id: client_api_key_id.clone(), + model_request_id: model_request_id.clone(), + }); + Box::pin(ready(Ok(enqueued))) + } + + fn restore( + &self, + recovery: ClientAdmissionRecovery, + ) -> BoxFuture<'_, Result> { + self.inner.restore(recovery) + } +} + +struct AdmissionRelease { + client_api_key_id: ClientApiKeyId, + model_request_id: ModelRequestId, +} + +pub struct ClientAdmissionReleaseWriter { + inner: Arc, + receiver: Mutex>, +} + +impl DaemonTask for ClientAdmissionReleaseWriter { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut receiver = self.receiver.lock().await; + loop { + let release = tokio::select! { + () = cancellation.cancelled() => return Ok(()), + release = receiver.recv() => release, + }; + let Some(release) = release else { + return Err(WorkerTaskError::safe( + "client admission release queue closed", + )); + }; + if let Err(error) = self + .inner + .release(&release.client_api_key_id, &release.model_request_id) + .await + { + tracing::warn!(%error, "Client admission 后台释放失败，依赖租约 TTL 收敛"); + } + } + }) + } +} + +/// Circuit decision 仍直读 Redis；success/failure feedback 只做有界入队。 +#[derive(Clone)] +pub struct BufferedProviderCircuitPort { + inner: Arc, + sender: mpsc::Sender, +} + +impl BufferedProviderCircuitPort { + #[must_use] + pub fn new(inner: Arc) -> (Self, ProviderCircuitFeedbackWriter) { + Self::with_capacity( + inner, + NonZeroUsize::new(DEFAULT_QUEUE_CAPACITY).expect("queue capacity is non-zero"), + ) + } + + #[must_use] + pub fn with_capacity( + inner: Arc, + capacity: NonZeroUsize, + ) -> (Self, ProviderCircuitFeedbackWriter) { + let (sender, receiver) = mpsc::channel(capacity.get()); + ( + Self { + inner: Arc::clone(&inner), + sender, + }, + ProviderCircuitFeedbackWriter { + inner, + receiver: Mutex::new(receiver), + }, + ) + } + + fn enqueue(&self, feedback: CircuitFeedback) { + if let Err(error) = self.sender.try_send(feedback) { + let reason = match error { + mpsc::error::TrySendError::Full(_) => "full", + mpsc::error::TrySendError::Closed(_) => "closed", + }; + tracing::warn!( + operation = "record_provider_circuit_feedback", + reason, + "Redis 协调队列不可用，已丢弃本次 circuit feedback" + ); + } + } +} + +impl ProviderCircuitPort for BufferedProviderCircuitPort { + fn decision<'a>( + &'a self, + provider_kind: &'a ProviderKind, + ) -> BoxFuture<'a, Result> { + self.inner.decision(provider_kind) + } + + fn observe_failure<'a>( + &'a self, + provider_kind: &'a ProviderKind, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> { + self.enqueue(CircuitFeedback::Failure(provider_kind.clone())); + Box::pin(ready(Ok(()))) + } + + fn observe_success<'a>( + &'a self, + provider_kind: &'a ProviderKind, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> { + self.enqueue(CircuitFeedback::Success(provider_kind.clone())); + Box::pin(ready(Ok(()))) + } +} + +enum CircuitFeedback { + Failure(ProviderKind), + Success(ProviderKind), +} + +pub struct ProviderCircuitFeedbackWriter { + inner: Arc, + receiver: Mutex>, +} + +impl DaemonTask for ProviderCircuitFeedbackWriter { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut receiver = self.receiver.lock().await; + loop { + let feedback = tokio::select! { + () = cancellation.cancelled() => return Ok(()), + feedback = receiver.recv() => feedback, + }; + let Some(feedback) = feedback else { + return Err(WorkerTaskError::safe( + "provider circuit feedback queue closed", + )); + }; + let result = match feedback { + CircuitFeedback::Failure(provider) => { + self.inner.observe_failure(&provider).await + } + CircuitFeedback::Success(provider) => { + self.inner.observe_success(&provider).await + } + }; + if let Err(error) = result { + tracing::warn!(%error, "Provider circuit feedback 后台写入失败"); + } + } + }) + } +} + +/// Continuation resolve 仍直读 Redis；新 pin 只做有界入队。 +#[derive(Clone)] +pub struct BufferedNativeContinuationPort { + inner: Arc, + sender: mpsc::Sender, +} + +impl BufferedNativeContinuationPort { + #[must_use] + pub fn new(inner: Arc) -> (Self, NativeContinuationWriter) { + Self::with_capacity( + inner, + NonZeroUsize::new(DEFAULT_QUEUE_CAPACITY).expect("queue capacity is non-zero"), + ) + } + + #[must_use] + pub fn with_capacity( + inner: Arc, + capacity: NonZeroUsize, + ) -> (Self, NativeContinuationWriter) { + let (sender, receiver) = mpsc::channel(capacity.get()); + ( + Self { + inner: Arc::clone(&inner), + sender, + }, + NativeContinuationWriter { + inner, + receiver: Mutex::new(receiver), + }, + ) + } + + fn enqueue(&self, pin: NativeContinuationPin) { + if let Err(error) = self.sender.try_send(pin) { + let reason = match error { + mpsc::error::TrySendError::Full(_) => "full", + mpsc::error::TrySendError::Closed(_) => "closed", + }; + tracing::warn!( + operation = "record_native_continuation", + reason, + "Redis 协调队列不可用，已丢弃本次 continuation affinity" + ); + } + } +} + +impl NativeContinuationPort for BufferedNativeContinuationPort { + fn resolve<'a>( + &'a self, + previous_response_id: &'a PreviousResponseId, + ) -> BoxFuture<'a, Result, NativeContinuationStoreError>> { + self.inner.resolve(previous_response_id) + } + + fn record<'a>( + &'a self, + pin: NativeContinuationPin, + ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> { + self.enqueue(pin); + Box::pin(ready(Ok(()))) + } +} + +pub struct NativeContinuationWriter { + inner: Arc, + receiver: Mutex>, +} + +impl DaemonTask for NativeContinuationWriter { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut receiver = self.receiver.lock().await; + loop { + let pin = tokio::select! { + () = cancellation.cancelled() => return Ok(()), + pin = receiver.recv() => pin, + }; + let Some(pin) = pin else { + return Err(WorkerTaskError::safe("native continuation queue closed")); + }; + if let Err(error) = self.inner.record(pin).await { + tracing::warn!(%error, "Continuation affinity 后台写入失败"); + } + } + }) + } +} diff --git a/backend/crates/gateway-store/src/redis/mod.rs b/backend/crates/gateway-store/src/redis/mod.rs index 1a32be9ef..775e32f4f 100644 --- a/backend/crates/gateway-store/src/redis/mod.rs +++ b/backend/crates/gateway-store/src/redis/mod.rs @@ -8,6 +8,7 @@ use sha2::{Digest, Sha256}; use async_trait::async_trait; mod client_admission; +mod coordination_buffer; mod credential_cooldown; mod credential_leases; mod credential_state; @@ -20,6 +21,7 @@ mod runtime_change; pub(crate) mod worker_lease; pub use client_admission::*; +pub use coordination_buffer::*; pub use credential_cooldown::*; pub use credential_leases::*; pub use credential_state::*; diff --git a/backend/crates/gateway-store/tests/coordination_buffer.rs b/backend/crates/gateway-store/tests/coordination_buffer.rs new file mode 100644 index 000000000..c5995e3ab --- /dev/null +++ b/backend/crates/gateway-store/tests/coordination_buffer.rs @@ -0,0 +1,247 @@ +use std::num::NonZeroUsize; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use futures::future::BoxFuture; +use gateway_core::engine::admission::{ + ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery, + ClientAdmissionRequest, ClientAdmissionRestoreResult, +}; +use gateway_core::engine::continuation::{ + NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId, +}; +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::engine::execution::{ + ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort, +}; +use gateway_core::engine::{CancellationToken, ModelRequestId}; +use gateway_core::policy::ClientApiKeyId; +use gateway_core::routing::ProviderKind; +use gateway_store::redis::{ + BufferedClientAdmissionPort, BufferedNativeContinuationPort, BufferedProviderCircuitPort, +}; + +#[derive(Default)] +struct RecordingCoordination { + operations: Mutex>, +} + +impl RecordingCoordination { + fn record(&self, operation: &'static str) { + self.operations + .lock() + .expect("coordination operations lock") + .push(operation); + } + + fn operations(&self) -> Vec<&'static str> { + self.operations + .lock() + .expect("coordination operations lock") + .clone() + } +} + +impl ClientAdmissionPort for RecordingCoordination { + fn admit( + &self, + _: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Ok(ClientAdmissionDecision::Granted) }) + } + + fn release<'a>( + &'a self, + _: &'a ClientApiKeyId, + _: &'a ModelRequestId, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + self.record("admission"); + Ok(true) + }) + } + + fn restore( + &self, + _: ClientAdmissionRecovery, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) }) + } +} + +impl ProviderCircuitPort for RecordingCoordination { + fn decision<'a>( + &'a self, + _: &'a ProviderKind, + ) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(ProviderCircuitDecision::Allow) }) + } + + fn observe_failure<'a>( + &'a self, + _: &'a ProviderKind, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> { + Box::pin(async move { + self.record("circuit_failure"); + Ok(()) + }) + } + + fn observe_success<'a>( + &'a self, + _: &'a ProviderKind, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> { + Box::pin(async move { + self.record("circuit_success"); + Ok(()) + }) + } +} + +impl NativeContinuationPort for RecordingCoordination { + fn resolve<'a>( + &'a self, + _: &'a PreviousResponseId, + ) -> BoxFuture<'a, Result, NativeContinuationStoreError>> { + Box::pin(async { Ok(None) }) + } + + fn record<'a>( + &'a self, + _: NativeContinuationPin, + ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> { + Box::pin(async move { + self.record("continuation"); + Ok(()) + }) + } +} + +#[tokio::test] +async fn full_redis_coordination_queues_should_drop_writes_without_waiting() { + let inner = Arc::new(RecordingCoordination::default()); + let (admissions, _admission_writer) = BufferedClientAdmissionPort::with_capacity( + inner.clone(), + NonZeroUsize::new(1).expect("capacity"), + ); + let (circuits, _circuit_writer) = BufferedProviderCircuitPort::with_capacity( + inner.clone(), + NonZeroUsize::new(1).expect("capacity"), + ); + let (continuation, _continuation_writer) = BufferedNativeContinuationPort::with_capacity( + inner.clone(), + NonZeroUsize::new(1).expect("capacity"), + ); + let client = ClientApiKeyId::new("key_buffer_test").expect("client key"); + let request = ModelRequestId::new("req_buffer_test").expect("request ID"); + let provider = ProviderKind::new("openai").expect("provider"); + + tokio::time::timeout(Duration::from_millis(50), async { + admissions + .release(&client, &request) + .await + .expect("first admission enqueue"); + admissions + .release(&client, &request) + .await + .expect("full admission queue remains fail-open"); + circuits + .observe_success(&provider) + .await + .expect("first circuit enqueue"); + circuits + .observe_failure(&provider) + .await + .expect("full circuit queue remains fail-open"); + continuation + .record(continuation_pin()) + .await + .expect("first continuation enqueue"); + continuation + .record(continuation_pin()) + .await + .expect("full continuation queue remains fail-open"); + }) + .await + .expect("coordination enqueue must never wait for Redis"); + + assert!(inner.operations().is_empty()); +} + +#[tokio::test] +async fn redis_coordination_writers_should_flush_each_side_effect() { + let inner = Arc::new(RecordingCoordination::default()); + let (admissions, admission_writer) = BufferedClientAdmissionPort::with_capacity( + inner.clone(), + NonZeroUsize::new(8).expect("capacity"), + ); + let (circuits, circuit_writer) = BufferedProviderCircuitPort::with_capacity( + inner.clone(), + NonZeroUsize::new(8).expect("capacity"), + ); + let (continuation, continuation_writer) = BufferedNativeContinuationPort::with_capacity( + inner.clone(), + NonZeroUsize::new(8).expect("capacity"), + ); + let client = ClientApiKeyId::new("key_writer_test").expect("client key"); + let request = ModelRequestId::new("req_writer_test").expect("request ID"); + let provider = ProviderKind::new("openai").expect("provider"); + admissions + .release(&client, &request) + .await + .expect("enqueue admission release"); + circuits + .observe_success(&provider) + .await + .expect("enqueue circuit feedback"); + continuation + .record(continuation_pin()) + .await + .expect("enqueue continuation"); + + let cancellation = CancellationToken::new(); + let tasks = [ + spawn_writer(Arc::new(admission_writer), cancellation.clone()), + spawn_writer(Arc::new(circuit_writer), cancellation.clone()), + spawn_writer(Arc::new(continuation_writer), cancellation.clone()), + ]; + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if inner.operations().len() == 3 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("writers should drain queued coordination writes"); + cancellation.cancel(); + for task in tasks { + task.await + .expect("coordination writer task") + .expect("coordination writer cancellation"); + } + + let mut operations = inner.operations(); + operations.sort_unstable(); + assert_eq!(operations, ["admission", "circuit_success", "continuation"]); +} + +fn spawn_writer( + writer: Arc, + cancellation: CancellationToken, +) -> tokio::task::JoinHandle> +where + T: gateway_core::task::DaemonTask + 'static, +{ + tokio::spawn(async move { writer.run(cancellation).await }) +} + +fn continuation_pin() -> NativeContinuationPin { + NativeContinuationPin::new( + PreviousResponseId::new("resp_buffer_test"), + PreviousResponseId::new("resp_buffer_test"), + ProviderKind::new("openai").expect("provider"), + ProviderAccountId::new("acct_buffer_test").expect("account"), + ) +} diff --git a/backend/crates/gateway-store/tests/execution_buffer.rs b/backend/crates/gateway-store/tests/execution_buffer.rs new file mode 100644 index 000000000..c56c8fd48 --- /dev/null +++ b/backend/crates/gateway-store/tests/execution_buffer.rs @@ -0,0 +1,156 @@ +use std::num::NonZeroUsize; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime}; + +use async_trait::async_trait; +use gateway_core::engine::{ + AttemptRecord, CancellationToken, ExecutionStore, IntermediateFailure, + ModelRequestFinalization, ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport, + UpstreamSendState, +}; +use gateway_core::error::StoreError; +use gateway_core::task::DaemonTask as _; +use gateway_store::postgres::BufferedExecutionStore; + +#[derive(Default)] +struct RecordingStore { + operations: Mutex>, +} + +impl RecordingStore { + fn record(&self, operation: &'static str) { + self.operations + .lock() + .expect("operations lock") + .push(operation); + } +} + +#[async_trait] +impl ExecutionStore for RecordingStore { + async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { + self.record("create"); + Ok(()) + } + + async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> { + self.record("attempt"); + Ok(()) + } + + async fn mark_send_state( + &self, + _: &ModelRequestId, + _: UpstreamSendState, + ) -> Result<(), StoreError> { + self.record("send"); + Ok(()) + } + + async fn mark_downstream_committed( + &self, + _: &ModelRequestId, + _: SystemTime, + _: Option, + ) -> Result<(), StoreError> { + self.record("commit"); + Ok(()) + } + + async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> { + self.record("status"); + Ok(()) + } + + async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> { + self.record("intermediate_failure"); + Ok(()) + } + + async fn record_probe_failure(&self, _: ProbeFailure) -> Result<(), StoreError> { + self.record("probe_failure"); + Ok(()) + } + + async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> { + self.record("finalize"); + Ok(()) + } + + async fn recover_expired(&self, _: SystemTime) -> Result { + self.record("recover"); + Ok(RecoveryReport::default()) + } +} + +#[tokio::test] +async fn full_observation_queue_never_waits_for_the_database() { + let inner = Arc::new(RecordingStore::default()); + let (store, _writer) = BufferedExecutionStore::with_capacity( + Arc::clone(&inner), + NonZeroUsize::new(1).expect("capacity"), + ); + let request_id = ModelRequestId::new("req_queue_full").expect("request id"); + + tokio::time::timeout( + Duration::from_millis(50), + store.mark_send_state(&request_id, UpstreamSendState::Sent), + ) + .await + .expect("first enqueue must not wait") + .expect("first enqueue is fail-open"); + tokio::time::timeout( + Duration::from_millis(50), + store.record_client_status(&request_id, 200), + ) + .await + .expect("full queue must not wait") + .expect("full queue is fail-open"); + + assert!(inner.operations.lock().expect("operations lock").is_empty()); +} + +#[tokio::test] +async fn observation_writer_persists_commands_in_enqueue_order() { + let inner = Arc::new(RecordingStore::default()); + let (store, writer) = BufferedExecutionStore::with_capacity( + Arc::clone(&inner), + NonZeroUsize::new(8).expect("capacity"), + ); + let request_id = ModelRequestId::new("req_queue_order").expect("request id"); + store + .mark_send_state(&request_id, UpstreamSendState::Sent) + .await + .expect("enqueue send state"); + store + .record_client_status(&request_id, 200) + .await + .expect("enqueue client status"); + + let cancellation = CancellationToken::new(); + let writer = Arc::new(writer); + let task = tokio::spawn({ + let writer = Arc::clone(&writer); + let cancellation = cancellation.clone(); + async move { writer.run(cancellation).await } + }); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if inner.operations.lock().expect("operations lock").len() == 2 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("writer should drain queued observations"); + cancellation.cancel(); + task.await + .expect("writer task") + .expect("writer cancellation"); + + assert_eq!( + *inner.operations.lock().expect("operations lock"), + ["send", "status"] + ); +} diff --git a/backend/crates/gateway-store/tests/main.rs b/backend/crates/gateway-store/tests/main.rs index 3ec752bac..f6799a403 100644 --- a/backend/crates/gateway-store/tests/main.rs +++ b/backend/crates/gateway-store/tests/main.rs @@ -1,3 +1,5 @@ +mod coordination_buffer; +mod execution_buffer; mod postgres; mod redis; mod support;