1#![deny(missing_docs)]
76
77use std::{
78 fmt,
79 future::Future,
80 io,
81 panic::{AssertUnwindSafe, catch_unwind},
82 sync::{Arc, Mutex, OnceLock},
83 time::Duration,
84};
85
86use axum::Router;
87pub use baukit_core::{BuildInfo, DeploymentEnvironment, ProcessKind, ServiceIdentity};
88use thiserror::Error;
89use tokio::{
90 net::TcpListener,
91 task::{JoinError, JoinHandle, JoinSet},
92 time::{Instant, timeout},
93};
94use tokio_util::sync::CancellationToken;
95
96pub const RUST_VERSION: &str = env!("BAUKIT_RUST_VERSION");
100
101#[macro_export]
107macro_rules! build_info {
108 () => {
109 $crate::BuildInfo::new(
110 env!("CARGO_PKG_VERSION"),
111 option_env!("GIT_COMMIT").unwrap_or("unknown"),
112 $crate::RUST_VERSION,
113 )
114 };
115}
116
117#[derive(Clone, Debug, Eq, PartialEq)]
123pub struct ServiceInfo {
124 name: String,
125 identity: ServiceIdentity,
126 build: BuildInfo,
127}
128
129impl ServiceInfo {
130 pub fn new(
132 product: impl Into<String>,
133 process: ProcessKind,
134 build: BuildInfo,
135 environment: DeploymentEnvironment,
136 ) -> Self {
137 let product = product.into();
138 let identity = ServiceIdentity::new(
139 product,
140 process,
141 build.version().to_owned(),
142 build.commit().to_owned(),
143 environment,
144 );
145 let name = identity.service_name();
146 Self {
147 name,
148 identity,
149 build,
150 }
151 }
152
153 pub fn name(&self) -> &str {
155 &self.name
156 }
157
158 pub fn product(&self) -> &str {
160 self.identity.product()
161 }
162
163 pub const fn process(&self) -> ProcessKind {
165 self.identity.process()
166 }
167
168 pub const fn build(&self) -> &BuildInfo {
170 &self.build
171 }
172
173 pub const fn telemetry_identity(&self) -> &ServiceIdentity {
175 &self.identity
176 }
177}
178
179struct ShutdownState {
180 root: CancellationToken,
181 started_at: OnceLock<Instant>,
182 drain_hooks: Mutex<Vec<DrainHook>>,
183}
184
185type DrainHook = Box<dyn FnOnce() + Send + 'static>;
186
187impl fmt::Debug for ShutdownState {
188 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
189 let hook_count = self
190 .drain_hooks
191 .lock()
192 .unwrap_or_else(std::sync::PoisonError::into_inner)
193 .len();
194 formatter
195 .debug_struct("ShutdownState")
196 .field("root", &self.root)
197 .field("started_at", &self.started_at)
198 .field("drain_hook_count", &hook_count)
199 .finish()
200 }
201}
202
203#[derive(Clone, Debug)]
209pub struct ShutdownToken {
210 observed: CancellationToken,
211 state: Arc<ShutdownState>,
212 drain_timeout: Duration,
213}
214
215impl ShutdownToken {
216 pub fn new(drain_timeout: Duration) -> Self {
218 let root = CancellationToken::new();
219 Self {
220 observed: root.clone(),
221 state: Arc::new(ShutdownState {
222 root,
223 started_at: OnceLock::new(),
224 drain_hooks: Mutex::new(Vec::new()),
225 }),
226 drain_timeout,
227 }
228 }
229
230 #[must_use]
234 pub fn child_token(&self) -> Self {
235 Self {
236 observed: self.observed.child_token(),
237 state: Arc::clone(&self.state),
238 drain_timeout: self.drain_timeout,
239 }
240 }
241
242 pub fn trigger(&self) -> bool {
244 let first = self.state.started_at.set(Instant::now()).is_ok();
245 if first {
246 let hooks = std::mem::take(
247 &mut *self
248 .state
249 .drain_hooks
250 .lock()
251 .unwrap_or_else(std::sync::PoisonError::into_inner),
252 );
253 for hook in hooks {
254 if catch_unwind(AssertUnwindSafe(hook)).is_err() {
255 tracing::error!("a shutdown drain hook panicked");
256 }
257 }
258 self.state.root.cancel();
259 }
260 first
261 }
262
263 pub fn on_drain<F>(&self, hook: F)
287 where
288 F: FnOnce() + Send + 'static,
289 {
290 let hook: DrainHook = Box::new(hook);
291 let late_hook = {
292 let mut hooks = self
293 .state
294 .drain_hooks
295 .lock()
296 .unwrap_or_else(std::sync::PoisonError::into_inner);
297 if self.state.started_at.get().is_none() {
298 hooks.push(hook);
299 None
300 } else {
301 Some(hook)
302 }
303 };
304 if let Some(hook) = late_hook
305 && catch_unwind(AssertUnwindSafe(hook)).is_err()
306 {
307 tracing::error!("a late shutdown drain hook panicked");
308 }
309 }
310
311 pub fn is_cancelled(&self) -> bool {
313 self.observed.is_cancelled()
314 }
315
316 pub async fn cancelled(&self) {
318 self.observed.cancelled().await;
319 }
320
321 pub const fn drain_timeout(&self) -> Duration {
323 self.drain_timeout
324 }
325
326 pub fn deadline(&self) -> Option<Instant> {
328 self.state
329 .started_at
330 .get()
331 .map(|started_at| *started_at + self.drain_timeout)
332 }
333
334 pub async fn run_during_drain<F>(&self, future: F) -> Result<F::Output, DrainTimeout>
344 where
345 F: Future,
346 {
347 self.cancelled().await;
348 let remaining = self.remaining();
349 timeout(remaining, future)
350 .await
351 .map_err(|_| DrainTimeout::new(self.drain_timeout))
352 }
353
354 pub async fn shutdown_on_signal(&self) -> io::Result<ShutdownSignal> {
360 let signal = wait_for_shutdown_signal().await?;
361 self.trigger();
362 Ok(signal)
363 }
364
365 #[must_use = "retain or join the signal listener so setup failures are observed"]
369 pub fn spawn_signal_listener(&self) -> tokio::task::JoinHandle<io::Result<ShutdownSignal>> {
370 let shutdown = self.clone();
371 tokio::spawn(async move { shutdown.shutdown_on_signal().await })
372 }
373
374 fn remaining(&self) -> Duration {
375 self.deadline()
376 .map(|deadline| deadline.saturating_duration_since(Instant::now()))
377 .unwrap_or(self.drain_timeout)
378 }
379}
380
381#[derive(Clone, Copy, Debug, Eq, PartialEq)]
383pub enum ShutdownSignal {
384 Interrupt,
386 Terminate,
388}
389
390#[cfg(unix)]
400pub async fn wait_for_shutdown_signal() -> io::Result<ShutdownSignal> {
401 use tokio::signal::unix::{SignalKind, signal};
402
403 let mut terminate = signal(SignalKind::terminate())?;
404 tokio::select! {
405 ctrl_c = tokio::signal::ctrl_c() => {
406 ctrl_c?;
407 Ok(ShutdownSignal::Interrupt)
408 }
409 received = terminate.recv() => received.map_or_else(
410 || Err(io::Error::other("SIGTERM signal stream ended")),
411 |_| Ok(ShutdownSignal::Terminate),
412 ),
413 }
414}
415
416#[cfg(not(unix))]
422pub async fn wait_for_shutdown_signal() -> io::Result<ShutdownSignal> {
423 tokio::signal::ctrl_c().await?;
424 Ok(ShutdownSignal::Interrupt)
425}
426
427#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
429#[error("shutdown drain deadline expired after {timeout:?}")]
430pub struct DrainTimeout {
431 timeout: Duration,
432}
433
434impl DrainTimeout {
435 const fn new(timeout: Duration) -> Self {
436 Self { timeout }
437 }
438
439 pub const fn timeout(self) -> Duration {
441 self.timeout
442 }
443}
444
445#[derive(Clone, Copy, Debug, Eq, PartialEq)]
447pub enum RestartPolicy {
448 FailProcess,
450 Restart {
452 max_restarts: usize,
454 },
455}
456
457pub struct TaskSupervisor {
459 shutdown: ShutdownToken,
460 tasks: JoinSet<()>,
461}
462
463impl TaskSupervisor {
464 pub fn new(shutdown: ShutdownToken) -> Self {
466 Self {
467 shutdown,
468 tasks: JoinSet::new(),
469 }
470 }
471
472 pub fn spawn<F, Fut>(&mut self, name: impl Into<String>, policy: RestartPolicy, factory: F)
477 where
478 F: Fn() -> Fut + Send + Sync + 'static,
479 Fut: Future<Output = ()> + Send + 'static,
480 {
481 let shutdown = self.shutdown.clone();
482 let name = name.into();
483 let factory = Arc::new(factory);
484 self.tasks.spawn(async move {
485 supervise_task(name, policy, factory, shutdown).await;
486 });
487 }
488
489 pub async fn join(mut self) -> Result<(), DrainTimeout> {
498 self.shutdown.cancelled().await;
499 let remaining = self.shutdown.remaining();
500 let joined = timeout(remaining, async {
501 while let Some(result) = self.tasks.join_next().await {
502 if let Err(error) = result {
503 tracing::error!(%error, "task supervisor failed while joining a worker");
504 }
505 }
506 })
507 .await;
508
509 if joined.is_err() {
510 self.tasks.abort_all();
511 while self.tasks.join_next().await.is_some() {}
512 return Err(DrainTimeout::new(self.shutdown.drain_timeout));
513 }
514
515 Ok(())
516 }
517
518 pub fn len(&self) -> usize {
520 self.tasks.len()
521 }
522
523 pub fn is_empty(&self) -> bool {
525 self.tasks.is_empty()
526 }
527}
528
529async fn supervise_task<F, Fut>(
530 name: String,
531 policy: RestartPolicy,
532 factory: Arc<F>,
533 shutdown: ShutdownToken,
534) where
535 F: Fn() -> Fut + Send + Sync + 'static,
536 Fut: Future<Output = ()> + Send + 'static,
537{
538 let mut restarts = 0_usize;
539 loop {
540 let mut attempt = AbortOnDropTask(tokio::spawn(factory()));
541 let outcome = tokio::select! {
542 _ = shutdown.cancelled() => {
543 if let Err(error) = (&mut attempt.0).await
544 && error.is_panic()
545 {
546 tracing::error!(task = %name, %error, "supervised task panicked while draining");
547 }
548 return;
549 }
550 result = &mut attempt.0 => result,
551 };
552
553 if shutdown.is_cancelled() {
554 return;
555 }
556
557 match &outcome {
558 Ok(()) => tracing::error!(task = %name, "supervised task exited unexpectedly"),
559 Err(error) if error.is_panic() => {
560 tracing::error!(task = %name, %error, "supervised task panicked")
561 }
562 Err(error) => tracing::error!(task = %name, %error, "supervised task was cancelled"),
563 }
564
565 match policy {
566 RestartPolicy::FailProcess => {
567 shutdown.trigger();
568 return;
569 }
570 RestartPolicy::Restart { max_restarts } if restarts < max_restarts => {
571 restarts += 1;
572 tracing::warn!(task = %name, restarts, max_restarts, "restarting supervised task");
573 }
574 RestartPolicy::Restart { max_restarts } => {
575 tracing::error!(task = %name, restarts, max_restarts, "task restart limit exhausted");
576 shutdown.trigger();
577 return;
578 }
579 }
580 }
581}
582
583struct AbortOnDropTask<T>(JoinHandle<T>);
584
585impl<T> Drop for AbortOnDropTask<T> {
586 fn drop(&mut self) {
587 self.0.abort();
588 }
589}
590
591#[derive(Clone, Copy, Debug, Eq, PartialEq)]
593pub enum ListenerKind {
594 Api,
596 Operations,
598}
599
600impl fmt::Display for ListenerKind {
601 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
602 formatter.write_str(match self {
603 Self::Api => "API",
604 Self::Operations => "operations",
605 })
606 }
607}
608
609#[derive(Debug, Error)]
611pub enum ListenerPairError {
612 #[error("{listener} listener failed: {source}")]
614 Listener {
615 listener: ListenerKind,
617 #[source]
619 source: io::Error,
620 },
621 #[error("listener task failed: {0}")]
623 Task(#[source] JoinError),
624 #[error("{0} listener exited before shutdown")]
626 UnexpectedExit(ListenerKind),
627 #[error(transparent)]
629 Drain(#[from] DrainTimeout),
630}
631
632type ListenerResult = Result<ListenerKind, (ListenerKind, io::Error)>;
633
634#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
636pub enum ShutdownOrder {
637 #[default]
639 Together,
640 OpsOutlivesApi,
642}
643
644pub async fn serve_listener_pair(
656 api_listener: TcpListener,
657 api_router: Router,
658 operations_listener: TcpListener,
659 operations_router: Router,
660 shutdown: ShutdownToken,
661) -> Result<(), ListenerPairError> {
662 serve_listener_pair_with_shutdown_order(
663 api_listener,
664 api_router,
665 operations_listener,
666 operations_router,
667 shutdown,
668 ShutdownOrder::Together,
669 )
670 .await
671}
672
673pub async fn serve_listener_pair_with_shutdown_order(
689 api_listener: TcpListener,
690 api_router: Router,
691 operations_listener: TcpListener,
692 operations_router: Router,
693 shutdown: ShutdownToken,
694 shutdown_order: ShutdownOrder,
695) -> Result<(), ListenerPairError> {
696 let mut listeners = JoinSet::new();
697 let operations_shutdown = match shutdown_order {
698 ShutdownOrder::Together => shutdown.observed.clone(),
699 ShutdownOrder::OpsOutlivesApi => CancellationToken::new(),
700 };
701 spawn_listener(
702 &mut listeners,
703 ListenerKind::Api,
704 api_listener,
705 api_router,
706 shutdown.observed.clone(),
707 );
708 spawn_listener(
709 &mut listeners,
710 ListenerKind::Operations,
711 operations_listener,
712 operations_router,
713 operations_shutdown.clone(),
714 );
715
716 let mut first_error = None;
717 while !listeners.is_empty() && !shutdown.is_cancelled() {
718 tokio::select! {
719 _ = shutdown.cancelled() => {}
720 result = listeners.join_next() => {
721 stop_operations_after_api(
722 result.as_ref(),
723 shutdown_order,
724 &operations_shutdown,
725 );
726 record_listener_result(result, &shutdown, &mut first_error);
727 }
728 }
729 }
730
731 if listeners.is_empty() {
732 return first_error.map_or(Ok(()), Err);
733 }
734
735 let remaining = shutdown.remaining();
736 let drained = timeout(remaining, async {
737 while let Some(result) = listeners.join_next().await {
738 stop_operations_after_api(Some(&result), shutdown_order, &operations_shutdown);
739 record_listener_result(Some(result), &shutdown, &mut first_error);
740 }
741 })
742 .await;
743
744 if drained.is_err() {
745 listeners.abort_all();
746 while listeners.join_next().await.is_some() {}
747 return Err(DrainTimeout::new(shutdown.drain_timeout).into());
748 }
749
750 first_error.map_or(Ok(()), Err)
751}
752
753fn spawn_listener(
754 listeners: &mut JoinSet<ListenerResult>,
755 kind: ListenerKind,
756 listener: TcpListener,
757 router: Router,
758 shutdown: CancellationToken,
759) {
760 listeners.spawn(async move {
761 axum::serve(listener, router)
762 .with_graceful_shutdown(async move { shutdown.cancelled().await })
763 .await
764 .map(|()| kind)
765 .map_err(|error| (kind, error))
766 });
767}
768
769fn stop_operations_after_api(
770 result: Option<&Result<ListenerResult, JoinError>>,
771 shutdown_order: ShutdownOrder,
772 operations_shutdown: &CancellationToken,
773) {
774 if shutdown_order != ShutdownOrder::OpsOutlivesApi {
775 return;
776 }
777
778 let listener = match result {
779 Some(Ok(Ok(kind) | Err((kind, _)))) => Some(*kind),
780 Some(Err(_)) | None => None,
781 };
782 if listener == Some(ListenerKind::Api) {
783 operations_shutdown.cancel();
784 }
785}
786
787fn record_listener_result(
788 result: Option<Result<ListenerResult, JoinError>>,
789 shutdown: &ShutdownToken,
790 first_error: &mut Option<ListenerPairError>,
791) {
792 let error = match result {
793 Some(Ok(Ok(kind))) if !shutdown.is_cancelled() => {
794 Some(ListenerPairError::UnexpectedExit(kind))
795 }
796 Some(Ok(Ok(_))) | None => None,
797 Some(Ok(Err((listener, source)))) => Some(ListenerPairError::Listener { listener, source }),
798 Some(Err(error)) => Some(ListenerPairError::Task(error)),
799 };
800
801 if let Some(error) = error {
802 tracing::error!(%error, "listener pair initiated process shutdown");
803 if first_error.is_none() {
804 *first_error = Some(error);
805 }
806 shutdown.trigger();
807 }
808}
809
810#[cfg(test)]
811mod tests;
812
813#[doc = include_str!("../README.md")]
815#[cfg(doctest)]
816struct ReadmeDoctests;