1use std::future::Future;
2use std::panic::{AssertUnwindSafe, catch_unwind};
3use std::pin::Pin;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, Mutex, MutexGuard};
6use std::time::{Duration, Instant};
7
8use talos_agent::session::{
9 RuntimeActiveTurnOutcome as AgentActiveTurnOutcome, RuntimeAdmissionClose,
10 RuntimeAdmissionControl,
11 RuntimeDurableReconciliationOutcome as AgentDurableReconciliationOutcome,
12 RuntimeShutdownTurnPolicy,
13};
14use talos_core::session::SessionOp;
15use thiserror::Error;
16use tokio::runtime::Handle;
17use tokio::sync::{Notify, mpsc};
18use tokio::task::{JoinError, JoinHandle};
19
20use crate::RuntimeResult;
21
22const LEGACY_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30);
23static NEXT_SHUTDOWN_PLAN_ID: AtomicU64 = AtomicU64::new(1);
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27#[non_exhaustive]
28pub enum ActiveTurnPolicy {
29 FinishCurrent {
31 grace: Duration,
33 },
34 Interrupt,
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
40#[non_exhaustive]
41pub enum ShutdownOptionsError {
42 #[error("shutdown total timeout must be greater than zero")]
44 ZeroTotalTimeout,
45 #[error("finish grace must be less than the total shutdown timeout")]
47 FinishGraceNotLessThanTotal,
48 #[error("shutdown total timeout exceeds the monotonic clock range")]
50 TotalTimeoutOutOfRange,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub struct ShutdownOptions {
56 total_timeout: Duration,
57 active_turn_policy: ActiveTurnPolicy,
58}
59
60impl ShutdownOptions {
61 pub fn new(
63 total_timeout: Duration,
64 active_turn_policy: ActiveTurnPolicy,
65 ) -> Result<Self, ShutdownOptionsError> {
66 if total_timeout.is_zero() {
67 return Err(ShutdownOptionsError::ZeroTotalTimeout);
68 }
69 if Instant::now().checked_add(total_timeout).is_none() {
70 return Err(ShutdownOptionsError::TotalTimeoutOutOfRange);
71 }
72 if let ActiveTurnPolicy::FinishCurrent { grace } = active_turn_policy
73 && grace >= total_timeout
74 {
75 return Err(ShutdownOptionsError::FinishGraceNotLessThanTotal);
76 }
77 Ok(Self {
78 total_timeout,
79 active_turn_policy,
80 })
81 }
82
83 pub fn interrupt(total_timeout: Duration) -> Result<Self, ShutdownOptionsError> {
85 Self::new(total_timeout, ActiveTurnPolicy::Interrupt)
86 }
87
88 pub fn finish_current(
90 total_timeout: Duration,
91 grace: Duration,
92 ) -> Result<Self, ShutdownOptionsError> {
93 Self::new(total_timeout, ActiveTurnPolicy::FinishCurrent { grace })
94 }
95
96 #[must_use]
98 pub const fn total_timeout(&self) -> Duration {
99 self.total_timeout
100 }
101
102 #[must_use]
104 pub const fn active_turn_policy(&self) -> ActiveTurnPolicy {
105 self.active_turn_policy
106 }
107
108 pub(crate) fn legacy_default() -> Self {
109 Self {
110 total_timeout: LEGACY_SHUTDOWN_TIMEOUT,
111 active_turn_policy: ActiveTurnPolicy::Interrupt,
112 }
113 }
114}
115
116#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
118pub struct ShutdownPlanId(u64);
119
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
122#[non_exhaustive]
123pub enum ShutdownActiveTurnOutcome {
124 Idle,
126 Finished,
128 InterruptedAndFinalized,
130 Failed,
132 Unreconciled,
134}
135
136#[derive(Debug, Clone, Copy, PartialEq, Eq)]
138#[non_exhaustive]
139pub enum ShutdownDurableOutcome {
140 Completed {
142 rejected_pending: u32,
144 },
145 Failed {
147 rejected_pending: u32,
149 },
150 NotRunDeadline,
152}
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156#[non_exhaustive]
157pub enum ShutdownActorOutcome {
158 Joined,
160 Failed,
162 Contained,
164}
165
166#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
171pub struct ShutdownFinalizerId(&'static str);
172
173impl ShutdownFinalizerId {
174 pub(crate) const fn new(value: &'static str) -> Self {
175 Self(value)
176 }
177
178 #[must_use]
180 pub const fn as_str(self) -> &'static str {
181 self.0
182 }
183}
184
185#[derive(Debug, Clone, Copy, PartialEq, Eq)]
187#[non_exhaustive]
188pub enum ShutdownFinalizerOutcome {
189 Completed,
191 Failed,
193 Panicked,
195 TimedOut,
197 NotRunDeadline,
199}
200
201#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub struct ShutdownFinalizerReport {
204 identifier: ShutdownFinalizerId,
205 outcome: ShutdownFinalizerOutcome,
206}
207
208impl ShutdownFinalizerReport {
209 #[must_use]
211 pub const fn identifier(&self) -> ShutdownFinalizerId {
212 self.identifier
213 }
214
215 #[must_use]
217 pub const fn outcome(&self) -> ShutdownFinalizerOutcome {
218 self.outcome
219 }
220}
221
222#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
224#[non_exhaustive]
225pub enum ShutdownFinalizerRegistryError {
226 #[error("duplicate shutdown finalizer identifier")]
228 DuplicateIdentifier,
229 #[error("duplicate shutdown finalizer order")]
231 DuplicateOrder,
232 #[error("shutdown finalizer cap must be greater than zero")]
234 ZeroCap,
235}
236
237#[derive(Debug, Clone, Copy, PartialEq, Eq)]
238pub(crate) struct RuntimeFinalizerError;
239
240pub(crate) type RuntimeFinalizerFuture =
241 Pin<Box<dyn Future<Output = Result<(), RuntimeFinalizerError>> + Send + 'static>>;
242
243pub(crate) trait RuntimeFinalizer: Send + Sync {
244 fn identifier(&self) -> ShutdownFinalizerId;
245 fn order(&self) -> u16;
246 fn cap(&self) -> Duration;
247 fn finalize(&self) -> RuntimeFinalizerFuture;
248}
249
250pub(crate) struct RuntimeFinalizerRegistry {
251 entries: Vec<Arc<dyn RuntimeFinalizer>>,
252}
253
254impl RuntimeFinalizerRegistry {
255 pub(crate) fn freeze(
256 mut entries: Vec<Arc<dyn RuntimeFinalizer>>,
257 ) -> Result<Self, ShutdownFinalizerRegistryError> {
258 for (index, entry) in entries.iter().enumerate() {
259 if entry.cap().is_zero() {
260 return Err(ShutdownFinalizerRegistryError::ZeroCap);
261 }
262 for other in &entries[..index] {
263 if other.identifier() == entry.identifier() {
264 return Err(ShutdownFinalizerRegistryError::DuplicateIdentifier);
265 }
266 if other.order() == entry.order() {
267 return Err(ShutdownFinalizerRegistryError::DuplicateOrder);
268 }
269 }
270 }
271 entries.sort_by_key(|entry| entry.order());
272 Ok(Self { entries })
273 }
274
275 fn not_run_deadline(&self) -> Vec<ShutdownFinalizerReport> {
276 self.entries
277 .iter()
278 .map(|entry| ShutdownFinalizerReport {
279 identifier: entry.identifier(),
280 outcome: ShutdownFinalizerOutcome::NotRunDeadline,
281 })
282 .collect()
283 }
284}
285
286#[derive(Debug, Clone, PartialEq, Eq)]
288#[non_exhaustive]
289pub struct ShutdownReport {
290 plan_id: ShutdownPlanId,
291 active_turn_policy: ActiveTurnPolicy,
292 elapsed: Duration,
293 deadline_exhausted: bool,
294 active_turn: ShutdownActiveTurnOutcome,
295 durable_reconciliation: ShutdownDurableOutcome,
296 finalizers: Vec<ShutdownFinalizerReport>,
297 actor: ShutdownActorOutcome,
298}
299
300impl ShutdownReport {
301 #[must_use]
303 pub const fn plan_id(&self) -> ShutdownPlanId {
304 self.plan_id
305 }
306
307 #[must_use]
309 pub const fn active_turn_policy(&self) -> ActiveTurnPolicy {
310 self.active_turn_policy
311 }
312
313 #[must_use]
315 pub const fn elapsed(&self) -> Duration {
316 self.elapsed
317 }
318
319 #[must_use]
321 pub const fn deadline_exhausted(&self) -> bool {
322 self.deadline_exhausted
323 }
324
325 #[must_use]
327 pub const fn active_turn(&self) -> ShutdownActiveTurnOutcome {
328 self.active_turn
329 }
330
331 #[must_use]
333 pub const fn durable_reconciliation(&self) -> ShutdownDurableOutcome {
334 self.durable_reconciliation
335 }
336
337 #[must_use]
339 pub fn finalizers(&self) -> &[ShutdownFinalizerReport] {
340 &self.finalizers
341 }
342
343 #[must_use]
345 pub const fn actor(&self) -> ShutdownActorOutcome {
346 self.actor
347 }
348
349 #[must_use]
351 pub fn is_complete(&self) -> bool {
352 !self.deadline_exhausted
353 && !matches!(self.active_turn, ShutdownActiveTurnOutcome::Unreconciled)
354 && matches!(
355 self.durable_reconciliation,
356 ShutdownDurableOutcome::Completed { .. }
357 )
358 && self
359 .finalizers
360 .iter()
361 .all(|entry| matches!(entry.outcome, ShutdownFinalizerOutcome::Completed))
362 && matches!(self.actor, ShutdownActorOutcome::Joined)
363 }
364}
365
366#[derive(Clone)]
368pub struct RuntimeShutdownHandle {
369 pub(crate) coordinator: Arc<ShutdownCoordinator>,
370}
371
372impl RuntimeShutdownHandle {
373 pub async fn shutdown(&self, options: ShutdownOptions) -> RuntimeResult<ShutdownReport> {
375 self.coordinator.shutdown(options).await
376 }
377}
378
379#[derive(Debug, Clone, Copy)]
380struct AcceptedPlan {
381 id: ShutdownPlanId,
382 options: ShutdownOptions,
383 accepted_at: Instant,
384 deadline: Instant,
385 active_at_fence: bool,
386}
387
388enum CoordinatorState {
389 Open,
390 Closing(AcceptedPlan),
391 Closed {
392 report: ShutdownReport,
393 actor_join_error: Option<JoinError>,
394 },
395}
396
397pub(crate) struct ShutdownCoordinator {
398 admission: RuntimeAdmissionControl,
399 command_tx: mpsc::Sender<SessionOp>,
400 actor_task: Mutex<Option<JoinHandle<()>>>,
401 state: Mutex<CoordinatorState>,
402 changed: Notify,
403 runtime: Handle,
404 finalizers: RuntimeFinalizerRegistry,
405}
406
407impl ShutdownCoordinator {
408 pub(crate) fn new(
409 admission: RuntimeAdmissionControl,
410 command_tx: mpsc::Sender<SessionOp>,
411 actor_task: JoinHandle<()>,
412 runtime: Handle,
413 finalizers: RuntimeFinalizerRegistry,
414 ) -> Arc<Self> {
415 Arc::new(Self {
416 admission,
417 command_tx,
418 actor_task: Mutex::new(Some(actor_task)),
419 state: Mutex::new(CoordinatorState::Open),
420 changed: Notify::new(),
421 runtime,
422 finalizers,
423 })
424 }
425
426 fn state(&self) -> MutexGuard<'_, CoordinatorState> {
427 self.state
428 .lock()
429 .unwrap_or_else(std::sync::PoisonError::into_inner)
430 }
431
432 fn actor_task(&self) -> MutexGuard<'_, Option<JoinHandle<()>>> {
433 self.actor_task
434 .lock()
435 .unwrap_or_else(std::sync::PoisonError::into_inner)
436 }
437
438 fn initiate(self: &Arc<Self>, options: ShutdownOptions) -> ShutdownPlanId {
439 let candidate = NEXT_SHUTDOWN_PLAN_ID.fetch_add(1, Ordering::Relaxed);
440 let policy = match options.active_turn_policy {
441 ActiveTurnPolicy::FinishCurrent { .. } => RuntimeShutdownTurnPolicy::FinishCurrent,
442 ActiveTurnPolicy::Interrupt => RuntimeShutdownTurnPolicy::Interrupt,
443 };
444 match self.admission.begin_shutdown(candidate, policy) {
445 RuntimeAdmissionClose::Existing { plan_id } => ShutdownPlanId(plan_id),
446 RuntimeAdmissionClose::Accepted { active_at_fence } => {
447 let accepted_at = Instant::now();
448 let deadline = accepted_at
449 .checked_add(options.total_timeout)
450 .unwrap_or(accepted_at);
451 let plan = AcceptedPlan {
452 id: ShutdownPlanId(candidate),
453 options,
454 accepted_at,
455 deadline,
456 active_at_fence,
457 };
458 *self.state() = CoordinatorState::Closing(plan);
459 self.changed.notify_waiters();
460 if matches!(options.active_turn_policy, ActiveTurnPolicy::Interrupt) {
461 self.admission.interrupt_active();
462 }
463 let coordinator = self.clone();
464 let spawn = catch_unwind(AssertUnwindSafe(|| {
465 self.runtime.spawn(async move {
466 coordinator.drive(plan).await;
467 })
468 }));
469 if spawn.is_err() {
470 self.publish_driver_failure(plan);
471 }
472 plan.id
473 }
474 }
475 }
476
477 pub(crate) async fn shutdown(
478 self: &Arc<Self>,
479 options: ShutdownOptions,
480 ) -> RuntimeResult<ShutdownReport> {
481 let plan_id = self.initiate(options);
482 self.wait_for_report(plan_id).await
483 }
484
485 pub(crate) fn initiate_default(self: &Arc<Self>) {
486 let _ = self.initiate(ShutdownOptions::legacy_default());
487 }
488
489 pub(crate) fn commit_reserved(
490 &self,
491 permit: mpsc::Permit<'_, SessionOp>,
492 op: SessionOp,
493 ) -> Result<(), SessionOp> {
494 self.admission.commit_reserved(permit, op)
495 }
496
497 pub(crate) fn is_admission_open(&self) -> bool {
498 self.admission.is_open()
499 }
500
501 async fn wait_for_report(&self, plan_id: ShutdownPlanId) -> RuntimeResult<ShutdownReport> {
502 loop {
503 let changed = self.changed.notified();
504 match &*self.state() {
505 CoordinatorState::Closed { report, .. } => return Ok(report.clone()),
506 CoordinatorState::Closing(plan) if plan.id == plan_id => {}
507 CoordinatorState::Open | CoordinatorState::Closing(_) => {}
508 }
509 changed.await;
510 }
511 }
512
513 pub(crate) fn take_actor_join_error(&self) -> Option<JoinError> {
514 match &mut *self.state() {
515 CoordinatorState::Closed {
516 actor_join_error, ..
517 } => actor_join_error.take(),
518 CoordinatorState::Open | CoordinatorState::Closing(_) => None,
519 }
520 }
521
522 async fn drive(self: Arc<Self>, plan: AcceptedPlan) {
523 let deadline = tokio::time::Instant::from_std(plan.deadline);
524 let send_result =
525 tokio::time::timeout_at(deadline, self.command_tx.send(SessionOp::Shutdown)).await;
526
527 if matches!(
528 plan.options.active_turn_policy,
529 ActiveTurnPolicy::FinishCurrent { .. }
530 ) && plan.active_at_fence
531 {
532 let ActiveTurnPolicy::FinishCurrent { grace } = plan.options.active_turn_policy else {
533 unreachable!("policy matched above")
534 };
535 let grace_end = plan.accepted_at.checked_add(grace).unwrap_or(plan.deadline);
536 let grace_deadline = tokio::time::Instant::from_std(plan.deadline.min(grace_end));
537 if tokio::time::timeout_at(grace_deadline, self.admission.wait_until_idle())
538 .await
539 .is_err()
540 {
541 self.admission.interrupt_active();
542 }
543 }
544
545 let mut actor_join_error = None;
546 let actor = self.actor_task().take();
547 let shutdown_barrier =
548 tokio::time::timeout_at(deadline, self.admission.wait_until_shutdown_barrier()).await;
549 let barrier_finished = shutdown_barrier.is_ok();
550
551 let snapshot = self.admission.snapshot();
552 let durable_reconciliation = match shutdown_barrier {
553 Ok(true) => match snapshot.durable_reconciliation {
554 AgentDurableReconciliationOutcome::Completed => ShutdownDurableOutcome::Completed {
555 rejected_pending: snapshot.rejected_pending,
556 },
557 AgentDurableReconciliationOutcome::Failed => ShutdownDurableOutcome::Failed {
558 rejected_pending: snapshot.rejected_pending,
559 },
560 AgentDurableReconciliationOutcome::Pending => ShutdownDurableOutcome::Failed {
561 rejected_pending: snapshot.rejected_pending,
562 },
563 },
564 Ok(false) => ShutdownDurableOutcome::Failed {
565 rejected_pending: snapshot.rejected_pending,
566 },
567 Err(_) => ShutdownDurableOutcome::NotRunDeadline,
568 };
569
570 let finalizers = if barrier_finished {
571 self.run_finalizers(deadline).await
572 } else {
573 self.finalizers.not_run_deadline()
574 };
575
576 let actor_outcome = if let Some(mut actor) = actor {
577 match tokio::time::timeout_at(deadline, &mut actor).await {
578 Ok(Ok(())) => ShutdownActorOutcome::Joined,
579 Ok(Err(error)) => {
580 actor_join_error = Some(error);
581 ShutdownActorOutcome::Failed
582 }
583 Err(_) => {
584 actor.abort();
585 ShutdownActorOutcome::Contained
586 }
587 }
588 } else {
589 ShutdownActorOutcome::Failed
590 };
591
592 let elapsed = plan.accepted_at.elapsed();
593 let deadline_exhausted = Instant::now() >= plan.deadline
594 || send_result.is_err()
595 || matches!(actor_outcome, ShutdownActorOutcome::Contained);
596 let snapshot = self.admission.snapshot();
597 let active_turn = if !plan.active_at_fence {
598 ShutdownActiveTurnOutcome::Idle
599 } else {
600 match snapshot.active_turn {
601 AgentActiveTurnOutcome::Idle => ShutdownActiveTurnOutcome::Idle,
602 AgentActiveTurnOutcome::Finished => ShutdownActiveTurnOutcome::Finished,
603 AgentActiveTurnOutcome::InterruptedAndFinalized => {
604 ShutdownActiveTurnOutcome::InterruptedAndFinalized
605 }
606 AgentActiveTurnOutcome::Failed => ShutdownActiveTurnOutcome::Failed,
607 AgentActiveTurnOutcome::Running => ShutdownActiveTurnOutcome::Unreconciled,
608 }
609 };
610 let report = ShutdownReport {
611 plan_id: plan.id,
612 active_turn_policy: plan.options.active_turn_policy,
613 elapsed,
614 deadline_exhausted,
615 active_turn,
616 durable_reconciliation,
617 finalizers,
618 actor: actor_outcome,
619 };
620 self.admission.mark_closed();
621 *self.state() = CoordinatorState::Closed {
622 report,
623 actor_join_error,
624 };
625 self.changed.notify_waiters();
626 }
627
628 async fn run_finalizers(&self, deadline: tokio::time::Instant) -> Vec<ShutdownFinalizerReport> {
629 let mut reports = Vec::with_capacity(self.finalizers.entries.len());
630 for (index, entry) in self.finalizers.entries.iter().enumerate() {
631 let now = tokio::time::Instant::now();
632 if now >= deadline {
633 reports.extend(self.finalizers.entries[index..].iter().map(|remaining| {
634 ShutdownFinalizerReport {
635 identifier: remaining.identifier(),
636 outcome: ShutdownFinalizerOutcome::NotRunDeadline,
637 }
638 }));
639 break;
640 }
641
642 let remaining = deadline.saturating_duration_since(now);
643 let finalizer_deadline = now + entry.cap().min(remaining);
644 let future = catch_unwind(AssertUnwindSafe(|| entry.finalize()));
645 let outcome = match future {
646 Err(_) => ShutdownFinalizerOutcome::Panicked,
647 Ok(future) => {
648 let spawned = catch_unwind(AssertUnwindSafe(|| self.runtime.spawn(future)));
649 match spawned {
650 Err(_) => ShutdownFinalizerOutcome::Panicked,
651 Ok(mut task) => {
652 match tokio::time::timeout_at(finalizer_deadline, &mut task).await {
653 Ok(Ok(Ok(()))) => ShutdownFinalizerOutcome::Completed,
654 Ok(Ok(Err(_))) => ShutdownFinalizerOutcome::Failed,
655 Ok(Err(error)) if error.is_panic() => {
656 ShutdownFinalizerOutcome::Panicked
657 }
658 Ok(Err(_)) => ShutdownFinalizerOutcome::Failed,
659 Err(_) => {
660 task.abort();
661 let _ = task.await;
662 ShutdownFinalizerOutcome::TimedOut
663 }
664 }
665 }
666 }
667 }
668 };
669 reports.push(ShutdownFinalizerReport {
670 identifier: entry.identifier(),
671 outcome,
672 });
673 }
674 reports
675 }
676
677 fn publish_driver_failure(&self, plan: AcceptedPlan) {
678 if let Some(actor) = self.actor_task().take() {
679 actor.abort();
680 }
681 let snapshot = self.admission.snapshot();
682 self.admission.mark_closed();
683 *self.state() = CoordinatorState::Closed {
684 report: ShutdownReport {
685 plan_id: plan.id,
686 active_turn_policy: plan.options.active_turn_policy,
687 elapsed: plan.accepted_at.elapsed(),
688 deadline_exhausted: true,
689 active_turn: if plan.active_at_fence {
690 ShutdownActiveTurnOutcome::Unreconciled
691 } else {
692 ShutdownActiveTurnOutcome::Idle
693 },
694 durable_reconciliation: match snapshot.durable_reconciliation {
695 AgentDurableReconciliationOutcome::Completed => {
696 ShutdownDurableOutcome::Completed {
697 rejected_pending: snapshot.rejected_pending,
698 }
699 }
700 AgentDurableReconciliationOutcome::Failed => ShutdownDurableOutcome::Failed {
701 rejected_pending: snapshot.rejected_pending,
702 },
703 AgentDurableReconciliationOutcome::Pending => {
704 ShutdownDurableOutcome::NotRunDeadline
705 }
706 },
707 finalizers: self.finalizers.not_run_deadline(),
708 actor: ShutdownActorOutcome::Contained,
709 },
710 actor_join_error: None,
711 };
712 self.changed.notify_waiters();
713 }
714}