mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix(workers): recover daemon health after successful work (#499)
Report per-run daemon progress to reset health failures and restart backoff after successful work, with stale-run isolation and recovery regression tests. Co-authored-by: Zyy <zyycn@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -11,7 +11,7 @@ use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::{
|
||||
lifecycle::CancellationToken,
|
||||
task::{DaemonTask, WorkerTaskError},
|
||||
task::{DaemonTask, WorkerDaemonContext, WorkerTaskError},
|
||||
};
|
||||
use tracing::{info, warn};
|
||||
|
||||
@@ -75,18 +75,20 @@ impl BackupTask {
|
||||
}
|
||||
|
||||
impl DaemonTask for BackupTask {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move { self.run_daemon(cancellation).await })
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move { self.run_daemon(context).await })
|
||||
}
|
||||
}
|
||||
|
||||
impl BackupTask {
|
||||
async fn run_daemon(&self, cancellation: CancellationToken) -> Result<(), WorkerTaskError> {
|
||||
async fn run_daemon(&self, context: WorkerDaemonContext) -> Result<(), WorkerTaskError> {
|
||||
let cancellation = context.cancellation();
|
||||
loop {
|
||||
if cancellation.is_cancelled() {
|
||||
return Ok(());
|
||||
}
|
||||
self.run_cycle(&cancellation).await?;
|
||||
self.run_cycle(cancellation).await?;
|
||||
context.succeeded();
|
||||
|
||||
tokio::select! {
|
||||
_ = cancellation.cancelled() => return Ok(()),
|
||||
|
||||
@@ -9,10 +9,7 @@ use std::{
|
||||
|
||||
use chrono::Utc;
|
||||
use futures::{FutureExt as _, StreamExt as _, future::BoxFuture, stream::FuturesUnordered};
|
||||
use gateway_core::{
|
||||
lifecycle::CancellationToken,
|
||||
task::{DaemonTask, WorkerTaskError},
|
||||
};
|
||||
use gateway_core::task::{DaemonTask, WorkerDaemonContext, WorkerTaskError};
|
||||
use tokio::{sync::Notify, time::Instant};
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -332,8 +329,9 @@ impl ImportTasksService for DefaultImportTasksService {
|
||||
pub(crate) struct ImportTaskWorker(pub Arc<DefaultImportTasksService>);
|
||||
|
||||
impl DaemonTask for ImportTaskWorker {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
let mut active = FuturesUnordered::new();
|
||||
// 固定一个取消 future,避免循环反复注册中立 CancellationToken 的 waiter
|
||||
let cancelled = cancellation.cancelled();
|
||||
@@ -354,6 +352,7 @@ impl DaemonTask for ImportTaskWorker {
|
||||
};
|
||||
active.push(self.0.execute(id, index, input));
|
||||
}
|
||||
context.succeeded();
|
||||
tokio::select! {
|
||||
() = &mut cancelled => {},
|
||||
_ = active.next(), if !active.is_empty() => {},
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
//! 供 backup::task 与 use_case::backup 的测试复用
|
||||
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
@@ -92,6 +93,7 @@ pub(crate) fn queued_record(seed: &BackupRecordSeed) -> BackupRecord {
|
||||
/// 内存版备份仓储
|
||||
pub(crate) struct FakeBackupRepository {
|
||||
settings: Mutex<BackupSettings>,
|
||||
settings_unavailable: AtomicBool,
|
||||
records: Mutex<Vec<BackupRecord>>,
|
||||
}
|
||||
|
||||
@@ -99,10 +101,16 @@ impl FakeBackupRepository {
|
||||
pub(crate) fn new(settings: BackupSettings) -> Self {
|
||||
Self {
|
||||
settings: Mutex::new(settings),
|
||||
settings_unavailable: AtomicBool::new(false),
|
||||
records: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn set_settings_unavailable(&self, unavailable: bool) {
|
||||
self.settings_unavailable
|
||||
.store(unavailable, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// 测试辅助:读取当前全部记录
|
||||
pub(crate) fn all_records(&self) -> Vec<BackupRecord> {
|
||||
self.records.lock().expect("records").clone()
|
||||
@@ -134,6 +142,13 @@ impl FakeBackupRepository {
|
||||
#[async_trait]
|
||||
impl BackupRepository for FakeBackupRepository {
|
||||
async fn load_settings(&self) -> AdminStoreResult<BackupSettings> {
|
||||
if self.settings_unavailable.load(Ordering::SeqCst) {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
"backup",
|
||||
"test store unavailable",
|
||||
));
|
||||
}
|
||||
Ok(self.settings.lock().expect("settings").clone())
|
||||
}
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ use gateway_admin::{
|
||||
ports::backup::{BackupRepository, DatabaseDumpPort, DumpRequest},
|
||||
};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::task::{DaemonTask, WorkerDaemonContext};
|
||||
|
||||
use super::support::{
|
||||
DUMP_CONTENT, FakeBackupRepository, FakeDumpPort, FakeObjectStore, backup_id,
|
||||
@@ -379,3 +380,40 @@ async fn manual_backup_preserves_primary_and_final_delete_failures() {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn daemon_reports_success_only_after_a_complete_cycle_and_recovers_after_store_failure() {
|
||||
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
|
||||
let task = Arc::new(BackupTask::new(
|
||||
repository.clone(),
|
||||
Arc::new(FakeDumpPort::new()),
|
||||
Arc::new(FakeObjectStore::new()),
|
||||
));
|
||||
let context = WorkerDaemonContext::new(CancellationToken::new());
|
||||
repository.set_settings_unavailable(true);
|
||||
let error = task.run(context.clone()).await.unwrap_err();
|
||||
assert_eq!(error.as_safe_str(), "backup store failed");
|
||||
assert_eq!(context.completed_cycles(), 0);
|
||||
|
||||
repository.set_settings_unavailable(false);
|
||||
let recovered = WorkerDaemonContext::new(CancellationToken::new());
|
||||
let worker = tokio::spawn({
|
||||
let task = task.clone();
|
||||
let context = recovered.clone();
|
||||
async move { task.run(context).await }
|
||||
});
|
||||
recovered.progress_since(0).await;
|
||||
assert_eq!(recovered.completed_cycles(), 1);
|
||||
repository.set_settings_unavailable(true);
|
||||
tokio::time::advance(std::time::Duration::from_secs(30)).await;
|
||||
assert_eq!(
|
||||
worker.await.unwrap().unwrap_err().as_safe_str(),
|
||||
"backup store failed"
|
||||
);
|
||||
assert_eq!(recovered.completed_cycles(), 1);
|
||||
|
||||
let cancelled = WorkerDaemonContext::new(CancellationToken::new());
|
||||
cancelled.cancellation().cancel();
|
||||
task.run(cancelled.clone()).await.unwrap();
|
||||
assert_eq!(cancelled.completed_cycles(), 0);
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@ use gateway_admin::{
|
||||
use gateway_core::{
|
||||
lifecycle::CancellationToken,
|
||||
routing::ProviderKind,
|
||||
task::{WorkerContribution, WorkerKind, WorkerRunnable},
|
||||
task::{WorkerContribution, WorkerDaemonContext, WorkerKind, WorkerRunnable},
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -75,7 +75,9 @@ fn start(bundle: &mut AdminBundle) -> (CancellationToken, tokio::task::JoinHandl
|
||||
let cancellation = CancellationToken::new();
|
||||
let token = cancellation.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
task.run(token).await.expect("import daemon exits cleanly");
|
||||
task.run(WorkerDaemonContext::new(token))
|
||||
.await
|
||||
.expect("import daemon exits cleanly");
|
||||
});
|
||||
(cancellation, handle)
|
||||
}
|
||||
|
||||
@@ -372,8 +372,12 @@ struct RuntimeSnapshotSubscriptionTask {
|
||||
}
|
||||
|
||||
impl DaemonTask for RuntimeSnapshotSubscriptionTask {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(
|
||||
&self,
|
||||
context: crate::task::WorkerDaemonContext,
|
||||
) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
let mut retry_delay = INITIAL_BACKOFF;
|
||||
loop {
|
||||
if cancellation.is_cancelled() {
|
||||
@@ -383,10 +387,11 @@ impl DaemonTask for RuntimeSnapshotSubscriptionTask {
|
||||
let mut subscription = match subscription {
|
||||
Ok(subscription) => {
|
||||
retry_delay = INITIAL_BACKOFF;
|
||||
context.succeeded();
|
||||
subscription
|
||||
}
|
||||
Err(_) => {
|
||||
wait_or_cancel(&cancellation, retry_delay).await;
|
||||
wait_or_cancel(cancellation, retry_delay).await;
|
||||
retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF);
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -6,8 +6,11 @@
|
||||
|
||||
use std::fmt;
|
||||
use std::num::NonZeroU64;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use event_listener::Event;
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::lifecycle::CancellationToken;
|
||||
@@ -341,11 +344,60 @@ pub trait ScheduledTask: Send + Sync {
|
||||
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>>;
|
||||
}
|
||||
|
||||
/// 一次 daemon 运行的取消与成功进度,重启时必须创建新上下文
|
||||
#[derive(Clone)]
|
||||
pub struct WorkerDaemonContext {
|
||||
cancellation: CancellationToken,
|
||||
progress: Arc<DaemonProgress>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct DaemonProgress {
|
||||
completed: AtomicU64,
|
||||
changed: Event,
|
||||
}
|
||||
|
||||
impl WorkerDaemonContext {
|
||||
#[must_use]
|
||||
pub fn new(cancellation: CancellationToken) -> Self {
|
||||
Self {
|
||||
cancellation,
|
||||
progress: Arc::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn cancellation(&self) -> &CancellationToken {
|
||||
&self.cancellation
|
||||
}
|
||||
|
||||
/// 仅在实际工作或依赖检查成功后通知 Host;进入 run 不代表恢复
|
||||
pub fn succeeded(&self) {
|
||||
if !self.cancellation.is_cancelled() {
|
||||
self.progress.completed.fetch_add(1, Ordering::Release);
|
||||
self.progress.changed.notify(usize::MAX);
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn completed_cycles(&self) -> u64 {
|
||||
self.progress.completed.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
/// 先注册再检查计数,避免报告与等待交错时丢失成功通知
|
||||
pub async fn progress_since(&self, completed: u64) {
|
||||
let changed = self.progress.changed.listen();
|
||||
if self.completed_cycles() == completed {
|
||||
changed.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 长驻任务
|
||||
///
|
||||
/// 实现负责连接级重连退避;Host 负责任务退出或 panic 后的重启退避
|
||||
pub trait DaemonTask: Send + Sync {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>>;
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>>;
|
||||
}
|
||||
|
||||
/// Host 可执行的两种任务形态,无法组合出“守护任务 + 周期调度”等非法状态
|
||||
|
||||
@@ -28,7 +28,8 @@ use gateway_core::runtime::{
|
||||
};
|
||||
use gateway_core::settings::SettingsValues;
|
||||
use gateway_core::task::{
|
||||
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerKind, WorkerRunnable,
|
||||
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDaemonContext, WorkerKind,
|
||||
WorkerRunnable,
|
||||
};
|
||||
|
||||
type ReadGate<T> = Arc<Mutex<Option<oneshot::Receiver<Result<T, SnapshotStoreError>>>>>;
|
||||
@@ -490,7 +491,7 @@ fn subscription_refresh_should_wait_for_inflight_compile_and_reload_authoritativ
|
||||
})
|
||||
.expect("subscription task");
|
||||
let cancellation = CancellationToken::new();
|
||||
let mut daemon = task.run(cancellation.clone());
|
||||
let mut daemon = task.run(WorkerDaemonContext::new(cancellation.clone()));
|
||||
let mut old = Box::pin(publisher.refresh());
|
||||
assert!(futures::poll!(old.as_mut()).is_pending());
|
||||
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
|
||||
|
||||
@@ -14,9 +14,9 @@ use gateway_core::health::{
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::task::{
|
||||
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
|
||||
WorkerDefinitionError, WorkerDisabledReason, WorkerId, WorkerKind, WorkerLeaderLeaseGuard,
|
||||
WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseRequest, WorkerRegistration,
|
||||
WorkerRunnable, WorkerSchedule, WorkerTaskError,
|
||||
WorkerDaemonContext, WorkerDefinitionError, WorkerDisabledReason, WorkerId, WorkerKind,
|
||||
WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseRequest,
|
||||
WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError,
|
||||
};
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
@@ -235,7 +235,7 @@ pub enum WorkerStartError {
|
||||
#[derive(Clone)]
|
||||
struct HealthEntry {
|
||||
snapshot: WorkerHealthSnapshot,
|
||||
success_freshness: Duration,
|
||||
success_freshness: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -253,11 +253,14 @@ impl WorkerHealthRegistry {
|
||||
.values()
|
||||
.map(|registration| {
|
||||
let freshness = match ®istration.runnable {
|
||||
WorkerRunnable::Scheduled { schedule, .. } => (*schedule)
|
||||
.interval()
|
||||
.saturating_mul(2)
|
||||
.saturating_add(schedule.maximum_backoff()),
|
||||
WorkerRunnable::Daemon { restart, .. } => restart.maximum_backoff(),
|
||||
WorkerRunnable::Scheduled { schedule, .. } => Some(
|
||||
(*schedule)
|
||||
.interval()
|
||||
.saturating_mul(2)
|
||||
.saturating_add(schedule.maximum_backoff()),
|
||||
),
|
||||
// daemon 可长期空闲或执行备份,重启退避不约束工作耗时
|
||||
WorkerRunnable::Daemon { .. } => None,
|
||||
};
|
||||
let key = WorkerHealthKey::Task(registration.id.clone());
|
||||
(
|
||||
@@ -275,7 +278,7 @@ impl WorkerHealthRegistry {
|
||||
key.clone(),
|
||||
HealthEntry {
|
||||
snapshot: empty_snapshot(key, WorkerRuntimeState::Disabled),
|
||||
success_freshness: Duration::ZERO,
|
||||
success_freshness: None,
|
||||
},
|
||||
)
|
||||
}));
|
||||
@@ -315,11 +318,11 @@ impl WorkerHealthRegistry {
|
||||
});
|
||||
}
|
||||
|
||||
fn succeeded(&self, id: &WorkerId) {
|
||||
fn succeeded(&self, id: &WorkerId, runtime_state: WorkerRuntimeState, cycles: u64) {
|
||||
self.update(id, |state| {
|
||||
state.state = WorkerRuntimeState::Idle;
|
||||
state.state = runtime_state;
|
||||
state.consecutive_failures = 0;
|
||||
state.completed_cycles = state.completed_cycles.saturating_add(1);
|
||||
state.completed_cycles = state.completed_cycles.saturating_add(cycles);
|
||||
state.last_success_at = Some(SystemTime::now());
|
||||
state.last_error = None;
|
||||
});
|
||||
@@ -361,8 +364,11 @@ impl WorkerHealthSource for WorkerHealthRegistry {
|
||||
snapshot.state,
|
||||
WorkerRuntimeState::Running | WorkerRuntimeState::Idle
|
||||
) && snapshot.last_success_at.is_some_and(|last_success| {
|
||||
now.duration_since(last_success)
|
||||
.is_ok_and(|age| age > entry.success_freshness)
|
||||
now.duration_since(last_success).is_ok_and(|age| {
|
||||
entry
|
||||
.success_freshness
|
||||
.is_some_and(|freshness| age > freshness)
|
||||
})
|
||||
});
|
||||
if stale {
|
||||
snapshot.state = WorkerRuntimeState::BackingOff;
|
||||
@@ -407,7 +413,7 @@ async fn supervise_scheduled(
|
||||
|
||||
let delay = match outcome {
|
||||
ScheduledOutcome::Succeeded => {
|
||||
health.succeeded(&id);
|
||||
health.succeeded(&id, WorkerRuntimeState::Idle, 1);
|
||||
backoff = schedule.initial_backoff();
|
||||
schedule.interval()
|
||||
}
|
||||
@@ -595,17 +601,33 @@ async fn supervise_daemon(
|
||||
break;
|
||||
}
|
||||
health.running(&id, None);
|
||||
let daemon_cancel = CancellationToken::new();
|
||||
let run = AssertUnwindSafe(task.run(daemon_cancel.clone())).catch_unwind();
|
||||
let context = WorkerDaemonContext::new(CancellationToken::new());
|
||||
let run = AssertUnwindSafe(task.run(context.clone())).catch_unwind();
|
||||
tokio::pin!(run);
|
||||
let result = tokio::select! {
|
||||
result = &mut run => Some(result),
|
||||
() = cancellation.cancelled() => {
|
||||
daemon_cancel.cancel();
|
||||
let _ = run.await;
|
||||
None
|
||||
let mut completed = 0;
|
||||
let result = loop {
|
||||
let outcome = tokio::select! {
|
||||
biased;
|
||||
() = cancellation.cancelled() => {
|
||||
context.cancellation().cancel();
|
||||
let _ = run.await;
|
||||
break None;
|
||||
}
|
||||
result = &mut run => Some(result),
|
||||
() = context.progress_since(completed) => None,
|
||||
};
|
||||
// 成功与退出可在同一 poll 内发生,先消费进度再记录本次退出故障
|
||||
let cycles = context.completed_cycles();
|
||||
if cycles > completed {
|
||||
health.succeeded(&id, WorkerRuntimeState::Running, cycles - completed);
|
||||
completed = cycles;
|
||||
backoff = restart.initial_backoff();
|
||||
}
|
||||
if let Some(result) = outcome {
|
||||
break Some(result);
|
||||
}
|
||||
};
|
||||
context.cancellation().cancel();
|
||||
let Some(result) = result else { break };
|
||||
let error = match result {
|
||||
Ok(Ok(())) => WorkerTaskError::safe("daemon exited unexpectedly"),
|
||||
|
||||
@@ -12,7 +12,7 @@ use gateway_admin::ports::system::{
|
||||
use gateway_core::diagnostics::{OperationalDiagnostics, OperationalFailure};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::task::{
|
||||
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerId, WorkerKind,
|
||||
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDaemonContext, WorkerId, WorkerKind,
|
||||
WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest,
|
||||
WorkerRegistration, WorkerRunnable, WorkerTaskError,
|
||||
};
|
||||
@@ -365,8 +365,9 @@ impl Drop for TaskLifetime {
|
||||
}
|
||||
|
||||
impl DaemonTask for WriterTask {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
let _lifetime = TaskLifetime(self.probe.dropped.clone());
|
||||
let mut received = self.received.lock().await;
|
||||
self.probe.started.cancel();
|
||||
@@ -394,8 +395,9 @@ impl DaemonTask for WriterTask {
|
||||
struct IdleTask;
|
||||
|
||||
impl DaemonTask for IdleTask {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
cancellation.cancelled().await;
|
||||
Ok(())
|
||||
})
|
||||
|
||||
@@ -11,10 +11,10 @@ use futures::future::BoxFuture;
|
||||
use gateway_core::health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerRuntimeState};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::task::{
|
||||
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDisabledReason,
|
||||
WorkerFencingToken, WorkerId, WorkerKind, WorkerLeaderLeaseGuard, WorkerLeaderLeasePort,
|
||||
WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest, WorkerRegistration,
|
||||
WorkerRunnable, WorkerSchedule, WorkerTaskError,
|
||||
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
|
||||
WorkerDaemonContext, WorkerDisabledReason, WorkerFencingToken, WorkerId, WorkerKind,
|
||||
WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseError,
|
||||
WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError,
|
||||
};
|
||||
use gateway_host::workers::{WorkerStartError, WorkerSupervisor};
|
||||
use tokio::sync::Notify;
|
||||
@@ -1023,3 +1023,187 @@ async fn worker_boundary_records_one_native_failure_and_recovers_health() {
|
||||
.contains("PRIVATE_WORKER_NATIVE_CAUSE")
|
||||
);
|
||||
}
|
||||
|
||||
// 每次运行保存自己的上下文,测试可模拟已退出任务的迟到通知
|
||||
struct ControlledDaemon {
|
||||
contexts: Arc<Mutex<Vec<WorkerDaemonContext>>>,
|
||||
commands: tokio::sync::Mutex<tokio::sync::mpsc::UnboundedReceiver<DaemonCommand>>,
|
||||
}
|
||||
|
||||
enum DaemonCommand {
|
||||
Succeed,
|
||||
Fail,
|
||||
SucceedThenFail,
|
||||
Panic,
|
||||
Exit,
|
||||
}
|
||||
|
||||
impl DaemonTask for ControlledDaemon {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
lock_unpoisoned(&self.contexts).push(context.clone());
|
||||
let mut commands = self.commands.lock().await;
|
||||
loop {
|
||||
let command = tokio::select! {
|
||||
() = context.cancellation().cancelled() => return Ok(()),
|
||||
command = commands.recv() => command.expect("daemon command"),
|
||||
};
|
||||
match command {
|
||||
DaemonCommand::Succeed => context.succeeded(),
|
||||
DaemonCommand::SucceedThenFail => {
|
||||
context.succeeded();
|
||||
return Err(WorkerTaskError::safe("test daemon failure"));
|
||||
}
|
||||
DaemonCommand::Fail => {
|
||||
return Err(WorkerTaskError::safe("test daemon failure"));
|
||||
}
|
||||
DaemonCommand::Panic => panic!("test daemon panic"),
|
||||
DaemonCommand::Exit => return Ok(()),
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn daemon_plan(task: impl DaemonTask + 'static) -> Vec<WorkerContribution> {
|
||||
complete_plan(vec![WorkerContribution::Registration(
|
||||
WorkerRegistration::try_new(
|
||||
WorkerId::try_new(WorkerKind::Backup, "backup").unwrap(),
|
||||
WorkerRunnable::Daemon {
|
||||
restart: DaemonRestartPolicy::try_new(
|
||||
Duration::from_millis(10),
|
||||
Duration::from_millis(20),
|
||||
)
|
||||
.unwrap(),
|
||||
task: Box::new(task),
|
||||
},
|
||||
)
|
||||
.unwrap(),
|
||||
)])
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn daemon_recovers_only_after_current_run_succeeds_and_resets_backoff() {
|
||||
let (commands, receiver) = tokio::sync::mpsc::unbounded_channel();
|
||||
let contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
let supervisor = supervisor();
|
||||
let id = WorkerId::try_new(WorkerKind::Backup, "backup").unwrap();
|
||||
supervisor
|
||||
.start(
|
||||
daemon_plan(ControlledDaemon {
|
||||
contexts: contexts.clone(),
|
||||
commands: tokio::sync::Mutex::new(receiver),
|
||||
}),
|
||||
Arc::new(FakeLeasePort::default()),
|
||||
Arc::new(RecordingDiagnostics::default()),
|
||||
)
|
||||
.unwrap();
|
||||
yield_until(|| lock_unpoisoned(&contexts).len() == 1).await;
|
||||
|
||||
// 连续失败逐步退避;仅重新进入 run 不能恢复健康
|
||||
for (attempt, delay) in [(1, 10), (2, 20)] {
|
||||
commands.send(DaemonCommand::Fail).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
|
||||
assert_eq!(task_health(&supervisor, &id).consecutive_failures, attempt);
|
||||
tokio::time::advance(Duration::from_millis(delay - 1)).await;
|
||||
assert_eq!(lock_unpoisoned(&contexts).len(), attempt as usize);
|
||||
tokio::time::advance(Duration::from_millis(1)).await;
|
||||
yield_until(|| lock_unpoisoned(&contexts).len() == attempt as usize + 1).await;
|
||||
assert_eq!(task_health(&supervisor, &id).consecutive_failures, attempt);
|
||||
}
|
||||
lock_unpoisoned(&contexts)[0].succeeded();
|
||||
tokio::task::yield_now().await;
|
||||
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 2);
|
||||
|
||||
commands.send(DaemonCommand::Succeed).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).consecutive_failures == 0).await;
|
||||
let recovered = task_health(&supervisor, &id);
|
||||
assert_eq!(recovered.state, WorkerRuntimeState::Running);
|
||||
assert_eq!(recovered.completed_cycles, 1);
|
||||
assert!(recovered.last_error.is_none());
|
||||
assert!(recovered.last_success_at.is_some());
|
||||
assert!(recovered.last_failure_at.is_some());
|
||||
|
||||
commands.send(DaemonCommand::Fail).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
|
||||
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 1);
|
||||
tokio::time::advance(Duration::from_millis(10)).await;
|
||||
yield_until(|| lock_unpoisoned(&contexts).len() == 4).await;
|
||||
commands.send(DaemonCommand::Fail).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
|
||||
tokio::time::advance(Duration::from_millis(20)).await;
|
||||
yield_until(|| lock_unpoisoned(&contexts).len() == 5).await;
|
||||
|
||||
// 同一 poll 中先成功再退出,必须保留最新故障并重新从初始退避开始
|
||||
commands.send(DaemonCommand::SucceedThenFail).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
|
||||
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 1);
|
||||
assert_eq!(task_health(&supervisor, &id).completed_cycles, 2);
|
||||
tokio::time::advance(Duration::from_millis(10)).await;
|
||||
yield_until(|| lock_unpoisoned(&contexts).len() == 6).await;
|
||||
supervisor.shutdown(Duration::from_secs(1)).await;
|
||||
lock_unpoisoned(&contexts)[5].succeeded();
|
||||
assert_eq!(
|
||||
task_health(&supervisor, &id).state,
|
||||
WorkerRuntimeState::Stopped
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn daemon_success_remains_healthy_while_idle_beyond_restart_backoff() {
|
||||
let (commands, receiver) = tokio::sync::mpsc::unbounded_channel();
|
||||
let supervisor = supervisor();
|
||||
let id = WorkerId::try_new(WorkerKind::Backup, "backup").unwrap();
|
||||
supervisor
|
||||
.start(
|
||||
daemon_plan(ControlledDaemon {
|
||||
contexts: Arc::default(),
|
||||
commands: tokio::sync::Mutex::new(receiver),
|
||||
}),
|
||||
Arc::new(FakeLeasePort::default()),
|
||||
Arc::new(RecordingDiagnostics::default()),
|
||||
)
|
||||
.unwrap();
|
||||
commands.send(DaemonCommand::Succeed).unwrap();
|
||||
wait_until(|| task_health(&supervisor, &id).completed_cycles == 1).await;
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
let health = task_health(&supervisor, &id);
|
||||
assert_eq!(health.state, WorkerRuntimeState::Running);
|
||||
assert_eq!(health.consecutive_failures, 0);
|
||||
supervisor.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn daemon_panic_and_unexpected_exit_remain_unhealthy_until_success() {
|
||||
for (command, error) in [
|
||||
(DaemonCommand::Panic, "daemon panicked"),
|
||||
(DaemonCommand::Exit, "daemon exited unexpectedly"),
|
||||
] {
|
||||
let (commands, receiver) = tokio::sync::mpsc::unbounded_channel();
|
||||
let contexts = Arc::new(Mutex::new(Vec::new()));
|
||||
let supervisor = supervisor();
|
||||
let id = WorkerId::try_new(WorkerKind::Backup, "backup").unwrap();
|
||||
supervisor
|
||||
.start(
|
||||
daemon_plan(ControlledDaemon {
|
||||
contexts: contexts.clone(),
|
||||
commands: tokio::sync::Mutex::new(receiver),
|
||||
}),
|
||||
Arc::new(FakeLeasePort::default()),
|
||||
Arc::new(RecordingDiagnostics::default()),
|
||||
)
|
||||
.unwrap();
|
||||
commands.send(command).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).state == WorkerRuntimeState::BackingOff).await;
|
||||
assert_eq!(
|
||||
task_health(&supervisor, &id).last_error.as_deref(),
|
||||
Some(error)
|
||||
);
|
||||
tokio::time::advance(Duration::from_millis(10)).await;
|
||||
yield_until(|| lock_unpoisoned(&contexts).len() == 2).await;
|
||||
assert_eq!(task_health(&supervisor, &id).consecutive_failures, 1);
|
||||
commands.send(DaemonCommand::Succeed).unwrap();
|
||||
yield_until(|| task_health(&supervisor, &id).consecutive_failures == 0).await;
|
||||
supervisor.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,8 +7,9 @@ use gateway_core::{
|
||||
lifecycle::CancellationToken,
|
||||
runtime::RuntimeSnapshotHandle,
|
||||
task::{
|
||||
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDefinitionError, WorkerId,
|
||||
WorkerKind, WorkerRegistration, WorkerRunnable, WorkerTaskError,
|
||||
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDaemonContext,
|
||||
WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable,
|
||||
WorkerTaskError,
|
||||
},
|
||||
};
|
||||
use gateway_plugin_sdk::Stage;
|
||||
@@ -52,8 +53,9 @@ impl PluginRuntime {
|
||||
}
|
||||
|
||||
impl DaemonTask for MaintenanceTask {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
let mut publications = self.snapshots.subscribe_publications();
|
||||
let mut active: Option<Arc<PreparedSet>> = None;
|
||||
let mut revision = None;
|
||||
@@ -91,6 +93,9 @@ impl DaemonTask for MaintenanceTask {
|
||||
}
|
||||
}
|
||||
revision = snapshot.as_ref().map(|s| s.revision());
|
||||
if snapshot.is_some() {
|
||||
context.succeeded();
|
||||
}
|
||||
// 发布变更优先于下一轮维护;不在快照发布锁中等待插件或数据库
|
||||
tokio::select! {
|
||||
biased;
|
||||
|
||||
@@ -16,7 +16,7 @@ use gateway_admin::{
|
||||
use gateway_core::{
|
||||
lifecycle::CancellationToken,
|
||||
policy::ClientApiKeyId,
|
||||
task::{WorkerContribution, WorkerRunnable},
|
||||
task::{WorkerContribution, WorkerDaemonContext, WorkerRunnable},
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::{sync::Notify, time::timeout};
|
||||
@@ -277,7 +277,11 @@ async fn published_maintenance_resets_budget_with_only_budget_permission() {
|
||||
};
|
||||
let stop = CancellationToken::new();
|
||||
let cancellation = stop.clone();
|
||||
let worker = tokio::spawn(async move { task.run(cancellation).await.unwrap() });
|
||||
let worker = tokio::spawn(async move {
|
||||
task.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
let result: Value = timeout(Duration::from_secs(15), async {
|
||||
loop {
|
||||
if let Ok(content) = std::fs::read_to_string(&marker)
|
||||
|
||||
@@ -9,6 +9,7 @@ use gateway_admin::{
|
||||
},
|
||||
ports::plugin_resources::PluginResourceAccess,
|
||||
};
|
||||
use gateway_core::task::WorkerDaemonContext;
|
||||
use gateway_core::{
|
||||
account::FastMode,
|
||||
lifecycle::CancellationToken,
|
||||
@@ -165,7 +166,9 @@ async fn start(
|
||||
let stop = CancellationToken::new();
|
||||
let cancellation = stop.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
task.run(cancellation).await.unwrap();
|
||||
task.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
(runtime, core, resources, stop, task)
|
||||
}
|
||||
|
||||
@@ -31,11 +31,10 @@ use gateway_core::{
|
||||
account::OpaqueProviderData,
|
||||
engine::budget::{ClientBudgetLimits, ClientBudgetStatus},
|
||||
engine::execution::ClientApiKeyUsageSink,
|
||||
lifecycle::CancellationToken,
|
||||
metering::Decimal,
|
||||
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
task::{DaemonTask, WorkerTaskError},
|
||||
task::{DaemonTask, WorkerDaemonContext, WorkerTaskError},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Transaction};
|
||||
@@ -626,8 +625,9 @@ impl ClientApiKeyUsageSink for PgClientApiKeyUsageSink {
|
||||
}
|
||||
|
||||
impl DaemonTask for PgClientApiKeyUsageWriter {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
loop {
|
||||
tokio::select! {
|
||||
() = cancellation.cancelled() => {
|
||||
@@ -649,6 +649,7 @@ impl DaemonTask for PgClientApiKeyUsageWriter {
|
||||
"client API key last-used flush failed",
|
||||
));
|
||||
}
|
||||
context.succeeded();
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -14,8 +14,7 @@ use gateway_core::engine::{
|
||||
ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport,
|
||||
};
|
||||
use gateway_core::error::{ProviderError, StoreError};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::task::{DaemonTask, WorkerTaskError};
|
||||
use gateway_core::task::{DaemonTask, WorkerDaemonContext, WorkerTaskError};
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use tokio::sync::{Mutex, Notify, mpsc};
|
||||
|
||||
@@ -486,7 +485,7 @@ where
|
||||
{
|
||||
fn run(
|
||||
&self,
|
||||
cancellation: CancellationToken,
|
||||
context: WorkerDaemonContext,
|
||||
) -> futures::future::BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let shutdown_deadline = OnceLock::new();
|
||||
@@ -496,7 +495,7 @@ where
|
||||
receiver,
|
||||
self.inner.as_ref(),
|
||||
self.state.as_ref(),
|
||||
cancellation.clone(),
|
||||
context.clone(),
|
||||
&shutdown_deadline,
|
||||
)
|
||||
}))
|
||||
@@ -511,12 +510,13 @@ async fn run_lane<S>(
|
||||
receiver: &Mutex<mpsc::UnboundedReceiver<QueuedExecutionObservation>>,
|
||||
store: &S,
|
||||
state: &ExecutionBufferState,
|
||||
cancellation: CancellationToken,
|
||||
context: WorkerDaemonContext,
|
||||
shutdown_deadline: &OnceLock<Instant>,
|
||||
) -> Result<(), WorkerTaskError>
|
||||
where
|
||||
S: ExecutionStore + OperationalDiagnostics + ?Sized,
|
||||
{
|
||||
let cancellation = context.cancellation();
|
||||
let mut receiver = receiver.lock().await;
|
||||
loop {
|
||||
let queued = tokio::select! {
|
||||
@@ -564,7 +564,11 @@ where
|
||||
drain_on_shutdown(lane, &mut receiver, store, state, deadline).await;
|
||||
return Ok(());
|
||||
},
|
||||
() = &mut persistence => {},
|
||||
succeeded = &mut persistence => {
|
||||
if succeeded {
|
||||
context.succeeded();
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -573,7 +577,7 @@ fn shared_shutdown_deadline(deadline: &OnceLock<Instant>) -> Instant {
|
||||
*deadline.get_or_init(|| Instant::now() + SHUTDOWN_DRAIN_TIMEOUT)
|
||||
}
|
||||
|
||||
async fn persist_queued<S>(mut queued: QueuedExecutionObservation, store: &S)
|
||||
async fn persist_queued<S>(mut queued: QueuedExecutionObservation, store: &S) -> bool
|
||||
where
|
||||
S: ExecutionStore + OperationalDiagnostics + ?Sized,
|
||||
{
|
||||
@@ -587,18 +591,21 @@ where
|
||||
request_id = ?request_id,
|
||||
"执行观测队列内部状态无效,已丢弃本次写入"
|
||||
);
|
||||
return;
|
||||
return false;
|
||||
};
|
||||
// Store 写入没有统一幂等键;超时可能表示已提交,不能在这里盲目重试并
|
||||
// 制造重复 ops_events
|
||||
// 失败会被计数并丢弃,数据面始终不等待补偿
|
||||
match write.persist(store).await {
|
||||
Ok(()) => state.record_persisted(),
|
||||
Ok(()) => {
|
||||
state.record_persisted();
|
||||
true
|
||||
}
|
||||
Err(error) => {
|
||||
state.record_write_failure();
|
||||
let stats = state.snapshot();
|
||||
if !stats.write_failure_total.is_power_of_two() {
|
||||
return;
|
||||
return false;
|
||||
}
|
||||
tracing::warn!(
|
||||
operation,
|
||||
@@ -607,6 +614,7 @@ where
|
||||
write_failure_total = stats.write_failure_total,
|
||||
"执行观测写入失败,数据面不受影响"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,7 +17,7 @@ use gateway_core::engine::admission::{
|
||||
};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::task::{DaemonTask, WorkerTaskError};
|
||||
use gateway_core::task::{DaemonTask, WorkerDaemonContext, WorkerTaskError};
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
|
||||
const DEFAULT_QUEUE_CAPACITY: usize = 4_096;
|
||||
@@ -139,8 +139,9 @@ pub struct ClientAdmissionReleaseWriter {
|
||||
}
|
||||
|
||||
impl DaemonTask for ClientAdmissionReleaseWriter {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
let mut receiver = self.receiver.lock().await;
|
||||
loop {
|
||||
let release = tokio::select! {
|
||||
@@ -155,7 +156,10 @@ impl DaemonTask for ClientAdmissionReleaseWriter {
|
||||
"client admission release queue closed",
|
||||
));
|
||||
};
|
||||
release_admission(self.inner.as_ref(), self.diagnostics.as_ref(), release).await;
|
||||
if release_admission(self.inner.as_ref(), self.diagnostics.as_ref(), release).await
|
||||
{
|
||||
context.succeeded();
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -185,7 +189,7 @@ async fn release_admission(
|
||||
port: &dyn ClientAdmissionPort,
|
||||
diagnostics: &dyn gateway_core::diagnostics::OperationalDiagnostics,
|
||||
release: AdmissionRelease,
|
||||
) {
|
||||
) -> bool {
|
||||
if let Err(source) = port
|
||||
.release(&release.client_api_key_id, &release.model_request_id)
|
||||
.await
|
||||
@@ -205,5 +209,7 @@ async fn release_admission(
|
||||
"admission diagnostic could not be recorded"
|
||||
);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
use super::*;
|
||||
use gateway_core::task::DaemonTask;
|
||||
use gateway_core::task::WorkerDaemonContext;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
|
||||
const COMMAND_DRAIN_TIMEOUT: Duration = Duration::from_secs(4);
|
||||
@@ -52,7 +53,7 @@ fn spawn_command_writer<T>(
|
||||
where
|
||||
T: DaemonTask + Send + 'static,
|
||||
{
|
||||
tokio::spawn(async move { writer.run(cancellation).await })
|
||||
tokio::spawn(async move { writer.run(WorkerDaemonContext::new(cancellation)).await })
|
||||
}
|
||||
|
||||
impl CommandStoreDrain {
|
||||
|
||||
@@ -12,8 +12,10 @@ use gateway_admin::{
|
||||
ports::store::ClientKeyStore as _,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::execution::ClientApiKeyUsageSink as _, lifecycle::CancellationToken,
|
||||
policy::ClientApiKeyId, task::DaemonTask as _,
|
||||
engine::execution::ClientApiKeyUsageSink as _,
|
||||
lifecycle::CancellationToken,
|
||||
policy::ClientApiKeyId,
|
||||
task::{DaemonTask as _, WorkerDaemonContext},
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
ClientApiKeyCursor, ClientApiKeyCursorValue, ClientApiKeyListQuery, ClientApiKeyRepository,
|
||||
@@ -592,7 +594,7 @@ async fn client_key_usage_daemon_should_flush_pending_touch_on_shutdown() {
|
||||
let task = {
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
tokio::spawn(async move { writer.run(cancellation).await })
|
||||
tokio::spawn(async move { writer.run(WorkerDaemonContext::new(cancellation)).await })
|
||||
};
|
||||
sink.record_used(&ClientApiKeyId::new("key_usage_daemon").expect("valid client API Key ID"));
|
||||
cancellation.cancel();
|
||||
|
||||
@@ -16,7 +16,7 @@ use gateway_core::error::{
|
||||
};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::routing::{ProviderKind, UpstreamModelId};
|
||||
use gateway_core::task::DaemonTask as _;
|
||||
use gateway_core::task::{DaemonTask as _, WorkerDaemonContext};
|
||||
use gateway_core::upstream::UpstreamSendState;
|
||||
use gateway_store::postgres::{BufferedExecutionStore, PgExecutionStore};
|
||||
use tokio::sync::{Notify, Semaphore};
|
||||
@@ -288,7 +288,7 @@ async fn observation_writer_persists_commands_in_enqueue_order() {
|
||||
let task = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
async move { writer.run(cancellation).await }
|
||||
async move { writer.run(WorkerDaemonContext::new(cancellation)).await }
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
@@ -338,7 +338,7 @@ async fn request_affine_lanes_run_independently_and_preserve_each_request_order(
|
||||
let task = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
async move { writer.run(cancellation).await }
|
||||
async move { writer.run(WorkerDaemonContext::new(cancellation)).await }
|
||||
});
|
||||
inner.wait_for_started(2).await;
|
||||
assert_eq!(inner.maximum_active.load(Ordering::Acquire), 2);
|
||||
@@ -383,7 +383,7 @@ async fn idle_wait_includes_a_dequeued_write_until_persistence_finishes() {
|
||||
let task = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
async move { writer.run(cancellation).await }
|
||||
async move { writer.run(WorkerDaemonContext::new(cancellation)).await }
|
||||
});
|
||||
inner.wait_for_started(1).await;
|
||||
|
||||
@@ -430,7 +430,7 @@ async fn global_item_budget_includes_writes_running_in_every_lane() {
|
||||
let task = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
async move { writer.run(cancellation).await }
|
||||
async move { writer.run(WorkerDaemonContext::new(cancellation)).await }
|
||||
});
|
||||
inner.wait_for_started(2).await;
|
||||
assert_eq!(store.stats().queued_items, 2);
|
||||
@@ -477,7 +477,7 @@ async fn all_lanes_share_one_shutdown_deadline_for_inflight_writes() {
|
||||
let task = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
async move { writer.run(cancellation).await }
|
||||
async move { writer.run(WorkerDaemonContext::new(cancellation)).await }
|
||||
});
|
||||
inner.wait_for_started(4).await;
|
||||
cancellation.cancel();
|
||||
@@ -512,7 +512,7 @@ async fn observation_writer_continues_after_a_database_write_failure() {
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.expect("shutdown drains observations");
|
||||
|
||||
@@ -545,7 +545,7 @@ async fn shutdown_drains_observations_already_accepted_by_the_queue() {
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.expect("shutdown drains observations");
|
||||
|
||||
@@ -653,7 +653,7 @@ async fn zero_attempt_create_and_finalize_drain_in_order_with_trace() {
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.expect("shutdown drains accepted commands");
|
||||
let stats = store.stats();
|
||||
@@ -703,7 +703,7 @@ async fn request_affine_lanes_persist_a_postgres_lifecycle_burst_without_orphans
|
||||
let task = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
async move { writer.run(cancellation).await }
|
||||
async move { writer.run(WorkerDaemonContext::new(cancellation)).await }
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(15), async {
|
||||
while store.stats().persisted_total != WRITES {
|
||||
@@ -760,7 +760,10 @@ async fn zero_attempt_dropped_finalize_is_fail_open_and_recoverable_at_deadline(
|
||||
assert_eq!(store.stats().dropped_total, 1);
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer.run(cancellation).await.expect("drain create only");
|
||||
writer
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.expect("drain create only");
|
||||
let repository = observability_repository(&database.pool);
|
||||
let pending = repository
|
||||
.usage_record_detail(request.id.as_str())
|
||||
@@ -842,7 +845,7 @@ async fn zero_attempt_dropped_create_does_not_make_finalize_an_upsert() {
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.expect("drain orphan finalization");
|
||||
assert_eq!(store.stats().persisted_total, 0);
|
||||
@@ -913,7 +916,7 @@ async fn zero_attempt_postgres_write_failures_are_not_retried_and_do_not_stop_th
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer
|
||||
.run(cancellation)
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.expect("continue draining after PG errors");
|
||||
let stats = store.stats();
|
||||
@@ -1023,7 +1026,10 @@ async fn operational_failure_uses_the_bounded_queue_without_creating_a_model_req
|
||||
assert_eq!(store.stats().dropped_total, 1);
|
||||
let cancellation = CancellationToken::new();
|
||||
cancellation.cancel();
|
||||
writer.run(cancellation).await.unwrap();
|
||||
writer
|
||||
.run(WorkerDaemonContext::new(cancellation))
|
||||
.await
|
||||
.unwrap();
|
||||
let row: (
|
||||
Option<String>,
|
||||
String,
|
||||
@@ -1105,7 +1111,10 @@ async fn progress_congestion_preserves_failure_reserve_without_waiting() {
|
||||
assert_eq!(store.stats().dropped_total, 1);
|
||||
let cancelled = CancellationToken::new();
|
||||
cancelled.cancel();
|
||||
writer.run(cancelled).await.unwrap();
|
||||
writer
|
||||
.run(WorkerDaemonContext::new(cancelled))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(store.stats().persisted_total, 5);
|
||||
let detail: String =
|
||||
sqlx::query_scalar("select error_details from ops_events where component = 'admin'")
|
||||
|
||||
@@ -13,6 +13,7 @@ use gateway_core::engine::admission::{
|
||||
};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::policy::{ClientApiKeyId, RateLimits};
|
||||
use gateway_core::task::WorkerDaemonContext;
|
||||
use gateway_store::redis::BufferedClientAdmissionPort;
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -280,5 +281,5 @@ fn spawn_writer<T>(
|
||||
where
|
||||
T: gateway_core::task::DaemonTask + 'static,
|
||||
{
|
||||
tokio::spawn(async move { writer.run(cancellation).await })
|
||||
tokio::spawn(async move { writer.run(WorkerDaemonContext::new(cancellation)).await })
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
use super::*;
|
||||
use crate::transport::profile::cli_release::CliReleaseService;
|
||||
use crate::transport::profile::platform_release::PlatformDesktopReleaseService;
|
||||
use gateway_core::task::WorkerDaemonContext;
|
||||
|
||||
pub(crate) struct ClientReleaseServices {
|
||||
pub desktop: Arc<CodexDesktopReleaseService>,
|
||||
@@ -299,11 +300,9 @@ impl ScheduledTask for OpenAiCatalogTask {
|
||||
}
|
||||
|
||||
impl DaemonTask for OpenAiCatalogEtagTask {
|
||||
fn run(
|
||||
&self,
|
||||
cancellation: gateway_core::lifecycle::CancellationToken,
|
||||
) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
fn run(&self, context: WorkerDaemonContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let cancellation = context.cancellation();
|
||||
loop {
|
||||
tokio::select! {
|
||||
() = cancellation.cancelled() => return Ok(()),
|
||||
@@ -314,6 +313,8 @@ impl DaemonTask for OpenAiCatalogEtagTask {
|
||||
error = %error,
|
||||
"OpenAI model catalog ETag refresh failed"
|
||||
);
|
||||
} else {
|
||||
context.succeeded();
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
@@ -1022,6 +1022,11 @@ Worker 由各 Bundle 贡献、由 Host 统一监督:
|
||||
以及账号冻结恢复 worker(容量熔断的自适应并发下调与到期探测解冻)
|
||||
- Provider:credential refresh、quota/catalog 健康、账号预热和官方版本/etag 检查
|
||||
|
||||
Daemon 每次运行持有独立的取消与成功进度上下文,实际工作或依赖检查成功后由 owner 上报。
|
||||
Host 据此清除连续失败与当前错误、重置重启退避,并保持 `Running`;仅重启不会恢复健康,
|
||||
已退出运行的通知不影响新运行。周期任务按调度间隔检查成功新鲜度,daemon 的空闲或长任务
|
||||
不受重启退避时长限制,运行中的依赖健康由相应探针和业务 owner 负责
|
||||
|
||||
预热与备份按部署时区解释本地时刻,缺失时刻跳过、重复时刻取较早一次。
|
||||
预热由 Store 在运行设置记录中持久化 UTC 执行游标,只领取晚于游标的时刻,重启或回拨不重复执行历史时段。
|
||||
游标独立于管理配置的替换与版本发布;取消仍沿既有 worker lease 生命周期释放。
|
||||
|
||||
Reference in New Issue
Block a user