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:
Zyy
2026-10-21 11:08:15 +08:00
committed by GitHub
co-authored by Zyy Codex
parent 52fb1e7180
commit ef11b72d8a
23 changed files with 463 additions and 95 deletions
@@ -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;
}
+53 -1
View File
@@ -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));
+46 -24
View File
@@ -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 &registration.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"),
+5 -3
View File
@@ -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(())
})
+188 -4
View File
@@ -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 -1
View File
@@ -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();
}
}
})
+5
View File
@@ -1022,6 +1022,11 @@ Worker 由各 Bundle 贡献、由 Host 统一监督:
以及账号冻结恢复 worker(容量熔断的自适应并发下调与到期探测解冻)
- Provider:credential refresh、quota/catalog 健康、账号预热和官方版本/etag 检查
Daemon 每次运行持有独立的取消与成功进度上下文,实际工作或依赖检查成功后由 owner 上报。
Host 据此清除连续失败与当前错误、重置重启退避,并保持 `Running`;仅重启不会恢复健康,
已退出运行的通知不影响新运行。周期任务按调度间隔检查成功新鲜度,daemon 的空闲或长任务
不受重启退避时长限制,运行中的依赖健康由相应探针和业务 owner 负责
预热与备份按部署时区解释本地时刻,缺失时刻跳过、重复时刻取较早一次。
预热由 Store 在运行设置记录中持久化 UTC 执行游标,只领取晚于游标的时刻,重启或回拨不重复执行历史时段。
游标独立于管理配置的替换与版本发布;取消仍沿既有 worker lease 生命周期释放。