1use std::collections::BTreeMap;
8use std::error::Error;
9use std::fmt;
10use std::future::Future;
11use std::panic::AssertUnwindSafe;
12use std::sync::Arc;
13use std::sync::atomic::{AtomicU8, Ordering};
14use std::time::Duration;
15
16use futures_util::FutureExt;
17use tokio::sync::Notify;
18use tokio::task::JoinSet;
19
20use crate::{TelemetryEventKind, TelemetryEventSink, TelemetryRecord};
21
22const ACCEPTING: u8 = 0;
23const STOPPING: u8 = 1;
24const ESCALATED: u8 = 2;
25
26pub const MIN_SHUTDOWN_DEADLINE: Duration = Duration::from_secs(1);
28pub const MAX_SHUTDOWN_DEADLINE: Duration = Duration::from_hours(1);
30pub const DEFAULT_SHUTDOWN_DEADLINE: Duration = Duration::from_secs(30);
32pub const MIN_TELEMETRY_FLUSH_DEADLINE: Duration = Duration::from_millis(100);
34pub const MAX_TELEMETRY_FLUSH_DEADLINE: Duration = Duration::from_mins(1);
36pub const DEFAULT_TELEMETRY_FLUSH_DEADLINE: Duration = Duration::from_secs(5);
38
39#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
41pub struct ShutdownDeadline(Duration);
42
43impl ShutdownDeadline {
44 pub fn new(value: Duration) -> Result<Self, ShutdownError> {
50 if !(MIN_SHUTDOWN_DEADLINE..=MAX_SHUTDOWN_DEADLINE).contains(&value) {
51 return Err(ShutdownError::InvalidShutdownDeadline);
52 }
53 Ok(Self(value))
54 }
55
56 #[must_use]
58 pub const fn get(self) -> Duration {
59 self.0
60 }
61}
62
63impl Default for ShutdownDeadline {
64 fn default() -> Self {
65 Self(DEFAULT_SHUTDOWN_DEADLINE)
66 }
67}
68
69#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
71pub struct TaskJoinDeadline(Duration);
72
73impl TaskJoinDeadline {
74 pub fn new(value: Duration, shutdown: ShutdownDeadline) -> Result<Self, ShutdownError> {
81 if !(MIN_SHUTDOWN_DEADLINE..=MAX_SHUTDOWN_DEADLINE).contains(&value) {
82 return Err(ShutdownError::InvalidTaskJoinDeadline);
83 }
84 if value > shutdown.get() {
85 return Err(ShutdownError::TaskJoinExceedsShutdown);
86 }
87 Ok(Self(value))
88 }
89
90 #[must_use]
92 pub const fn get(self) -> Duration {
93 self.0
94 }
95}
96
97impl Default for TaskJoinDeadline {
98 fn default() -> Self {
99 Self(DEFAULT_SHUTDOWN_DEADLINE)
100 }
101}
102
103#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
105pub struct TelemetryFlushDeadline(Duration);
106
107impl TelemetryFlushDeadline {
108 pub fn new(value: Duration) -> Result<Self, ShutdownError> {
115 if !(MIN_TELEMETRY_FLUSH_DEADLINE..=MAX_TELEMETRY_FLUSH_DEADLINE).contains(&value) {
116 return Err(ShutdownError::InvalidTelemetryFlushDeadline);
117 }
118 Ok(Self(value))
119 }
120
121 #[must_use]
123 pub const fn get(self) -> Duration {
124 self.0
125 }
126}
127
128impl Default for TelemetryFlushDeadline {
129 fn default() -> Self {
130 Self(DEFAULT_TELEMETRY_FLUSH_DEADLINE)
131 }
132}
133
134#[derive(Debug)]
135struct SignalState {
136 state: AtomicU8,
137 notify: Notify,
138}
139
140#[derive(Clone, Debug)]
142pub struct ShutdownSignal {
143 state: Arc<SignalState>,
144}
145
146impl ShutdownSignal {
147 fn new() -> Self {
148 Self {
149 state: Arc::new(SignalState {
150 state: AtomicU8::new(ACCEPTING),
151 notify: Notify::new(),
152 }),
153 }
154 }
155
156 #[must_use]
158 pub fn request_shutdown(&self) -> ShutdownRequest {
159 loop {
160 let current = self.state.state.load(Ordering::Acquire);
161 let (next, outcome) = match current {
162 ACCEPTING => (STOPPING, ShutdownRequest::Initiated),
163 STOPPING => (ESCALATED, ShutdownRequest::Escalated),
164 _ => return ShutdownRequest::AlreadyEscalated,
165 };
166 if self
167 .state
168 .state
169 .compare_exchange(current, next, Ordering::AcqRel, Ordering::Acquire)
170 .is_ok()
171 {
172 self.state.notify.notify_waiters();
173 return outcome;
174 }
175 }
176 }
177
178 #[must_use]
180 pub fn is_shutdown_requested(&self) -> bool {
181 self.state.state.load(Ordering::Acquire) >= STOPPING
182 }
183
184 #[must_use]
186 pub fn is_escalated(&self) -> bool {
187 self.state.state.load(Ordering::Acquire) >= ESCALATED
188 }
189
190 pub async fn cancelled(&self) {
192 self.wait_for(STOPPING).await;
193 }
194
195 async fn escalated(&self) {
196 self.wait_for(ESCALATED).await;
197 }
198
199 async fn wait_for(&self, target: u8) {
200 loop {
201 let notified = self.state.notify.notified();
202 if self.state.state.load(Ordering::Acquire) >= target {
203 return;
204 }
205 notified.await;
206 }
207 }
208
209 fn begin_shutdown(&self) {
210 if self
211 .state
212 .state
213 .compare_exchange(ACCEPTING, STOPPING, Ordering::AcqRel, Ordering::Acquire)
214 .is_ok()
215 {
216 self.state.notify.notify_waiters();
217 }
218 }
219
220 pub fn ensure_accepting(&self) -> Result<(), ShutdownError> {
226 if self.is_shutdown_requested() {
227 Err(ShutdownError::ShuttingDown)
228 } else {
229 Ok(())
230 }
231 }
232}
233
234#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
236#[non_exhaustive]
237pub enum ShutdownRequest {
238 Initiated,
240 Escalated,
242 AlreadyEscalated,
244}
245
246#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
248#[non_exhaustive]
249pub enum ShutdownTaskPhase {
250 Tasklet,
252 ChunkReadProcess,
254 ChunkWrite,
256 Transaction,
258 RetryBackoff,
260 FlowDecision,
262}
263
264#[derive(Clone, Copy, Debug, Eq, PartialEq)]
266pub struct UnjoinedPhase {
267 phase: ShutdownTaskPhase,
268 count: usize,
269}
270
271impl UnjoinedPhase {
272 #[must_use]
274 pub const fn phase(self) -> ShutdownTaskPhase {
275 self.phase
276 }
277
278 #[must_use]
280 pub const fn count(self) -> usize {
281 self.count
282 }
283}
284
285#[derive(Clone, Debug, Eq, PartialEq)]
287#[non_exhaustive]
288pub enum DrainResult {
289 Complete {
291 panicked_tasks: usize,
293 },
294 Incomplete {
296 unjoined_tasks: usize,
298 phases: Vec<UnjoinedPhase>,
300 panicked_tasks: usize,
302 escalated: bool,
304 },
305}
306
307#[derive(Clone, Copy, Debug, Eq, PartialEq)]
309#[non_exhaustive]
310pub enum ShutdownHookStatus {
311 Completed,
313 Failed,
315 DeadlineExceeded,
317}
318
319#[derive(Clone, Copy, Debug, Eq, PartialEq)]
321#[non_exhaustive]
322pub enum TelemetryFlushStatus {
323 Completed {
325 dropped_events: u64,
327 },
328 Failed,
330 DeadlineExceeded,
332}
333
334#[derive(Clone, Debug, Eq, PartialEq)]
336pub struct ShutdownReport {
337 drain: DrainResult,
338 persistence: ShutdownHookStatus,
339 telemetry: TelemetryFlushStatus,
340 repository_close: ShutdownHookStatus,
341}
342
343#[derive(Clone, Copy, Debug, Eq, PartialEq)]
345pub struct ShutdownHookError;
346
347impl fmt::Display for ShutdownHookError {
348 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
349 formatter.write_str("shutdown hook failed")
350 }
351}
352
353impl Error for ShutdownHookError {}
354
355impl ShutdownReport {
356 #[must_use]
358 pub const fn drain(&self) -> &DrainResult {
359 &self.drain
360 }
361
362 #[must_use]
364 pub const fn persistence(&self) -> ShutdownHookStatus {
365 self.persistence
366 }
367
368 #[must_use]
370 pub const fn telemetry(&self) -> TelemetryFlushStatus {
371 self.telemetry
372 }
373
374 #[must_use]
376 pub const fn repository_close(&self) -> ShutdownHookStatus {
377 self.repository_close
378 }
379}
380
381pub struct ShutdownCoordinator {
383 signal: ShutdownSignal,
384 shutdown_deadline: ShutdownDeadline,
385 task_join_deadline: TaskJoinDeadline,
386 telemetry_deadline: TelemetryFlushDeadline,
387 tasks: JoinSet<(ShutdownTaskPhase, bool)>,
388 phases: BTreeMap<ShutdownTaskPhase, usize>,
389 event_sink: Option<Arc<dyn TelemetryEventSink>>,
390}
391
392impl fmt::Debug for ShutdownCoordinator {
393 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
394 formatter
395 .debug_struct("ShutdownCoordinator")
396 .field("shutdown_deadline", &self.shutdown_deadline)
397 .field("task_join_deadline", &self.task_join_deadline)
398 .field("telemetry_deadline", &self.telemetry_deadline)
399 .field("owned_tasks", &self.tasks.len())
400 .finish_non_exhaustive()
401 }
402}
403
404impl ShutdownCoordinator {
405 pub fn new(
412 shutdown_deadline: ShutdownDeadline,
413 task_join_deadline: TaskJoinDeadline,
414 telemetry_deadline: TelemetryFlushDeadline,
415 ) -> Result<Self, ShutdownError> {
416 if task_join_deadline.get() > shutdown_deadline.get() {
417 return Err(ShutdownError::TaskJoinExceedsShutdown);
418 }
419 Ok(Self {
420 signal: ShutdownSignal::new(),
421 shutdown_deadline,
422 task_join_deadline,
423 telemetry_deadline,
424 tasks: JoinSet::new(),
425 phases: BTreeMap::new(),
426 event_sink: None,
427 })
428 }
429
430 #[must_use]
432 pub fn with_event_sink(mut self, sink: Arc<dyn TelemetryEventSink>) -> Self {
433 self.event_sink = Some(sink);
434 self
435 }
436
437 #[must_use]
439 pub fn signal(&self) -> ShutdownSignal {
440 self.signal.clone()
441 }
442
443 pub fn spawn<F>(&mut self, phase: ShutdownTaskPhase, future: F) -> Result<(), ShutdownError>
452 where
453 F: Future<Output = ()> + Send + 'static,
454 {
455 self.signal.ensure_accepting()?;
456 *self.phases.entry(phase).or_default() += 1;
457 self.tasks.spawn(async move {
458 let panicked = AssertUnwindSafe(future).catch_unwind().await.is_err();
459 (phase, panicked)
460 });
461 Ok(())
462 }
463
464 pub async fn shutdown<P, PF, T, TF, C, CF>(
473 &mut self,
474 persist: P,
475 flush_telemetry: T,
476 close_repository: C,
477 ) -> ShutdownReport
478 where
479 P: FnOnce() -> PF,
480 PF: Future<Output = Result<(), ShutdownHookError>>,
481 T: FnOnce() -> TF,
482 TF: Future<Output = Result<u64, ShutdownHookError>>,
483 C: FnOnce() -> CF,
484 CF: Future<Output = Result<(), ShutdownHookError>>,
485 {
486 self.signal.begin_shutdown();
489 crate::telemetry::emit_safely(
490 self.event_sink.as_ref(),
491 &TelemetryRecord::shutdown(TelemetryEventKind::ShutdownRequested, "requested", 0),
492 );
493 crate::telemetry::emit_safely(
494 self.event_sink.as_ref(),
495 &TelemetryRecord::shutdown(TelemetryEventKind::ShutdownIntakeStopped, "stopped", 0),
496 );
497 let started = tokio::time::Instant::now();
498 let correctness_end = started + self.shutdown_deadline.get();
499 let join_end = started + self.task_join_deadline.get();
500 let mut panicked_tasks = 0;
501 let mut escalated = false;
502
503 while !self.tasks.is_empty() {
504 tokio::select! {
505 joined = self.tasks.join_next() => {
506 if let Some(Ok((phase, panicked))) = joined {
507 panicked_tasks += usize::from(panicked);
508 decrement_phase(&mut self.phases, phase);
509 }
510 }
511 () = tokio::time::sleep_until(join_end) => break,
512 () = self.signal.escalated() => {
513 escalated = true;
514 break;
515 }
516 }
517 }
518
519 let drain = if self.tasks.is_empty() {
520 DrainResult::Complete { panicked_tasks }
521 } else {
522 DrainResult::Incomplete {
523 unjoined_tasks: self.tasks.len(),
524 phases: self
525 .phases
526 .iter()
527 .map(|(phase, count)| UnjoinedPhase {
528 phase: *phase,
529 count: *count,
530 })
531 .collect(),
532 panicked_tasks,
533 escalated,
534 }
535 };
536 match &drain {
537 DrainResult::Complete { .. } => crate::telemetry::emit_safely(
538 self.event_sink.as_ref(),
539 &TelemetryRecord::shutdown(
540 TelemetryEventKind::ShutdownDrainCompleted,
541 "complete",
542 0,
543 ),
544 ),
545 DrainResult::Incomplete { unjoined_tasks, .. } => crate::telemetry::emit_safely(
546 self.event_sink.as_ref(),
547 &TelemetryRecord::shutdown(
548 TelemetryEventKind::ShutdownDeadlineExceeded,
549 "incomplete",
550 *unjoined_tasks,
551 ),
552 ),
553 }
554
555 let persisted = persist().await;
560 let persistence = if tokio::time::Instant::now() > correctness_end {
561 ShutdownHookStatus::DeadlineExceeded
562 } else {
563 match persisted {
564 Ok(()) => ShutdownHookStatus::Completed,
565 Err(_) => ShutdownHookStatus::Failed,
566 }
567 };
568 let telemetry =
569 match tokio::time::timeout(self.telemetry_deadline.get(), flush_telemetry()).await {
570 Ok(Ok(dropped_events)) => TelemetryFlushStatus::Completed { dropped_events },
571 Ok(Err(_)) => TelemetryFlushStatus::Failed,
572 Err(_) => TelemetryFlushStatus::DeadlineExceeded,
573 };
574 let repository_close = match close_repository().await {
575 Ok(()) => ShutdownHookStatus::Completed,
576 Err(_) => ShutdownHookStatus::Failed,
577 };
578
579 ShutdownReport {
580 drain,
581 persistence,
582 telemetry,
583 repository_close,
584 }
585 }
586}
587
588impl Default for ShutdownCoordinator {
589 fn default() -> Self {
590 Self {
591 signal: ShutdownSignal::new(),
592 shutdown_deadline: ShutdownDeadline::default(),
593 task_join_deadline: TaskJoinDeadline::default(),
594 telemetry_deadline: TelemetryFlushDeadline::default(),
595 tasks: JoinSet::new(),
596 phases: BTreeMap::new(),
597 event_sink: None,
598 }
599 }
600}
601
602fn decrement_phase(phases: &mut BTreeMap<ShutdownTaskPhase, usize>, phase: ShutdownTaskPhase) {
603 if let Some(count) = phases.get_mut(&phase) {
604 *count = count.saturating_sub(1);
605 if *count == 0 {
606 phases.remove(&phase);
607 }
608 }
609}
610
611#[derive(Clone, Copy, Debug, Eq, PartialEq)]
613#[non_exhaustive]
614pub enum ShutdownError {
615 ShuttingDown,
617 InvalidShutdownDeadline,
619 InvalidTaskJoinDeadline,
621 TaskJoinExceedsShutdown,
623 InvalidTelemetryFlushDeadline,
625}
626
627impl fmt::Display for ShutdownError {
628 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
629 match self {
630 Self::ShuttingDown => formatter.write_str("runtime intake is shutting down"),
631 Self::InvalidShutdownDeadline => {
632 formatter.write_str("shutdown deadline must be between 1 second and 1 hour")
633 }
634 Self::InvalidTaskJoinDeadline => {
635 formatter.write_str("task join deadline must be between 1 second and 1 hour")
636 }
637 Self::TaskJoinExceedsShutdown => {
638 formatter.write_str("task join deadline cannot exceed shutdown deadline")
639 }
640 Self::InvalidTelemetryFlushDeadline => formatter
641 .write_str("telemetry flush deadline must be between 100 ms and 60 seconds"),
642 }
643 }
644}
645
646impl Error for ShutdownError {}
647
648#[cfg(test)]
649mod tests {
650 #![allow(clippy::expect_used)]
651 use std::sync::{Arc, Mutex};
652
653 use super::*;
654
655 #[test]
656 fn deadlines_enforce_accepted_bounds_and_relationship() {
657 let shutdown = ShutdownDeadline::new(Duration::from_secs(2)).expect("valid deadline");
658 assert_eq!(
659 TaskJoinDeadline::new(Duration::from_secs(3), shutdown),
660 Err(ShutdownError::TaskJoinExceedsShutdown)
661 );
662 assert_eq!(
663 TelemetryFlushDeadline::new(Duration::from_millis(99)),
664 Err(ShutdownError::InvalidTelemetryFlushDeadline)
665 );
666 }
667
668 #[test]
669 fn first_request_stops_intake_and_second_escalates() {
670 let coordinator = ShutdownCoordinator::default();
671 let signal = coordinator.signal();
672 assert_eq!(signal.ensure_accepting(), Ok(()));
673 assert_eq!(signal.request_shutdown(), ShutdownRequest::Initiated);
674 assert_eq!(signal.ensure_accepting(), Err(ShutdownError::ShuttingDown));
675 assert_eq!(signal.request_shutdown(), ShutdownRequest::Escalated);
676 assert!(signal.is_escalated());
677 }
678
679 #[tokio::test]
680 async fn phases_and_hooks_complete_in_fixed_order() {
681 let events = Arc::new(Mutex::new(Vec::new()));
682 let mut coordinator = ShutdownCoordinator::default();
683 let task_events = Arc::clone(&events);
684 coordinator
685 .spawn(ShutdownTaskPhase::Tasklet, async move {
686 task_events.lock().expect("events lock").push("task");
687 })
688 .expect("intake is open");
689
690 let persist_events = Arc::clone(&events);
691 let telemetry_events = Arc::clone(&events);
692 let close_events = Arc::clone(&events);
693 let report = coordinator
694 .shutdown(
695 || async move {
696 persist_events.lock().expect("events lock").push("persist");
697 Ok::<_, ShutdownHookError>(())
698 },
699 || async move {
700 telemetry_events
701 .lock()
702 .expect("events lock")
703 .push("telemetry");
704 Ok::<_, ShutdownHookError>(0)
705 },
706 || async move {
707 close_events.lock().expect("events lock").push("close");
708 Ok::<_, ShutdownHookError>(())
709 },
710 )
711 .await;
712
713 assert_eq!(report.drain(), &DrainResult::Complete { panicked_tasks: 0 });
714 assert_eq!(
715 *events.lock().expect("events lock"),
716 vec!["task", "persist", "telemetry", "close"]
717 );
718 }
719
720 #[tokio::test]
721 async fn an_existing_first_request_is_not_treated_as_escalation() {
722 let mut coordinator = ShutdownCoordinator::default();
723 let signal = coordinator.signal();
724 assert_eq!(signal.request_shutdown(), ShutdownRequest::Initiated);
725
726 let report = coordinator
727 .shutdown(
728 || async { Ok::<_, ShutdownHookError>(()) },
729 || async { Ok::<_, ShutdownHookError>(0) },
730 || async { Ok::<_, ShutdownHookError>(()) },
731 )
732 .await;
733
734 assert!(!signal.is_escalated());
735 assert_eq!(report.drain(), &DrainResult::Complete { panicked_tasks: 0 });
736 }
737
738 #[tokio::test]
739 async fn escalation_reports_every_unjoined_phase_without_detaching() {
740 let mut coordinator = ShutdownCoordinator::default();
741 coordinator
742 .spawn(ShutdownTaskPhase::Transaction, std::future::pending())
743 .expect("intake is open");
744 let signal = coordinator.signal();
745 let escalator = signal.clone();
746 tokio::spawn(async move {
747 escalator.cancelled().await;
748 let _ = escalator.request_shutdown();
749 });
750
751 let report = coordinator
752 .shutdown(
753 || async { Ok::<_, ShutdownHookError>(()) },
754 || async { Ok::<_, ShutdownHookError>(0) },
755 || async { Ok::<_, ShutdownHookError>(()) },
756 )
757 .await;
758
759 assert_eq!(
760 report.drain(),
761 &DrainResult::Incomplete {
762 unjoined_tasks: 1,
763 phases: vec![UnjoinedPhase {
764 phase: ShutdownTaskPhase::Transaction,
765 count: 1,
766 }],
767 panicked_tasks: 0,
768 escalated: true,
769 }
770 );
771 }
772}