From 89a35e9dbeb972b50008c7272a87f8eed5121d8c Mon Sep 17 00:00:00 2001 From: Vika Date: Sun, 9 Jul 2023 18:27:52 +0300 Subject: webmentions/supervisor: ability to cancel a background task This requires the background task to be cancellation-safe, as it is dropped after receiving a cancellation event. Perhaps in the future there will be a supervisor version that may forward the cancellation to the task itself. --- kittybox-rs/src/webmentions/mod.rs | 67 ++++++++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 13 deletions(-) (limited to 'kittybox-rs/src') diff --git a/kittybox-rs/src/webmentions/mod.rs b/kittybox-rs/src/webmentions/mod.rs index d798c50..cffd064 100644 --- a/kittybox-rs/src/webmentions/mod.rs +++ b/kittybox-rs/src/webmentions/mod.rs @@ -36,25 +36,66 @@ async fn accept_webmention>( } } -pub fn router>(queue: Q) -> axum::Router { +pub fn router>(queue: Q, cancellation_token: tokio_util::sync::CancellationToken) -> (axum::Router, SupervisedTask) { // Automatically spawn a background task to handle webmentions - tokio::task::spawn(supervised_webmentions_task(queue.clone())); + let bgtask_handle = supervised_webmentions_task(queue.clone(), cancellation_token); - axum::Router::new() + let router = axum::Router::new() .route("/.kittybox/webmention", axum::routing::post(accept_webmention::) ) - .layer(Extension(queue)) + .layer(Extension(queue)); + + (router, bgtask_handle) +} + +#[derive(thiserror::Error, Debug)] +pub enum SupervisorError { + #[error("the task was explicitly cancelled")] + Cancelled } -async fn supervisor> + Send + 'static, F: FnMut() -> A>(mut f: F) -> std::convert::Infallible { - loop { - match tokio::task::spawn(f()).await { - Err(e) => tracing::error!("background task exited unexpectedly: {}", e), - Ok(Err(e)) => tracing::error!("background task returned error: {}", e), - Ok(Ok(_)) => unreachable!("task's Ok is Infallible") +pub type SupervisedTask = tokio::task::JoinHandle>; + +pub fn supervisor(mut f: F, cancellation_token: tokio_util::sync::CancellationToken) -> SupervisedTask +where + E: std::error::Error + std::fmt::Debug + Send + 'static, + A: std::future::Future> + Send + 'static, + F: FnMut() -> A + Send + 'static +{ + + let supervisor_future = async move { + loop { + // Don't spawn the task if we are already cancelled, but + // have somehow missed it (probably because the task + // crashed and we immediately received a cancellation + // request after noticing the crashed task) + if cancellation_token.is_cancelled() { + return Err(SupervisorError::Cancelled) + } + let task = tokio::task::spawn(f()); + tokio::select! { + _ = cancellation_token.cancelled() => { + tracing::info!("Shutdown of background task {:?} requested.", std::any::type_name::()); + return Err(SupervisorError::Cancelled) + } + task_result = task => match task_result { + Err(e) => tracing::error!("background task {:?} exited unexpectedly: {}", std::any::type_name::(), e), + Ok(Err(e)) => tracing::error!("background task {:?} returned error: {}", std::any::type_name::(), e), + Ok(Ok(_)) => unreachable!("task's Ok is Infallible") + } + } + tracing::debug!("Sleeping for a little while to back-off..."); + tokio::time::sleep(std::time::Duration::from_secs(5)).await; } - } + }; + #[cfg(not(tokio_unstable))] + return tokio::task::spawn(supervisor_future); + #[cfg(tokio_unstable)] + return tokio::task::Builder::new() + .name(format!("supervisor for background task {}", std::any::type_name::()).as_str()) + .spawn(supervisor_future) + .unwrap(); } async fn process_webmentions_from_queue>(queue: Q) -> Result { @@ -74,6 +115,6 @@ async fn process_webmentions_from_queue>(queue: Q) -> Re unreachable!() } -async fn supervised_webmentions_task>(queue: Q) { - supervisor::(|| process_webmentions_from_queue(queue.clone())).await; +fn supervised_webmentions_task>(queue: Q, cancellation_token: tokio_util::sync::CancellationToken) -> SupervisedTask { + supervisor::(move || process_webmentions_from_queue(queue.clone()), cancellation_token) } -- cgit 1.4.1