mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user