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 <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-29 00:11:26 +08:00
co-authored by Codex
parent 0c90ed30fb
commit cc42e78334
8 changed files with 1118 additions and 20 deletions
+90 -20
View File
@@ -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<StoreBundle> {
}),
);
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<dyn gateway_core::engine::admission::ClientAdmissionPort> = Arc::new(
redis::RedisClientAdmissionRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?,
);
let circuits: Arc<dyn gateway_core::engine::execution::ProviderCircuitPort> =
Arc::new(redis::RedisProviderCircuitRepository::new(
redis_connection.clone(),
REDIS_NAMESPACE,
redis::ProviderCircuitPolicy::default(),
)?),
)?);
let continuation: Arc<dyn gateway_core::engine::continuation::NativeContinuationPort> =
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<StoreBundle> {
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<StoreBundle> {
fn store_worker_contributions(
execution: Arc<postgres::PgExecutionStore>,
execution_writer: postgres::ExecutionObservationWriter<postgres::PgExecutionStore>,
admission_release_writer: redis::ClientAdmissionReleaseWriter,
circuit_feedback_writer: redis::ProviderCircuitFeedbackWriter,
continuation_writer: redis::NativeContinuationWriter,
retention: Arc<postgres::PgRetentionRepository>,
) -> StoreResult<Vec<WorkerContribution>> {
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)?,
),
])
}
@@ -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<S: ?Sized> {
inner: Arc<S>,
sender: mpsc::Sender<ExecutionObservationWrite>,
}
impl<S: ?Sized> BufferedExecutionStore<S> {
#[must_use]
pub fn new(inner: Arc<S>) -> (Self, ExecutionObservationWriter<S>) {
Self::with_capacity(
inner,
NonZeroUsize::new(DEFAULT_QUEUE_CAPACITY).expect("queue capacity is non-zero"),
)
}
#[must_use]
pub fn with_capacity(
inner: Arc<S>,
capacity: NonZeroUsize,
) -> (Self, ExecutionObservationWriter<S>) {
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<S> ExecutionStore for BufferedExecutionStore<S>
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<u16>,
) -> 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<RecoveryReport, StoreError> {
self.inner.recover_expired(now).await
}
}
/// 由 Host 监督的单消费者,保证同一请求的观测命令按入队顺序落库。
pub struct ExecutionObservationWriter<S: ?Sized> {
inner: Arc<S>,
receiver: Mutex<mpsc::Receiver<ExecutionObservationWrite>>,
}
impl<S> DaemonTask for ExecutionObservationWriter<S>
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<NewModelRequest>),
Attempt(Box<AttemptRecord>),
CreateWithAttempt(Box<(NewModelRequest, AttemptRecord)>),
MarkSendState {
request_id: ModelRequestId,
state: UpstreamSendState,
},
MarkDownstreamCommitted {
request_id: ModelRequestId,
committed_at: SystemTime,
client_status_code: Option<u16>,
},
RecordClientStatus {
request_id: ModelRequestId,
client_status_code: u16,
},
IntermediateFailure(Box<IntermediateFailure>),
ProbeFailure(Box<ProbeFailure>),
Finalize(Box<ModelRequestFinalization>),
}
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<S>(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,
}
}
}
@@ -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::*;
@@ -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<dyn ClientAdmissionPort>,
sender: mpsc::Sender<AdmissionRelease>,
}
impl BufferedClientAdmissionPort {
#[must_use]
pub fn new(inner: Arc<dyn ClientAdmissionPort>) -> (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<dyn ClientAdmissionPort>,
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<ClientAdmissionDecision, ClientAdmissionError>> {
self.inner.admit(request)
}
fn release<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
model_request_id: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
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<ClientAdmissionRestoreResult, ClientAdmissionError>> {
self.inner.restore(recovery)
}
}
struct AdmissionRelease {
client_api_key_id: ClientApiKeyId,
model_request_id: ModelRequestId,
}
pub struct ClientAdmissionReleaseWriter {
inner: Arc<dyn ClientAdmissionPort>,
receiver: Mutex<mpsc::Receiver<AdmissionRelease>>,
}
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<dyn ProviderCircuitPort>,
sender: mpsc::Sender<CircuitFeedback>,
}
impl BufferedProviderCircuitPort {
#[must_use]
pub fn new(inner: Arc<dyn ProviderCircuitPort>) -> (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<dyn ProviderCircuitPort>,
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<ProviderCircuitDecision, ProviderCircuitError>> {
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<dyn ProviderCircuitPort>,
receiver: Mutex<mpsc::Receiver<CircuitFeedback>>,
}
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<dyn NativeContinuationPort>,
sender: mpsc::Sender<NativeContinuationPin>,
}
impl BufferedNativeContinuationPort {
#[must_use]
pub fn new(inner: Arc<dyn NativeContinuationPort>) -> (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<dyn NativeContinuationPort>,
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<Option<NativeContinuationPin>, 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<dyn NativeContinuationPort>,
receiver: Mutex<mpsc::Receiver<NativeContinuationPin>>,
}
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 后台写入失败");
}
}
})
}
}
@@ -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::*;
@@ -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<Vec<&'static str>>,
}
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<ClientAdmissionDecision, ClientAdmissionError>> {
Box::pin(async { Ok(ClientAdmissionDecision::Granted) })
}
fn release<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async move {
self.record("admission");
Ok(true)
})
}
fn restore(
&self,
_: ClientAdmissionRecovery,
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
}
}
impl ProviderCircuitPort for RecordingCoordination {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
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<Option<NativeContinuationPin>, 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<T>(
writer: Arc<T>,
cancellation: CancellationToken,
) -> tokio::task::JoinHandle<Result<(), gateway_core::task::WorkerTaskError>>
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"),
)
}
@@ -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<Vec<&'static str>>,
}
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<u16>,
) -> 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<RecoveryReport, StoreError> {
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"]
);
}
@@ -1,3 +1,5 @@
mod coordination_buffer;
mod execution_buffer;
mod postgres;
mod redis;
mod support;