use std::future::Future;
use tokio_util::sync::CancellationToken;
#[derive(Debug, Clone, Default)]
pub struct ShutdownHandle {
token: CancellationToken,
}
impl ShutdownHandle {
pub fn new() -> Self {
Self {
token: CancellationToken::new(),
}
}
pub fn from_token(token: CancellationToken) -> Self {
Self { token }
}
pub fn shutdown(&self) {
self.token.cancel();
}
pub fn is_shutdown_requested(&self) -> bool {
self.token.is_cancelled()
}
pub fn cancelled(&self) -> impl Future<Output = ()> + Send + 'static {
self.token.clone().cancelled_owned()
}
pub fn token(&self) -> CancellationToken {
self.token.clone()
}
}
impl From<CancellationToken> for ShutdownHandle {
fn from(token: CancellationToken) -> Self {
Self::from_token(token)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn it_starts_in_not_shutdown_state() {
let handle = ShutdownHandle::new();
assert!(!handle.is_shutdown_requested());
}
#[test]
fn it_reports_shutdown_after_trigger() {
let handle = ShutdownHandle::new();
handle.shutdown();
assert!(handle.is_shutdown_requested());
}
#[test]
fn it_is_idempotent_on_repeated_shutdown() {
let handle = ShutdownHandle::new();
handle.shutdown();
handle.shutdown();
assert!(handle.is_shutdown_requested());
}
#[test]
fn it_shares_state_across_clones() {
let original = ShutdownHandle::new();
let cloned = original.clone();
cloned.shutdown();
assert!(original.is_shutdown_requested());
}
#[tokio::test]
async fn it_resolves_cancelled_after_shutdown() {
let handle = ShutdownHandle::new();
let waiter = handle.clone();
let task = tokio::spawn(async move { waiter.cancelled().await });
handle.shutdown();
task.await.unwrap();
}
#[test]
fn it_returns_a_clone_of_the_underlying_token() {
let handle = ShutdownHandle::new();
let token = handle.token();
token.cancel();
assert!(handle.is_shutdown_requested());
}
#[test]
fn it_constructs_from_existing_token() {
let token = CancellationToken::new();
let handle = ShutdownHandle::from_token(token.clone());
token.cancel();
assert!(handle.is_shutdown_requested());
}
#[test]
fn it_constructs_via_from_impl() {
let token = CancellationToken::new();
let handle: ShutdownHandle = token.clone().into();
token.cancel();
assert!(handle.is_shutdown_requested());
}
#[test]
fn it_yields_a_fresh_handle_when_defaulted() {
let handle = ShutdownHandle::default();
assert!(!handle.is_shutdown_requested());
}
}