#![deny(missing_docs)]
use std::{
fmt,
future::Future,
io,
panic::{AssertUnwindSafe, catch_unwind},
sync::{Arc, Mutex, OnceLock},
time::Duration,
};
use axum::Router;
pub use baukit_core::{BuildInfo, DeploymentEnvironment, ProcessKind, ServiceIdentity};
use thiserror::Error;
use tokio::{
net::TcpListener,
task::{JoinError, JoinHandle, JoinSet},
time::{Instant, timeout},
};
use tokio_util::sync::CancellationToken;
pub const RUST_VERSION: &str = env!("BAUKIT_RUST_VERSION");
#[macro_export]
macro_rules! build_info {
() => {
$crate::BuildInfo::new(
env!("CARGO_PKG_VERSION"),
option_env!("GIT_COMMIT").unwrap_or("unknown"),
$crate::RUST_VERSION,
)
};
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ServiceInfo {
name: String,
identity: ServiceIdentity,
build: BuildInfo,
}
impl ServiceInfo {
pub fn new(
product: impl Into<String>,
process: ProcessKind,
build: BuildInfo,
environment: DeploymentEnvironment,
) -> Self {
let product = product.into();
let identity = ServiceIdentity::new(
product,
process,
build.version().to_owned(),
build.commit().to_owned(),
environment,
);
let name = identity.service_name();
Self {
name,
identity,
build,
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn product(&self) -> &str {
self.identity.product()
}
pub const fn process(&self) -> ProcessKind {
self.identity.process()
}
pub const fn build(&self) -> &BuildInfo {
&self.build
}
pub const fn telemetry_identity(&self) -> &ServiceIdentity {
&self.identity
}
}
struct ShutdownState {
root: CancellationToken,
started_at: OnceLock<Instant>,
drain_hooks: Mutex<Vec<DrainHook>>,
}
type DrainHook = Box<dyn FnOnce() + Send + 'static>;
impl fmt::Debug for ShutdownState {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
let hook_count = self
.drain_hooks
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.len();
formatter
.debug_struct("ShutdownState")
.field("root", &self.root)
.field("started_at", &self.started_at)
.field("drain_hook_count", &hook_count)
.finish()
}
}
#[derive(Clone, Debug)]
pub struct ShutdownToken {
observed: CancellationToken,
state: Arc<ShutdownState>,
drain_timeout: Duration,
}
impl ShutdownToken {
pub fn new(drain_timeout: Duration) -> Self {
let root = CancellationToken::new();
Self {
observed: root.clone(),
state: Arc::new(ShutdownState {
root,
started_at: OnceLock::new(),
drain_hooks: Mutex::new(Vec::new()),
}),
drain_timeout,
}
}
#[must_use]
pub fn child_token(&self) -> Self {
Self {
observed: self.observed.child_token(),
state: Arc::clone(&self.state),
drain_timeout: self.drain_timeout,
}
}
pub fn trigger(&self) -> bool {
let first = self.state.started_at.set(Instant::now()).is_ok();
if first {
let hooks = std::mem::take(
&mut *self
.state
.drain_hooks
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner),
);
for hook in hooks {
if catch_unwind(AssertUnwindSafe(hook)).is_err() {
tracing::error!("a shutdown drain hook panicked");
}
}
self.state.root.cancel();
}
first
}
pub fn on_drain<F>(&self, hook: F)
where
F: FnOnce() + Send + 'static,
{
let hook: DrainHook = Box::new(hook);
let late_hook = {
let mut hooks = self
.state
.drain_hooks
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if self.state.started_at.get().is_none() {
hooks.push(hook);
None
} else {
Some(hook)
}
};
if let Some(hook) = late_hook
&& catch_unwind(AssertUnwindSafe(hook)).is_err()
{
tracing::error!("a late shutdown drain hook panicked");
}
}
pub fn is_cancelled(&self) -> bool {
self.observed.is_cancelled()
}
pub async fn cancelled(&self) {
self.observed.cancelled().await;
}
pub const fn drain_timeout(&self) -> Duration {
self.drain_timeout
}
pub fn deadline(&self) -> Option<Instant> {
self.state
.started_at
.get()
.map(|started_at| *started_at + self.drain_timeout)
}
pub async fn run_during_drain<F>(&self, future: F) -> Result<F::Output, DrainTimeout>
where
F: Future,
{
self.cancelled().await;
let remaining = self.remaining();
timeout(remaining, future)
.await
.map_err(|_| DrainTimeout::new(self.drain_timeout))
}
pub async fn shutdown_on_signal(&self) -> io::Result<ShutdownSignal> {
let signal = wait_for_shutdown_signal().await?;
self.trigger();
Ok(signal)
}
#[must_use = "retain or join the signal listener so setup failures are observed"]
pub fn spawn_signal_listener(&self) -> tokio::task::JoinHandle<io::Result<ShutdownSignal>> {
let shutdown = self.clone();
tokio::spawn(async move { shutdown.shutdown_on_signal().await })
}
fn remaining(&self) -> Duration {
self.deadline()
.map(|deadline| deadline.saturating_duration_since(Instant::now()))
.unwrap_or(self.drain_timeout)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ShutdownSignal {
Interrupt,
Terminate,
}
#[cfg(unix)]
pub async fn wait_for_shutdown_signal() -> io::Result<ShutdownSignal> {
use tokio::signal::unix::{SignalKind, signal};
let mut terminate = signal(SignalKind::terminate())?;
tokio::select! {
ctrl_c = tokio::signal::ctrl_c() => {
ctrl_c?;
Ok(ShutdownSignal::Interrupt)
}
received = terminate.recv() => received.map_or_else(
|| Err(io::Error::other("SIGTERM signal stream ended")),
|_| Ok(ShutdownSignal::Terminate),
),
}
}
#[cfg(not(unix))]
pub async fn wait_for_shutdown_signal() -> io::Result<ShutdownSignal> {
tokio::signal::ctrl_c().await?;
Ok(ShutdownSignal::Interrupt)
}
#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
#[error("shutdown drain deadline expired after {timeout:?}")]
pub struct DrainTimeout {
timeout: Duration,
}
impl DrainTimeout {
const fn new(timeout: Duration) -> Self {
Self { timeout }
}
pub const fn timeout(self) -> Duration {
self.timeout
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RestartPolicy {
FailProcess,
Restart {
max_restarts: usize,
},
}
pub struct TaskSupervisor {
shutdown: ShutdownToken,
tasks: JoinSet<()>,
}
impl TaskSupervisor {
pub fn new(shutdown: ShutdownToken) -> Self {
Self {
shutdown,
tasks: JoinSet::new(),
}
}
pub fn spawn<F, Fut>(&mut self, name: impl Into<String>, policy: RestartPolicy, factory: F)
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let shutdown = self.shutdown.clone();
let name = name.into();
let factory = Arc::new(factory);
self.tasks.spawn(async move {
supervise_task(name, policy, factory, shutdown).await;
});
}
pub async fn join(mut self) -> Result<(), DrainTimeout> {
self.shutdown.cancelled().await;
let remaining = self.shutdown.remaining();
let joined = timeout(remaining, async {
while let Some(result) = self.tasks.join_next().await {
if let Err(error) = result {
tracing::error!(%error, "task supervisor failed while joining a worker");
}
}
})
.await;
if joined.is_err() {
self.tasks.abort_all();
while self.tasks.join_next().await.is_some() {}
return Err(DrainTimeout::new(self.shutdown.drain_timeout));
}
Ok(())
}
pub fn len(&self) -> usize {
self.tasks.len()
}
pub fn is_empty(&self) -> bool {
self.tasks.is_empty()
}
}
async fn supervise_task<F, Fut>(
name: String,
policy: RestartPolicy,
factory: Arc<F>,
shutdown: ShutdownToken,
) where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let mut restarts = 0_usize;
loop {
let mut attempt = AbortOnDropTask(tokio::spawn(factory()));
let outcome = tokio::select! {
_ = shutdown.cancelled() => {
if let Err(error) = (&mut attempt.0).await
&& error.is_panic()
{
tracing::error!(task = %name, %error, "supervised task panicked while draining");
}
return;
}
result = &mut attempt.0 => result,
};
if shutdown.is_cancelled() {
return;
}
match &outcome {
Ok(()) => tracing::error!(task = %name, "supervised task exited unexpectedly"),
Err(error) if error.is_panic() => {
tracing::error!(task = %name, %error, "supervised task panicked")
}
Err(error) => tracing::error!(task = %name, %error, "supervised task was cancelled"),
}
match policy {
RestartPolicy::FailProcess => {
shutdown.trigger();
return;
}
RestartPolicy::Restart { max_restarts } if restarts < max_restarts => {
restarts += 1;
tracing::warn!(task = %name, restarts, max_restarts, "restarting supervised task");
}
RestartPolicy::Restart { max_restarts } => {
tracing::error!(task = %name, restarts, max_restarts, "task restart limit exhausted");
shutdown.trigger();
return;
}
}
}
}
struct AbortOnDropTask<T>(JoinHandle<T>);
impl<T> Drop for AbortOnDropTask<T> {
fn drop(&mut self) {
self.0.abort();
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ListenerKind {
Api,
Operations,
}
impl fmt::Display for ListenerKind {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Api => "API",
Self::Operations => "operations",
})
}
}
#[derive(Debug, Error)]
pub enum ListenerPairError {
#[error("{listener} listener failed: {source}")]
Listener {
listener: ListenerKind,
#[source]
source: io::Error,
},
#[error("listener task failed: {0}")]
Task(#[source] JoinError),
#[error("{0} listener exited before shutdown")]
UnexpectedExit(ListenerKind),
#[error(transparent)]
Drain(#[from] DrainTimeout),
}
type ListenerResult = Result<ListenerKind, (ListenerKind, io::Error)>;
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum ShutdownOrder {
#[default]
Together,
OpsOutlivesApi,
}
pub async fn serve_listener_pair(
api_listener: TcpListener,
api_router: Router,
operations_listener: TcpListener,
operations_router: Router,
shutdown: ShutdownToken,
) -> Result<(), ListenerPairError> {
serve_listener_pair_with_shutdown_order(
api_listener,
api_router,
operations_listener,
operations_router,
shutdown,
ShutdownOrder::Together,
)
.await
}
pub async fn serve_listener_pair_with_shutdown_order(
api_listener: TcpListener,
api_router: Router,
operations_listener: TcpListener,
operations_router: Router,
shutdown: ShutdownToken,
shutdown_order: ShutdownOrder,
) -> Result<(), ListenerPairError> {
let mut listeners = JoinSet::new();
let operations_shutdown = match shutdown_order {
ShutdownOrder::Together => shutdown.observed.clone(),
ShutdownOrder::OpsOutlivesApi => CancellationToken::new(),
};
spawn_listener(
&mut listeners,
ListenerKind::Api,
api_listener,
api_router,
shutdown.observed.clone(),
);
spawn_listener(
&mut listeners,
ListenerKind::Operations,
operations_listener,
operations_router,
operations_shutdown.clone(),
);
let mut first_error = None;
while !listeners.is_empty() && !shutdown.is_cancelled() {
tokio::select! {
_ = shutdown.cancelled() => {}
result = listeners.join_next() => {
stop_operations_after_api(
result.as_ref(),
shutdown_order,
&operations_shutdown,
);
record_listener_result(result, &shutdown, &mut first_error);
}
}
}
if listeners.is_empty() {
return first_error.map_or(Ok(()), Err);
}
let remaining = shutdown.remaining();
let drained = timeout(remaining, async {
while let Some(result) = listeners.join_next().await {
stop_operations_after_api(Some(&result), shutdown_order, &operations_shutdown);
record_listener_result(Some(result), &shutdown, &mut first_error);
}
})
.await;
if drained.is_err() {
listeners.abort_all();
while listeners.join_next().await.is_some() {}
return Err(DrainTimeout::new(shutdown.drain_timeout).into());
}
first_error.map_or(Ok(()), Err)
}
fn spawn_listener(
listeners: &mut JoinSet<ListenerResult>,
kind: ListenerKind,
listener: TcpListener,
router: Router,
shutdown: CancellationToken,
) {
listeners.spawn(async move {
axum::serve(listener, router)
.with_graceful_shutdown(async move { shutdown.cancelled().await })
.await
.map(|()| kind)
.map_err(|error| (kind, error))
});
}
fn stop_operations_after_api(
result: Option<&Result<ListenerResult, JoinError>>,
shutdown_order: ShutdownOrder,
operations_shutdown: &CancellationToken,
) {
if shutdown_order != ShutdownOrder::OpsOutlivesApi {
return;
}
let listener = match result {
Some(Ok(Ok(kind) | Err((kind, _)))) => Some(*kind),
Some(Err(_)) | None => None,
};
if listener == Some(ListenerKind::Api) {
operations_shutdown.cancel();
}
}
fn record_listener_result(
result: Option<Result<ListenerResult, JoinError>>,
shutdown: &ShutdownToken,
first_error: &mut Option<ListenerPairError>,
) {
let error = match result {
Some(Ok(Ok(kind))) if !shutdown.is_cancelled() => {
Some(ListenerPairError::UnexpectedExit(kind))
}
Some(Ok(Ok(_))) | None => None,
Some(Ok(Err((listener, source)))) => Some(ListenerPairError::Listener { listener, source }),
Some(Err(error)) => Some(ListenerPairError::Task(error)),
};
if let Some(error) = error {
tracing::error!(%error, "listener pair initiated process shutdown");
if first_error.is_none() {
*first_error = Some(error);
}
shutdown.trigger();
}
}
#[cfg(test)]
mod tests;
#[doc = include_str!("../README.md")]
#[cfg(doctest)]
struct ReadmeDoctests;