diff --git a/.github/workflows/audit.yml b/.github/workflows/audit.yml index 7bfbcb7..6e51373 100644 --- a/.github/workflows/audit.yml +++ b/.github/workflows/audit.yml @@ -12,6 +12,6 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - - uses: actions-rs/audit-check@v1 + - uses: rustsec/audit-check@v1.4.1 with: token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2ae9ff9..3d073aa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -48,6 +48,23 @@ jobs: - name: Run cargo test run: cargo test -- --test-threads 1 + msrv: + name: Minimum Supported Rust Version + runs-on: ubuntu-latest + env: + RUSTFLAGS: "-D warnings" + steps: + - name: Checkout sources + uses: actions/checkout@v3 + + - name: Install MSRV toolchain + uses: dtolnay/rust-toolchain@1.63.0 + + #- uses: Swatinem/rust-cache@v1 + + - name: Run cargo build + run: cargo build + lints: name: Lints runs-on: ubuntu-latest @@ -115,7 +132,7 @@ jobs: runs-on: ubuntu-latest environment: production if: github.event_name == 'release' - needs: [build, test, lints, docs, leaks] + needs: [build, test, msrv, lints, docs, leaks] steps: - name: Checkout sources uses: actions/checkout@v3 diff --git a/Cargo.toml b/Cargo.toml index f9c70b6..07e8143 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,8 +1,9 @@ [package] name = "tokio-graceful-shutdown" authors = ["Finomnis "] -version = "0.13.0" -edition = "2018" +version = "0.14.0" +edition = "2021" +rust-version = "1.63" license = "MIT OR Apache-2.0" readme = "README.md" repository = "https://github.com/Finomnis/tokio-graceful-shutdown" @@ -20,39 +21,35 @@ exclude = [ ] [dependencies] -# Error definitions -thiserror = "1.0.32" -miette = "5.3.0" +tracing = { version = "0.1.37", default-features = false } -# For async utilities -tokio = { version = "1.20.1", default-features = false, features = [ +tokio = { version = "1.32.0", default-features = false, features = [ "signal", "rt", "macros", "time", ] } -tokio-util = { version = "0.7.2", default-features = false } -futures = "0.3.23" -async-recursion = "1.0.0" -pin-project-lite = "0.2.9" +tokio-util = { version = "0.7.8", default-features = false } -# For 'IntoSubsystem' trait -async-trait = "0.1.57" - -# For logging -log = "0.4.17" +pin-project-lite = "0.2.13" +thiserror = "1.0.49" +miette = "5.10.0" +async-trait = "0.1.73" +atomic = "0.6.0" +bytemuck = { version = "1.14.0", features = ["derive"] } [dev-dependencies] # Error propagation -anyhow = "1.0.61" +anyhow = "1.0.75" eyre = "0.6.8" -miette = { version = "5.3.0", features = ["fancy"] } +miette = { version = "5.10.0", features = ["fancy"] } # Logging -env_logger = "0.10.0" +tracing-subscriber = "0.3.17" +tracing-test = "0.2.4" # Tokio -tokio = { version = "1.20.1", features = ["full"] } +tokio = { version = "1.32.0", features = ["full"] } # Hyper example hyper = { version = "0.14.20", features = ["full"] } diff --git a/README.md b/README.md index b5445da..c3ff462 100644 --- a/README.md +++ b/README.md @@ -40,17 +40,19 @@ This subsystem can now be executed like this: ```rust #[tokio::main] async fn main() -> Result<()> { - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)) + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } ``` The `Toplevel` object is the root object of the subsystem tree. -Subsystems can then be started using the `start()` functionality of the toplevel object. +Subsystems can then be started in it using the `start()` method +of its `SubsystemHandle` object. The `catch_signals()` method signals the `Toplevel` object to listen for SIGINT/SIGTERM/Ctrl+C and initiate a shutdown thereafter. diff --git a/examples/01_normal_shutdown.rs b/examples/01_normal_shutdown.rs index 7f5dacc..7d717df 100644 --- a/examples/01_normal_shutdown.rs +++ b/examples/01_normal_shutdown.rs @@ -7,40 +7,42 @@ //! If custom arguments for the subsystem coroutines are required, //! a struct has to be used instead, as seen in other examples. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(400)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } async fn subsys2(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem2 started."); + tracing::info!("Subsystem2 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem2 ..."); + tracing::info!("Shutting down Subsystem2 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem2 stopped."); + tracing::info!("Subsystem2 stopped."); Ok(()) } #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .start("Subsys2", subsys2) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + s.start(SubsystemBuilder::new("Subsys2", subsys2)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/02_structs.rs b/examples/02_structs.rs index e675b1a..089da5d 100644 --- a/examples/02_structs.rs +++ b/examples/02_structs.rs @@ -2,15 +2,14 @@ //! custom parameters to be passed to the subsystem. //! //! There are two ways of using structs as subsystems, by either -//! wrapping them in an async closure, or by implementing the +//! wrapping them in a closure, or by implementing the //! IntoSubsystem trait. Note, though, that the IntoSubsystem //! trait requires an additional dependency, `async-trait`. use async_trait::async_trait; -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{IntoSubsystem, SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{IntoSubsystem, SubsystemBuilder, SubsystemHandle, Toplevel}; struct Subsystem1 { arg: u32, @@ -18,11 +17,11 @@ struct Subsystem1 { impl Subsystem1 { async fn run(self, subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started. Extra argument: {}", self.arg); + tracing::info!("Subsystem1 started. Extra argument: {}", self.arg); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } } @@ -34,11 +33,11 @@ struct Subsystem2 { #[async_trait] impl IntoSubsystem for Subsystem2 { async fn run(self, subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem2 started. Extra argument: {}", self.arg); + tracing::info!("Subsystem2 started. Extra argument: {}", self.arg); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem2 ..."); + tracing::info!("Shutting down Subsystem2 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem2 stopped."); + tracing::info!("Subsystem2 stopped."); Ok(()) } } @@ -46,17 +45,20 @@ impl IntoSubsystem for Subsystem2 { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); let subsys1 = Subsystem1 { arg: 42 }; let subsys2 = Subsystem2 { arg: 69 }; - // Create toplevel - Toplevel::new() - .start("Subsys1", |a| subsys1.run(a)) - .start("Subsys2", subsys2.into_subsystem()) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", |a| subsys1.run(a))); + s.start(SubsystemBuilder::new("Subsys2", subsys2.into_subsystem())); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/03_shutdown_timeout.rs b/examples/03_shutdown_timeout.rs index 3ce386f..5063d28 100644 --- a/examples/03_shutdown_timeout.rs +++ b/examples/03_shutdown_timeout.rs @@ -4,30 +4,32 @@ //! so the subsystem gets cancelled and the program returns an appropriate //! error code. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(2000)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(500)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(500)) + .await + .map_err(Into::into) } diff --git a/examples/04_subsystem_finished.rs b/examples/04_subsystem_finished.rs index cb2ddef..7c09c96 100644 --- a/examples/04_subsystem_finished.rs +++ b/examples/04_subsystem_finished.rs @@ -1,19 +1,18 @@ -//! This subsystem demonstrates that subsystems can also stop +//! This example demonstrates that subsystems can also stop //! prematurely. //! //! Returning Ok(()) from a subsystem indicates that the subsystem //! stopped intentionally, and no further measures by the runtime are performed. //! (unless there are no more subsystems left, in that case TopLevel would shut down anyway) -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); // Task ends without an error. This should not cause the main program to shutdown, // because Subsys2 is still running. @@ -28,14 +27,17 @@ async fn subsys2(subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .start("Subsys2", subsys2) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + s.start(SubsystemBuilder::new("Subsys2", subsys2)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/05_subsystem_finished_with_error.rs b/examples/05_subsystem_finished_with_error.rs index 5b72f32..c6f3009 100644 --- a/examples/05_subsystem_finished_with_error.rs +++ b/examples/05_subsystem_finished_with_error.rs @@ -6,18 +6,17 @@ //! As expected, this is a graceful shutdown, giving other subsystems //! the chance to also shut down gracefully. -use env_logger::{Builder, Env}; use miette::{miette, Result}; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); // Task ends with an error. This should cause the main program to shutdown. - Err(miette!("Subsystem1 threw an error.")) + Err(miette!("Subsystem1 failed intentionally.")) } async fn subsys2(subsys: SubsystemHandle) -> Result<()> { @@ -28,14 +27,17 @@ async fn subsys2(subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .start("Subsys2", subsys2) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + s.start(SubsystemBuilder::new("Subsys2", subsys2)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/06_nested_subsystems.rs b/examples/06_nested_subsystems.rs index e32a1b8..795e0e8 100644 --- a/examples/06_nested_subsystems.rs +++ b/examples/06_nested_subsystems.rs @@ -1,40 +1,42 @@ //! This example demonstrates how one subsystem can launch another //! nested subsystem. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { - subsys.start("Subsys2", subsys2); - log::info!("Subsystem1 started."); + subsys.start(SubsystemBuilder::new("Subsys2", subsys2)); + tracing::info!("Subsystem1 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } async fn subsys2(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem2 started."); + tracing::info!("Subsystem2 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem2 ..."); + tracing::info!("Shutting down Subsystem2 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem2 stopped."); + tracing::info!("Subsystem2 stopped."); Ok(()) } #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/07_nested_error.rs b/examples/07_nested_error.rs index 6b52fe3..ff0c711 100644 --- a/examples/07_nested_error.rs +++ b/examples/07_nested_error.rs @@ -2,38 +2,40 @@ //! a graceful shutdown is performed and other subsystems get the chance //! to clean up. -use env_logger::{Builder, Env}; use miette::{miette, Result}; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { - subsys.start("Subsys2", subsys2); - log::info!("Subsystem1 started."); + subsys.start(SubsystemBuilder::new("Subsys2", subsys2)); + tracing::info!("Subsystem1 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } async fn subsys2(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem2 started."); + tracing::info!("Subsystem2 started."); sleep(Duration::from_millis(500)).await; - Err(miette!("Subsystem2 threw an error.")) + Err(miette!("Subsystem2 failed intentionally.")) } #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/08_panic_handling.rs b/examples/08_panic_handling.rs index 6295997..d690ac0 100644 --- a/examples/08_panic_handling.rs +++ b/examples/08_panic_handling.rs @@ -4,23 +4,22 @@ //! A normal program shutdown is performed, and other subsystems get the //! chance to clean up their work. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { - subsys.start("Subsys2", subsys2); - log::info!("Subsystem1 started."); + subsys.start(SubsystemBuilder::new("Subsys2", subsys2)); + tracing::info!("Subsystem1 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } async fn subsys2(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem2 started."); + tracing::info!("Subsystem2 started."); sleep(Duration::from_millis(500)).await; panic!("Subsystem2 panicked!") @@ -29,13 +28,16 @@ async fn subsys2(_subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/09_task_cancellation.rs b/examples/09_task_cancellation.rs index 9ff3706..b7ac464 100644 --- a/examples/09_task_cancellation.rs +++ b/examples/09_task_cancellation.rs @@ -1,5 +1,6 @@ //! This example demonstrates how to implement a clean shutdown -//! of a subsystem. +//! of a subsystem, through the example of a countdown that +//! gets cancelled on shutdown. //! //! There are two options to cancel tasks on shutdown: //! - with [tokio::select] @@ -7,10 +8,11 @@ //! //! In this case we go with `cancel_on_shutdown()`, but `tokio::select` would be equally viable. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{errors::CancelledByShutdown, FutureExt, SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{ + errors::CancelledByShutdown, FutureExt, SubsystemBuilder, SubsystemHandle, Toplevel, +}; struct CountdownSubsystem {} impl CountdownSubsystem { @@ -20,20 +22,20 @@ impl CountdownSubsystem { async fn countdown(&self) { for i in (1..10).rev() { - log::info!("Countdown: {}", i); + tracing::info!("Countdown: {}", i); sleep(Duration::from_millis(1000)).await; } } async fn run(self, subsys: SubsystemHandle) -> Result<()> { - log::info!("Starting countdown ..."); + tracing::info!("Starting countdown ..."); match self.countdown().cancel_on_shutdown(&subsys).await { Ok(()) => { - log::info!("Countdown finished."); + tracing::info!("Countdown finished."); } Err(CancelledByShutdown) => { - log::info!("Countdown cancelled."); + tracing::info!("Countdown cancelled."); } } @@ -44,13 +46,18 @@ impl CountdownSubsystem { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Countdown", |h| CountdownSubsystem::new().run(h)) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Countdown", |h| { + CountdownSubsystem::new().run(h) + })); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/10_request_shutdown.rs b/examples/10_request_shutdown.rs index acbc880..6aceaa5 100644 --- a/examples/10_request_shutdown.rs +++ b/examples/10_request_shutdown.rs @@ -1,10 +1,11 @@ //! This example demonstrates how a subsystem can initiate //! a shutdown. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{errors::CancelledByShutdown, FutureExt, SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{ + errors::CancelledByShutdown, FutureExt, SubsystemBuilder, SubsystemHandle, Toplevel, +}; struct CountdownSubsystem {} impl CountdownSubsystem { @@ -14,7 +15,7 @@ impl CountdownSubsystem { async fn countdown(&self) { for i in (1..10).rev() { - log::info!("Shutting down in: {}", i); + tracing::info!("Shutting down in: {}", i); sleep(Duration::from_millis(1000)).await; } } @@ -22,7 +23,7 @@ impl CountdownSubsystem { async fn run(self, subsys: SubsystemHandle) -> Result<()> { match self.countdown().cancel_on_shutdown(&subsys).await { Ok(()) => subsys.request_shutdown(), - Err(CancelledByShutdown) => log::info!("Countdown cancelled."), + Err(CancelledByShutdown) => tracing::info!("Countdown cancelled."), } Ok(()) @@ -32,13 +33,18 @@ impl CountdownSubsystem { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Countdown", |h| CountdownSubsystem::new().run(h)) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Countdown", |h| { + CountdownSubsystem::new().run(h) + })); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/11_double_panic.rs b/examples/11_double_panic.rs index 5079cda..a38bc4e 100644 --- a/examples/11_double_panic.rs +++ b/examples/11_double_panic.rs @@ -7,47 +7,49 @@ //! There is no real programming knowledge to be gained here, this example is just //! to demonstrate the robustness of the system. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { - subsys.start("Subsys2", subsys2); - subsys.start("Subsys3", subsys3); - log::info!("Subsystem1 started."); + subsys.start(SubsystemBuilder::new("Subsys2", subsys2)); + subsys.start(SubsystemBuilder::new("Subsys3", subsys3)); + tracing::info!("Subsystem1 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem1 ..."); + tracing::info!("Shutting down Subsystem1 ..."); sleep(Duration::from_millis(200)).await; panic!("Subsystem1 panicked!"); } async fn subsys2(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem2 started."); + tracing::info!("Subsystem2 started."); sleep(Duration::from_millis(500)).await; panic!("Subsystem2 panicked!") } async fn subsys3(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem3 started."); + tracing::info!("Subsystem3 started."); subsys.on_shutdown_requested().await; - log::info!("Shutting down Subsystem3 ..."); + tracing::info!("Shutting down Subsystem3 ..."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem3 shut down successfully."); + tracing::info!("Subsystem3 shut down successfully."); Ok(()) } #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/12_subsystem_auto_restart.rs b/examples/12_subsystem_auto_restart.rs index 3c98341..f4369cf 100644 --- a/examples/12_subsystem_auto_restart.rs +++ b/examples/12_subsystem_auto_restart.rs @@ -4,43 +4,41 @@ //! This isn't really a usecase related to this library, but seems to be used regularly, //! so I included it anyway. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{ErrorAction, SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(subsys: SubsystemHandle) -> Result<()> { // This subsystem panics every two seconds. // It should get restarted constantly. - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); tokio::select! { _ = subsys.on_shutdown_requested() => (), _ = sleep(Duration::from_secs(2)) => { panic!("Subsystem1 panicked!"); } }; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Ok(()) } async fn subsys1_keepalive(subsys: SubsystemHandle) -> Result<()> { loop { - let subsys_result = Toplevel::nested(&subsys, "") - .start("Subsys1", subsys1) - .handle_shutdown_requests(Duration::from_millis(50)) - .await; - - if let Err(err) = &subsys_result { - log::error!("Subsystem1 failed: {}", err); - } - - if subsys.is_shutdown_requested() { + let nested_subsys = subsys.start( + SubsystemBuilder::new("Subsys1", subsys1) + .on_failure(ErrorAction::CatchAndLocalShutdown) + .on_panic(ErrorAction::CatchAndLocalShutdown), + ); + + if let Err(err) = nested_subsys.join().await { + tracing::error!("Subsystem1 failed: {:?}", miette::Report::from(err)); + } else { break; } - log::info!("Restarting subsystem1 ..."); + tracing::info!("Restarting subsystem1 ..."); } Ok(()) @@ -49,13 +47,16 @@ async fn subsys1_keepalive(subsys: SubsystemHandle) -> Result<()> { #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Subsys1Keepalive", subsys1_keepalive) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1Keepalive", subsys1_keepalive)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/13_partial_shutdown.rs b/examples/13_partial_shutdown.rs index b6d4fad..8f92b0d 100644 --- a/examples/13_partial_shutdown.rs +++ b/examples/13_partial_shutdown.rs @@ -3,45 +3,49 @@ //! Subsys1 will perform a partial shutdown after 5 seconds, which will in turn //! shut down Subsys2 and Subsys3, leaving Subsys1 running. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{ErrorAction, SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys3(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsys3 started."); + tracing::info!("Subsys3 started."); subsys.on_shutdown_requested().await; - log::info!("Subsys3 stopped."); + tracing::info!("Subsys3 stopped."); Ok(()) } async fn subsys2(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsys2 started."); - subsys.start("Subsys3", subsys3); + tracing::info!("Subsys2 started."); + subsys.start(SubsystemBuilder::new("Subsys3", subsys3)); subsys.on_shutdown_requested().await; - log::info!("Subsys2 stopped."); + tracing::info!("Subsys2 stopped."); Ok(()) } async fn subsys1(subsys: SubsystemHandle) -> Result<()> { // This subsystem shuts down the nested subsystem after 5 seconds. - log::info!("Subsys1 started."); + tracing::info!("Subsys1 started."); - log::info!("Starting nested subsystem ..."); - let nested_subsys = subsys.start("Subsys2", subsys2); - log::info!("Nested subsystem started."); + tracing::info!("Starting nested subsystem ..."); + let nested_subsys = subsys.start(SubsystemBuilder::new("Subsys2", subsys2)); + tracing::info!("Nested subsystem started."); tokio::select! { _ = subsys.on_shutdown_requested() => (), - _ = sleep(Duration::from_secs(5)) => { - log::info!("Shutting down nested subsystem ..."); - subsys.perform_partial_shutdown(nested_subsys).await?; - log::info!("Nested subsystem shut down."); + _ = sleep(Duration::from_secs(1)) => { + tracing::info!("Shutting down nested subsystem ..."); + // Redirect errors during shutdown to the local `.join()` call + nested_subsys.change_failure_action(ErrorAction::CatchAndLocalShutdown); + nested_subsys.change_panic_action(ErrorAction::CatchAndLocalShutdown); + // Perform shutdown + nested_subsys.initiate_shutdown(); + nested_subsys.join().await?; + tracing::info!("Nested subsystem shut down."); subsys.on_shutdown_requested().await; } }; - log::info!("Subsys1 stopped."); + tracing::info!("Subsys1 stopped."); Ok(()) } @@ -49,13 +53,16 @@ async fn subsys1(subsys: SubsystemHandle) -> Result<()> { #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/14_partial_shutdown_error.rs b/examples/14_partial_shutdown_error.rs index d2be52e..93a36d1 100644 --- a/examples/14_partial_shutdown_error.rs +++ b/examples/14_partial_shutdown_error.rs @@ -4,46 +4,50 @@ //! shutdown, but instead it will be delivered to the task that initiated //! the partial shutdown. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{ErrorAction, SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys3(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsys3 started."); + tracing::info!("Subsys3 started."); subsys.on_shutdown_requested().await; panic!("Subsystem3 threw an error!") } async fn subsys2(subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsys2 started."); - subsys.start("Subsys3", subsys3); + tracing::info!("Subsys2 started."); + subsys.start(SubsystemBuilder::new("Subsys3", subsys3)); subsys.on_shutdown_requested().await; - log::info!("Subsys2 stopped."); + tracing::info!("Subsys2 stopped."); Ok(()) } async fn subsys1(subsys: SubsystemHandle) -> Result<()> { // This subsystem shuts down the nested subsystem after 5 seconds. - log::info!("Subsys1 started."); + tracing::info!("Subsys1 started."); - log::info!("Starting nested subsystem ..."); - let nested_subsys = subsys.start("Subsys2", subsys2); - log::info!("Nested subsystem started."); + tracing::info!("Starting nested subsystem ..."); + let nested_subsys = subsys.start(SubsystemBuilder::new("Subsys2", subsys2)); + tracing::info!("Nested subsystem started."); tokio::select! { _ = subsys.on_shutdown_requested() => (), _ = sleep(Duration::from_secs(1)) => { - log::info!("Shutting down nested subsystem ..."); - if let Err(err) = subsys.perform_partial_shutdown(nested_subsys).await{ - log::warn!("Partial shutdown failed: {}", err); + tracing::info!("Shutting down nested subsystem ..."); + // Redirect errors during shutdown to the local `.join()` call + nested_subsys.change_failure_action(ErrorAction::CatchAndLocalShutdown); + nested_subsys.change_panic_action(ErrorAction::CatchAndLocalShutdown); + // Perform shutdown + nested_subsys.initiate_shutdown(); + if let Err(err) = nested_subsys.join().await { + tracing::warn!("Error during nested subsystem shutdown: {:?}", miette::Report::from(err)); }; - log::info!("Nested subsystem shut down."); + tracing::info!("Nested subsystem shut down."); subsys.on_shutdown_requested().await; } }; - log::info!("Subsys1 stopped."); + tracing::info!("Subsys1 stopped."); Ok(()) } @@ -51,13 +55,16 @@ async fn subsys1(subsys: SubsystemHandle) -> Result<()> { #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/15_without_miette.rs b/examples/15_without_miette.rs index 7ff2a69..330ff9c 100644 --- a/examples/15_without_miette.rs +++ b/examples/15_without_miette.rs @@ -1,10 +1,9 @@ //! This example shows how to use this library with std::error::Error instead of miette::Error -use env_logger::{Builder, Env}; use std::error::Error; use std::fmt; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; #[derive(Debug, Clone)] struct MyError; @@ -18,9 +17,9 @@ impl fmt::Display for MyError { impl Error for MyError {} async fn subsys1(_subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); // Task ends with an error. This should cause the main program to shutdown. Err(MyError {}) @@ -29,13 +28,16 @@ async fn subsys1(_subsys: SubsystemHandle) -> Result<(), MyError> { #[tokio::main] async fn main() -> Result<(), Box> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/16_with_anyhow.rs b/examples/16_with_anyhow.rs index 1463a29..74e9370 100644 --- a/examples/16_with_anyhow.rs +++ b/examples/16_with_anyhow.rs @@ -1,14 +1,13 @@ //! This example shows how to use this library with anyhow instead of miette use anyhow::{anyhow, Result}; -use env_logger::{Builder, Env}; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); // Task ends with an error. This should cause the main program to shutdown. Err(anyhow!("Subsystem1 threw an error.")) @@ -17,13 +16,16 @@ async fn subsys1(_subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/17_with_eyre.rs b/examples/17_with_eyre.rs index dd54f05..148a22b 100644 --- a/examples/17_with_eyre.rs +++ b/examples/17_with_eyre.rs @@ -1,14 +1,13 @@ //! This example shows how to use this library with eyre instead of miette -use env_logger::{Builder, Env}; use eyre::{eyre, Result}; use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; async fn subsys1(_subsys: SubsystemHandle) -> Result<()> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); sleep(Duration::from_millis(500)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); // Task ends with an error. This should cause the main program to shutdown. Err(eyre!("Subsystem1 threw an error.")) @@ -17,13 +16,16 @@ async fn subsys1(_subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); - // Create toplevel - Toplevel::new() - .start("Subsys1", subsys1) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(1000)) - .await - .map_err(Into::into) + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(1000)) + .await + .map_err(Into::into) } diff --git a/examples/18_error_type_passthrough.rs b/examples/18_error_type_passthrough.rs index e8fd494..4b18c4e 100644 --- a/examples/18_error_type_passthrough.rs +++ b/examples/18_error_type_passthrough.rs @@ -1,11 +1,10 @@ //! This example shows to pass custom error types all the way through to the top, //! to recover them from the return value of `handle_shutdown_requests`. -use env_logger::{Builder, Env}; use tokio::time::{sleep, Duration}; use tokio_graceful_shutdown::{ errors::{GracefulShutdownError, SubsystemError}, - IntoSubsystem, SubsystemHandle, Toplevel, + IntoSubsystem, SubsystemBuilder, SubsystemHandle, Toplevel, }; #[derive(Debug, thiserror::Error)] @@ -17,33 +16,33 @@ enum MyError { } async fn subsys1(_subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem1 started."); + tracing::info!("Subsystem1 started."); sleep(Duration::from_millis(200)).await; - log::info!("Subsystem1 stopped."); + tracing::info!("Subsystem1 stopped."); Err(MyError::WithData(42)) } async fn subsys2(_subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem2 started."); + tracing::info!("Subsystem2 started."); sleep(Duration::from_millis(200)).await; - log::info!("Subsystem2 stopped."); + tracing::info!("Subsystem2 stopped."); Err(MyError::WithoutData) } async fn subsys3(_subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem3 started."); + tracing::info!("Subsystem3 started."); sleep(Duration::from_millis(200)).await; - log::info!("Subsystem3 stopped."); + tracing::info!("Subsystem3 stopped."); panic!("This subsystem panicked."); } async fn subsys4(_subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem4 started."); + tracing::info!("Subsystem4 started."); sleep(Duration::from_millis(1000)).await; - log::info!("Subsystem4 stopped."); + tracing::info!("Subsystem4 stopped."); // This subsystem would end normally but takes too long and therefore // will time out. @@ -51,9 +50,9 @@ async fn subsys4(_subsys: SubsystemHandle) -> Result<(), MyError> { } async fn subsys5(_subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem5 started."); + tracing::info!("Subsystem5 started."); sleep(Duration::from_millis(200)).await; - log::info!("Subsystem5 stopped."); + tracing::info!("Subsystem5 stopped."); // This subsystem ended normally and should not show up in the list of // subsystem errors. @@ -69,9 +68,9 @@ struct Subsys6; #[async_trait::async_trait] impl IntoSubsystem for Subsys6 { async fn run(self, _subsys: SubsystemHandle) -> Result<(), MyError> { - log::info!("Subsystem6 started."); + tracing::info!("Subsystem6 started."); sleep(Duration::from_millis(200)).await; - log::info!("Subsystem6 stopped."); + tracing::info!("Subsystem6 stopped."); Err(MyError::WithData(69)) } @@ -80,48 +79,48 @@ impl IntoSubsystem for Subsys6 { #[tokio::main] async fn main() -> Result<(), miette::Report> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - let errors = Toplevel::::new() - .start("Subsys1", subsys1) - .start("Subsys2", subsys2) - .start("Subsys3", subsys3) - .start("Subsys4", subsys4) - .start("Subsys5", subsys5) - .start("Subsys6", Subsys6.into_subsystem()) - .catch_signals() - .handle_shutdown_requests(Duration::from_millis(500)) - .await; + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + let errors = Toplevel::::new(|s| async move { + s.start(SubsystemBuilder::new("Subsys1", subsys1)); + s.start(SubsystemBuilder::new("Subsys2", subsys2)); + s.start(SubsystemBuilder::new("Subsys3", subsys3)); + s.start(SubsystemBuilder::new("Subsys4", subsys4)); + s.start(SubsystemBuilder::new("Subsys5", subsys5)); + s.start(SubsystemBuilder::new("Subsys6", Subsys6.into_subsystem())); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(500)) + .await; if let Err(e) = &errors { match e { GracefulShutdownError::SubsystemsFailed(_) => { - log::warn!("Subsystems failed.") + tracing::warn!("Subsystems failed.") } GracefulShutdownError::ShutdownTimeout(_) => { - log::warn!("Shutdown timed out.") + tracing::warn!("Shutdown timed out.") } }; for subsystem_error in e.get_subsystem_errors() { match subsystem_error { SubsystemError::Failed(name, e) => { - log::warn!(" Subsystem '{}' failed.", name); + tracing::warn!(" Subsystem '{}' failed.", name); match e.get_error() { MyError::WithData(data) => { - log::warn!(" It failed with MyError::WithData({})", data) + tracing::warn!(" It failed with MyError::WithData({})", data) } MyError::WithoutData => { - log::warn!(" It failed with MyError::WithoutData") + tracing::warn!(" It failed with MyError::WithoutData") } } } - SubsystemError::Cancelled(name) => { - log::warn!(" Subsystem '{}' was cancelled.", name) - } SubsystemError::Panicked(name) => { - log::warn!(" Subsystem '{}' panicked.", name) + tracing::warn!(" Subsystem '{}' panicked.", name) } } } diff --git a/examples/hyper.rs b/examples/hyper.rs index 5e655ad..72348d4 100644 --- a/examples/hyper.rs +++ b/examples/hyper.rs @@ -7,10 +7,9 @@ //! hyper's graceful shutdown waits for all connections to be closed naturally //! instead of terminating them. -use env_logger::{Builder, Env}; use miette::{miette, Result}; use tokio::time::Duration; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; use std::convert::Infallible; @@ -34,7 +33,7 @@ async fn hyper_subsystem(subsys: SubsystemHandle) -> Result<()> { let addr = ([127, 0, 0, 1], 12345).into(); let server = Server::bind(&addr).serve(make_svc); - log::info!("Listening on http://{}", addr); + tracing::info!("Listening on http://{}", addr); // This is the connection between our crate and hyper. // Hyper already anticipated our use case and provides a very @@ -48,13 +47,16 @@ async fn hyper_subsystem(subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Hyper", hyper_subsystem) - .catch_signals() - .handle_shutdown_requests(Duration::from_secs(60)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Hyper", hyper_subsystem)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_secs(60)) + .await + .map_err(Into::into) } diff --git a/examples/warp.rs b/examples/warp.rs index 7c5aa76..96f36df 100644 --- a/examples/warp.rs +++ b/examples/warp.rs @@ -5,10 +5,9 @@ //! warp's graceful shutdown waits for all connections to be closed naturally //! instead of terminating them. -use env_logger::{Builder, Env}; use miette::Result; use tokio::time::Duration; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; use warp::Filter; @@ -19,10 +18,10 @@ async fn warp_subsystem(subsys: SubsystemHandle) -> Result<()> { let (addr, server) = warp::serve(routes).bind_with_graceful_shutdown(([127, 0, 0, 1], 12345), async move { subsys.on_shutdown_requested().await; - log::info!("Starting server shutdown ..."); + tracing::info!("Starting server shutdown ..."); }); - log::info!("Listening on http://{}", addr); + tracing::info!("Listening on http://{}", addr); server.await; @@ -32,13 +31,16 @@ async fn warp_subsystem(subsys: SubsystemHandle) -> Result<()> { #[tokio::main] async fn main() -> Result<()> { // Init logging - Builder::from_env(Env::default().default_filter_or("debug")).init(); - - // Create toplevel - Toplevel::new() - .start("Warp", warp_subsystem) - .catch_signals() - .handle_shutdown_requests(Duration::from_secs(60)) - .await - .map_err(Into::into) + tracing_subscriber::fmt() + .with_max_level(tracing::Level::TRACE) + .init(); + + // Setup and execute subsystem tree + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("Warp", warp_subsystem)); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_secs(60)) + .await + .map_err(Into::into) } diff --git a/src/error_action.rs b/src/error_action.rs new file mode 100644 index 0000000..63a37ba --- /dev/null +++ b/src/error_action.rs @@ -0,0 +1,38 @@ +use bytemuck::NoUninit; + +/// Possible ways a subsystem can react to errors. +/// +/// An error will propagate upwards in the subsystem tree until +/// it reaches a subsystem that won't forward it to its parent. +/// +/// If an error reaches the [`Toplevel`](crate::Toplevel), a global shutdown will be initiated. +/// +/// Also see: +/// - [`SubsystemBuilder::on_failure`](crate::SubsystemBuilder::on_failure) +/// - [`SubsystemBuilder::on_panic`](crate::SubsystemBuilder::on_panic) +/// - [`NestedSubsystem::change_failure_action`](crate::NestedSubsystem::change_failure_action) +/// - [`NestedSubsystem::change_panic_action`](crate::NestedSubsystem::change_panic_action) +/// +#[derive(Clone, Copy, Debug, Eq, PartialEq, NoUninit)] +#[repr(u8)] +pub enum ErrorAction { + /// Pass the error on to the parent subsystem, but don't react to it. + Forward, + /// Store the error so it can be retrieved through + /// [`NestedSubsystem::join`](crate::NestedSubsystem::join), + /// then initiate a shutdown of the subsystem and its children. + /// Do not forward the error to the parent subsystem. + CatchAndLocalShutdown, +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn derive_traits() { + let x = ErrorAction::CatchAndLocalShutdown; + #[allow(clippy::clone_on_copy)] + let y = x.clone(); + assert!(y == x); + } +} diff --git a/src/errors.rs b/src/errors.rs index 686606b..65871d7 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -1,5 +1,7 @@ //! All the errors that can be caused by this crate. +use std::sync::Arc; + use miette::Diagnostic; use thiserror::Error; @@ -11,22 +13,24 @@ use crate::ErrTypeTraits; pub enum GracefulShutdownError { /// At least one subsystem caused an error. #[error("at least one subsystem returned an error")] - SubsystemsFailed(#[related] Vec>), + #[diagnostic(code(graceful_shutdown::failed))] + SubsystemsFailed(#[related] Box<[SubsystemError]>), /// The shutdown did not finish within the given timeout. #[error("shutdown timed out")] - ShutdownTimeout(#[related] Vec>), + #[diagnostic(code(graceful_shutdown::timeout))] + ShutdownTimeout(#[related] Box<[SubsystemError]>), } impl GracefulShutdownError { /// Converts the error into a list of subsystem errors that occurred. - pub fn into_subsystem_errors(self) -> Vec> { + pub fn into_subsystem_errors(self) -> Box<[SubsystemError]> { match self { GracefulShutdownError::SubsystemsFailed(rel) => rel, GracefulShutdownError::ShutdownTimeout(rel) => rel, } } /// Queries the list of subsystem errors that occurred. - pub fn get_subsystem_errors(&self) -> &Vec> { + pub fn get_subsystem_errors(&self) -> &[SubsystemError] { match self { GracefulShutdownError::SubsystemsFailed(rel) => rel, GracefulShutdownError::ShutdownTimeout(rel) => rel, @@ -34,21 +38,14 @@ impl GracefulShutdownError { } } -/// This enum contains all the possible errors that a partial shutdown +/// This enum contains all the possible errors that joining a subsystem /// could cause. #[derive(Debug, Error, Diagnostic)] -pub enum PartialShutdownError { +pub enum SubsystemJoinError { /// At least one subsystem caused an error. + #[diagnostic(code(graceful_shutdown::subsystem_join::failed))] #[error("at least one subsystem returned an error")] - SubsystemsFailed(#[related] Vec>), - /// The given nested subsystem does not seem to be a child of - /// the parent subsystem. - #[error("unable to find nested subsystem in given subsystem")] - SubsystemNotFound, - /// A partial shutdown can not be performed because the entire program - /// is already shutting down. - #[error("unable to perform partial shutdown, the program is already shutting down")] - AlreadyShuttingDown, + SubsystemsFailed(#[related] Arc<[SubsystemError]>), } /// A wrapper type that carries the errors returned by subsystems. @@ -100,14 +97,13 @@ impl std::error::Error for SubsystemFailure where #[derive(Debug, Error, Diagnostic)] pub enum SubsystemError { /// The subsystem returned an error value. Carries the actual error as the second argument. + #[diagnostic(code(graceful_shutdown::subsystem::failed))] #[error("Error in subsystem '{0}'")] - Failed(String, #[source] SubsystemFailure), - /// The subsystem was cancelled. Should only happen if the shutdown timeout is exceeded. - #[error("Subsystem '{0}' was aborted")] - Cancelled(String), + Failed(Arc, #[source] SubsystemFailure), /// The subsystem panicked. + #[diagnostic(code(graceful_shutdown::subsystem::panicked))] #[error("Subsystem '{0}' panicked")] - Panicked(String), + Panicked(Arc), } impl SubsystemError { @@ -119,7 +115,6 @@ impl SubsystemError { pub fn name(&self) -> &str { match self { SubsystemError::Failed(name, _) => name, - SubsystemError::Cancelled(name) => name, SubsystemError::Panicked(name) => name, } } @@ -148,12 +143,9 @@ mod tests { #[test] fn errors_can_be_converted_to_diagnostic() { - examine_report(GracefulShutdownError::ShutdownTimeout::(vec![]).into()); - examine_report(GracefulShutdownError::SubsystemsFailed::(vec![]).into()); - examine_report(PartialShutdownError::AlreadyShuttingDown::.into()); - examine_report(PartialShutdownError::SubsystemNotFound::.into()); - examine_report(PartialShutdownError::SubsystemsFailed::(vec![]).into()); - examine_report(SubsystemError::Cancelled::("".into()).into()); + examine_report(GracefulShutdownError::ShutdownTimeout::(Box::new([])).into()); + examine_report(GracefulShutdownError::SubsystemsFailed::(Box::new([])).into()); + examine_report(SubsystemJoinError::SubsystemsFailed::(Arc::new([])).into()); examine_report(SubsystemError::Panicked::("".into()).into()); examine_report( SubsystemError::Failed::("".into(), SubsystemFailure("".into())).into(), @@ -164,18 +156,18 @@ mod tests { #[test] fn extract_related_from_graceful_shutdown_error() { let related = || { - vec![ - SubsystemError::Cancelled("a".into()), + Box::new([ + SubsystemError::Failed("a".into(), SubsystemFailure(String::from("A").into())), SubsystemError::Panicked("b".into()), - ] + ]) }; - let matches_related = |data: &Vec>| { + let matches_related = |data: &[SubsystemError]| { let mut iter = data.iter(); let elem = iter.next().unwrap(); assert_eq!(elem.name(), "a"); - assert!(matches!(elem, SubsystemError::Cancelled(_))); + assert!(matches!(elem, SubsystemError::Failed(_, _))); let elem = iter.next().unwrap(); assert_eq!(elem.name(), "b"); diff --git a/src/exit_state.rs b/src/exit_state.rs deleted file mode 100644 index 1fe6b88..0000000 --- a/src/exit_state.rs +++ /dev/null @@ -1,63 +0,0 @@ -use crate::{errors::SubsystemError, ErrTypeTraits}; - -pub struct SubprocessExitState { - pub name: String, - pub exit_state: String, - pub raw_result: Result<(), SubsystemError>, -} - -impl SubprocessExitState { - pub fn new( - name: &str, - exit_state: &str, - raw_result: Result<(), SubsystemError>, - ) -> Self { - Self { - name: name.to_string(), - exit_state: exit_state.to_string(), - raw_result, - } - } -} - -pub type ShutdownResults = Vec>; - -pub fn join_shutdown_results( - mut left: ShutdownResults, - right: Vec>, -) -> ShutdownResults { - for mut right_element in right { - left.append(&mut right_element); - } - - left -} - -pub fn prettify_exit_states( - exit_states: &[SubprocessExitState], -) -> Vec { - let max_subprocess_name_length = exit_states - .iter() - .map(|code| code.name.len()) - .max() - .unwrap_or(0); - - let mut exit_states = exit_states.iter().collect::>(); - exit_states.sort_by_key(|el| el.name.clone()); - - exit_states - .iter() - .map( - |SubprocessExitState { - name, - exit_state, - raw_result: _, - }| { - let required_padding_length = max_subprocess_name_length - name.len(); - let padding = " ".repeat(required_padding_length); - - name.clone() + &padding + " => " + exit_state - }, - ) - .collect::>() -} diff --git a/src/future_ext.rs b/src/future_ext.rs index 89b457b..f3c2ee7 100644 --- a/src/future_ext.rs +++ b/src/future_ext.rs @@ -89,7 +89,7 @@ impl FutureExt for T { type Future = T; fn cancel_on_shutdown(self, subsys: &SubsystemHandle) -> CancelOnShutdownFuture<'_, T> { - let cancellation = subsys.local_shutdown_token().wait_for_shutdown(); + let cancellation = subsys.get_cancellation_token().cancelled(); CancelOnShutdownFuture { future: self, diff --git a/src/into_subsystem.rs b/src/into_subsystem.rs index 1fd6efe..14b814a 100644 --- a/src/into_subsystem.rs +++ b/src/into_subsystem.rs @@ -26,7 +26,7 @@ type SubsystemFunction = /// ``` /// use miette::Result; /// use tokio::time::Duration; -/// use tokio_graceful_shutdown::{IntoSubsystem, SubsystemHandle, Toplevel}; +/// use tokio_graceful_shutdown::{IntoSubsystem, SubsystemBuilder, SubsystemHandle, Toplevel}; /// /// struct MySubsystem; /// @@ -41,12 +41,15 @@ type SubsystemFunction = /// #[tokio::main] /// async fn main() -> Result<()> { /// // Create toplevel -/// Toplevel::new() -/// .start("Subsys1", MySubsystem{}.into_subsystem()) -/// .catch_signals() -/// .handle_shutdown_requests(Duration::from_millis(500)) -/// .await -/// .map_err(Into::into) +/// Toplevel::new(|s| async move { +/// s.start(SubsystemBuilder::new( +/// "Subsys1", MySubsystem{}.into_subsystem() +/// )); +/// }) +/// .catch_signals() +/// .handle_shutdown_requests(Duration::from_millis(500)) +/// .await +/// .map_err(Into::into) /// } /// ``` /// @@ -63,11 +66,11 @@ where /// Returning an error automatically initiates a shutdown. /// /// For more information about subsystem functions, see - /// [`Toplevel::start()`](crate::Toplevel::start) and [`SubsystemHandle::start()`](crate::SubsystemHandle::start). + /// [`SubsystemHandle::start()`](crate::SubsystemHandle::start). async fn run(self, subsys: SubsystemHandle) -> Result<(), Err>; /// Converts the object into a type that can be passed into - /// [`Toplevel::start()`](crate::Toplevel::start) and [`SubsystemHandle::start()`](crate::SubsystemHandle::start). + /// [`SubsystemHandle::start()`](crate::SubsystemHandle::start). fn into_subsystem(self) -> Box> { Box::new(|handle: SubsystemHandle| { Box::pin(async move { self.run(handle).await }) diff --git a/src/lib.rs b/src/lib.rs index 64bd125..e1d4a89 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,13 +21,12 @@ //! //! ``` //! use miette::Result; -//! use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; -//! use env_logger::{Builder, Env}; +//! use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; //! use tokio::time::{sleep, Duration}; //! //! async fn countdown() { //! for i in (1..=5).rev() { -//! log::info!("Shutting down in: {}", i); +//! tracing::info!("Shutting down in: {}", i); //! sleep(Duration::from_millis(1000)).await; //! } //! } @@ -35,7 +34,7 @@ //! async fn countdown_subsystem(subsys: SubsystemHandle) -> Result<()> { //! tokio::select! { //! _ = subsys.on_shutdown_requested() => { -//! log::info!("Countdown cancelled."); +//! tracing::info!("Countdown cancelled."); //! }, //! _ = countdown() => { //! subsys.request_shutdown(); @@ -48,34 +47,36 @@ //! #[tokio::main] //! async fn main() -> Result<()> { //! // Init logging -//! Builder::from_env(Env::default().default_filter_or("debug")).init(); -//! -//! // Create toplevel -//! Toplevel::new() -//! .start("Countdown", countdown_subsystem) -//! .catch_signals() -//! .handle_shutdown_requests(Duration::from_millis(1000)) -//! .await -//! .map_err(Into::into) +//! tracing_subscriber::fmt() +//! .with_max_level(tracing::Level::TRACE) +//! .init(); +//! +//! // Setup and execute subsystem tree +//! Toplevel::new(|s| async move { +//! s.start(SubsystemBuilder::new("Countdown", countdown_subsystem)); +//! }) +//! .catch_signals() +//! .handle_shutdown_requests(Duration::from_millis(1000)) +//! .await +//! .map_err(Into::into) //! } //! ``` //! -//! There are a couple of things to note here. //! -//! For one, the [`Toplevel`] object represents the root object of the subsystem tree +//! The [`Toplevel`] object represents the root object of the subsystem tree //! and is the main entry point of how to interact with this crate. -//! Subsystems can then be started using the [`start()`](Toplevel::start) functionality of the toplevel object. +//! Creating a [`Toplevel`] object initially spawns a simple subsystem, which can then +//! spawn further subsystems recursively. //! //! The [`catch_signals()`](Toplevel::catch_signals) method signals the `Toplevel` object to listen for SIGINT/SIGTERM/Ctrl+C and initiate a shutdown thereafter. //! //! [`handle_shutdown_requests()`](Toplevel::handle_shutdown_requests) is the final and most important method of `Toplevel`. It idles until the program enters the shutdown mode. Then, it collects all the return values of the subsystems, determines the global error state and makes sure the shutdown happens within the given timeout. //! Lastly, it returns an error value that can be directly used as a return code for `main()`. //! -//! Further, the way to register and start a new submodule ist to provide -//! a submodule function/lambda to [`Toplevel::start`] or -//! [`SubsystemHandle::start`]. +//! Further, the way to register and start a new submodule is to provide +//! a submodule function/lambda to [`SubsystemHandle::start`]. //! If additional arguments shall to be provided to the submodule, it is necessary to create -//! a submodule `struct`. Further details can be seen in the `examples` folder of the repository. +//! a submodule `struct`. Further details can be seen in the `examples` directory of the repository. //! //! Finally, you can see the [`SubsystemHandle`] object that gets provided to the subsystem. //! It is the main way of the subsystem to communicate with this crate. @@ -83,6 +84,7 @@ //! to initiate a shutdown. //! +#![deny(unreachable_pub)] #![deny(missing_docs)] #![doc( issue_tracker_base_url = "https://github.com/Finomnis/tokio-graceful-shutdown/issues", @@ -104,20 +106,20 @@ impl ErrTypeTraits for T where } pub mod errors; -mod exit_state; + +mod error_action; mod future_ext; mod into_subsystem; mod runner; -mod shutdown_token; mod signal_handling; mod subsystem; mod toplevel; mod utils; -use shutdown_token::ShutdownToken; - +pub use error_action::ErrorAction; pub use future_ext::FutureExt; pub use into_subsystem::IntoSubsystem; pub use subsystem::NestedSubsystem; +pub use subsystem::SubsystemBuilder; pub use subsystem::SubsystemHandle; pub use toplevel::Toplevel; diff --git a/src/runner.rs b/src/runner.rs index 6dc01f6..1354057 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -1,122 +1,112 @@ +//! The SubsystemRunner is a little tricky, so here some explanation. +//! +//! A two-layer `tokio::spawn` is required to make this work reliably; the inner `spawn` is the actual subsystem, +//! and the outer `spawn` carries out the duty of propagating the `StopReason` and cleaning up. +//! +//! Further, everything in here reacts properly to being dropped, including +//! the runner itself, who cancels the subsystem on drop. + +use std::{future::Future, sync::Arc}; + use crate::{ errors::{SubsystemError, SubsystemFailure}, - utils::ShutdownGuard, - ErrTypeTraits, ShutdownToken, + ErrTypeTraits, SubsystemHandle, }; -use std::{future::Future, sync::Arc}; -use tokio::task::{JoinError, JoinHandle}; -use tokio_util::sync::CancellationToken; -pub struct SubsystemRunner { - outer_joinhandle: JoinHandle>>, - cancellation_token: CancellationToken, +mod alive_guard; +pub(crate) use self::alive_guard::AliveGuard; + +pub(crate) struct SubsystemRunner { + aborthandle: tokio::task::AbortHandle, } -/// Dropping the SubsystemRunner cancels the task. -/// -/// In consequence, this means that dropping the Toplevel object cancels all tasks. -impl Drop for SubsystemRunner { +impl SubsystemRunner { + pub(crate) fn new( + name: Arc, + subsystem: Subsys, + subsystem_handle: SubsystemHandle, + guard: AliveGuard, + ) -> Self + where + Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, + Fut: 'static + Future> + Send, + Err: Into, + { + let future = async { run_subsystem(name, subsystem, subsystem_handle, guard).await }; + let aborthandle = tokio::spawn(future).abort_handle(); + SubsystemRunner { aborthandle } + } +} + +impl Drop for SubsystemRunner { fn drop(&mut self) { - self.abort(); + self.aborthandle.abort() } } -impl SubsystemRunner { - async fn handle_subsystem( - mut inner_joinhandle: JoinHandle>, - shutdown_token: ShutdownToken, - local_shutdown_token: ShutdownToken, - name: String, - cancellation_token: CancellationToken, - shutdown_guard: Arc, - ) -> Result<(), SubsystemError> { - /// Maps the complicated return value of the subsystem joinhandle to an appropriate error - fn map_subsystem_result( - name: &str, - result: Result, JoinError>, - ) -> Result<(), SubsystemError> { - match result { - Ok(Ok(())) => Ok(()), - Ok(Err(e)) => Err(SubsystemError::Failed( - name.to_string(), - SubsystemFailure(e), - )), - Err(e) => Err(if e.is_cancelled() { - SubsystemError::Cancelled(name.to_string()) - } else { - SubsystemError::Panicked(name.to_string()) - }), - } - } +async fn run_subsystem( + name: Arc, + subsystem: Subsys, + mut subsystem_handle: SubsystemHandle, + guard: AliveGuard, +) where + Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, + Fut: 'static + Future> + Send, + Err: Into, +{ + let mut redirected_subsystem_handle = subsystem_handle.delayed_clone(); - let joinhandle_ref = &mut inner_joinhandle; - let result = tokio::select! { - result = joinhandle_ref => { - map_subsystem_result(&name, result) - }, - _ = cancellation_token.cancelled() => { - inner_joinhandle.abort(); - map_subsystem_result(&name, inner_joinhandle.await) - } - }; + let future = async { subsystem(subsystem_handle).await.map_err(|e| e.into()) }; + let join_handle = tokio::spawn(future); - match &result { - Ok(()) | Err(SubsystemError::Cancelled(_)) => {} - Err(SubsystemError::Failed(name, e)) => { - log::error!("Error in subsystem '{}': {:?}", name, e); - if !local_shutdown_token.is_shutting_down() { - shutdown_token.shutdown(); - } + // Abort on drop + guard.on_cancel({ + let abort_handle = join_handle.abort_handle(); + let name = Arc::clone(&name); + move || { + if !abort_handle.is_finished() { + tracing::warn!("Subsystem cancelled: '{}'", name); } - Err(SubsystemError::Panicked(name)) => { - log::error!("Subsystem '{}' panicked", name); - if !local_shutdown_token.is_shutting_down() { - shutdown_token.shutdown(); - } - } - }; - - drop(shutdown_guard); - - result - } - - pub fn new> + Send>( - name: String, - shutdown_token: ShutdownToken, - local_shutdown_token: ShutdownToken, - cancellation_token: CancellationToken, - subsystem_future: Fut, - shutdown_guard: Arc, - ) -> Self { - // Spawn to nested tasks. - // This enables us to catch panics, as panics get returned through a JoinHandle. - let inner_joinhandle = tokio::spawn(subsystem_future); - let outer_joinhandle = tokio::spawn(Self::handle_subsystem( - inner_joinhandle, - shutdown_token, - local_shutdown_token, - name, - cancellation_token.clone(), - shutdown_guard, - )); + abort_handle.abort(); + } + }); - Self { - outer_joinhandle, - cancellation_token, + let failure = match join_handle.await { + Ok(Ok(())) => None, + Ok(Err(e)) => Some(SubsystemError::Failed(name, SubsystemFailure(e))), + Err(e) => { + if e.is_panic() { + Some(SubsystemError::Panicked(name)) + } else { + // Don't do anything in case of a cancellation; + // cancellations can't be forwarded (because the + // current function we are in will be cancelled + // simultaneously) + None + } } - } + }; - pub async fn join(&mut self) -> Result<(), SubsystemError> { - // Safety: we are in full control over the outer_joinhandle and the - // code it runs. Therefore, if this either returns a panic or a cancelled, - // it's a programming error on our side. - // Therefore using unwrap() here is the correct way of handling it. - // (this and the fact that unreachable code would decrease our test coverage) - (&mut self.outer_joinhandle).await.unwrap() - } + // Retrieve the handle that was passed into the subsystem. + // Originally it was intended to pass the handle as reference, but due + // to complications (https://stackoverflow.com/questions/77172947/async-lifetime-issues-of-pass-by-reference-parameters) + // it was decided to pass ownership instead. + // + // It is still important that the handle does not leak out of the subsystem. + let subsystem_handle = match redirected_subsystem_handle.try_recv() { + Ok(s) => s, + Err(_) => panic!("The SubsystemHandle object must not be leaked out of the subsystem!"), + }; - pub fn abort(&self) { - self.cancellation_token.cancel(); + // Raise potential errors + let joiner_token = subsystem_handle.joiner_token; + if let Some(failure) = failure { + joiner_token.raise_failure(failure); } + + // Wait for children to finish before we destroy the `SubsystemHandle` object. + // Otherwise the children would be cancelled immediately. + // + // This is the main mechanism that forwards a cancellation to all the children. + joiner_token.downgrade().join().await; } diff --git a/src/runner/alive_guard.rs b/src/runner/alive_guard.rs new file mode 100644 index 0000000..00cfddf --- /dev/null +++ b/src/runner/alive_guard.rs @@ -0,0 +1,120 @@ +use std::sync::{Arc, Mutex}; + +struct Inner { + finished_callback: Option>, + cancelled_callback: Option>, +} + +/// Allows registering callback functions that will get called on destruction. +/// +/// This struct is the mechanism that manages lifetime of parents and children +/// in the subsystem tree. It allows for cancellation of the subsytem on drop, +/// and for automatic deregistering in the parent when the child is finished. +pub(crate) struct AliveGuard { + inner: Arc>, +} +impl Clone for AliveGuard { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } + } +} + +impl AliveGuard { + pub(crate) fn new() -> Self { + Self { + inner: Arc::new(Mutex::new(Inner { + finished_callback: None, + cancelled_callback: None, + })), + } + } + + pub(crate) fn on_cancel(&self, cancelled_callback: impl FnOnce() + 'static + Send) { + let mut inner = self.inner.lock().unwrap(); + assert!(inner.cancelled_callback.is_none()); + inner.cancelled_callback = Some(Box::new(cancelled_callback)); + } + + pub(crate) fn on_finished(&self, finished_callback: impl FnOnce() + 'static + Send) { + let mut inner = self.inner.lock().unwrap(); + assert!(inner.finished_callback.is_none()); + inner.finished_callback = Some(Box::new(finished_callback)); + } +} + +impl Drop for Inner { + fn drop(&mut self) { + if let Some(finished_callback) = self.finished_callback.take() { + finished_callback(); + } else { + tracing::error!("No `finished` callback was registered in AliveGuard! This should not happen, please report this at https://github.com/Finomnis/tokio-graceful-shutdown/issues."); + } + + if let Some(cancelled_callback) = self.cancelled_callback.take() { + cancelled_callback() + } + } +} + +#[cfg(test)] +mod tests { + + use std::sync::atomic::{AtomicU32, Ordering}; + + use super::*; + + #[test] + fn finished_callback() { + let alive_guard = AliveGuard::new(); + + let counter = Arc::new(AtomicU32::new(0)); + let counter2 = Arc::clone(&counter); + + alive_guard.on_finished(move || { + counter2.fetch_add(1, Ordering::Relaxed); + }); + + drop(alive_guard); + + assert_eq!(counter.load(Ordering::Relaxed), 1); + } + + #[test] + fn cancel_callback() { + let alive_guard = AliveGuard::new(); + + let counter = Arc::new(AtomicU32::new(0)); + let counter2 = Arc::clone(&counter); + + alive_guard.on_finished(|| {}); + alive_guard.on_cancel(move || { + counter2.fetch_add(1, Ordering::Relaxed); + }); + + drop(alive_guard); + + assert_eq!(counter.load(Ordering::Relaxed), 1); + } + + #[test] + fn both_callbacks() { + let alive_guard = AliveGuard::new(); + + let counter = Arc::new(AtomicU32::new(0)); + let counter2 = Arc::clone(&counter); + let counter3 = Arc::clone(&counter); + + alive_guard.on_finished(move || { + counter2.fetch_add(1, Ordering::Relaxed); + }); + alive_guard.on_cancel(move || { + counter3.fetch_add(1, Ordering::Relaxed); + }); + + drop(alive_guard); + + assert_eq!(counter.load(Ordering::Relaxed), 2); + } +} diff --git a/src/shutdown_token.rs b/src/shutdown_token.rs deleted file mode 100644 index feaf487..0000000 --- a/src/shutdown_token.rs +++ /dev/null @@ -1,123 +0,0 @@ -use tokio_util::sync::{CancellationToken, WaitForCancellationFuture}; - -#[derive(Clone)] -#[doc(hidden)] -pub struct ShutdownToken { - token: CancellationToken, - is_toplevel: bool, -} - -pub fn create_shutdown_token() -> ShutdownToken { - ShutdownToken { - token: CancellationToken::new(), - is_toplevel: true, - } -} - -impl ShutdownToken { - pub fn shutdown(&self) { - if !self.token.is_cancelled() { - if self.is_toplevel { - log::info!("Initiating shutdown ..."); - } else { - log::debug!("Initiating partial shutdown ..."); - } - self.token.cancel() - } - } - - pub fn wait_for_shutdown(&self) -> WaitForCancellationFuture<'_> { - self.token.cancelled() - } - - pub fn is_shutting_down(&self) -> bool { - self.token.is_cancelled() - } - - pub fn child_token(&self) -> Self { - Self { - token: self.token.child_token(), - is_toplevel: false, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - use std::sync::atomic::{AtomicBool, Ordering}; - use tokio::time::{sleep, Duration}; - - #[tokio::test] - async fn triggers_correctly() { - let finished = AtomicBool::new(false); - - let token1 = create_shutdown_token(); - let token2 = token1.clone(); - - let stoppee = async { - token2.wait_for_shutdown().await; - finished.store(true, Ordering::SeqCst); - }; - - let stopper = async { - sleep(Duration::from_millis(100)).await; - assert!(!finished.load(Ordering::SeqCst)); - assert!(!token1.is_shutting_down()); - assert!(!token2.is_shutting_down()); - - token1.shutdown(); - sleep(Duration::from_millis(100)).await; - - assert!(finished.load(Ordering::SeqCst)); - assert!(token1.is_shutting_down()); - assert!(token2.is_shutting_down()); - }; - - tokio::join!(stopper, stoppee); - } - - #[tokio::test] - async fn triggers_correctly_on_partial() { - let finished = AtomicBool::new(false); - - let token1 = create_shutdown_token().child_token(); - let token2 = token1.clone(); - - let stoppee = async { - token2.wait_for_shutdown().await; - finished.store(true, Ordering::SeqCst); - }; - - let stopper = async { - sleep(Duration::from_millis(100)).await; - assert!(!finished.load(Ordering::SeqCst)); - assert!(!token1.is_shutting_down()); - assert!(!token2.is_shutting_down()); - - token1.shutdown(); - sleep(Duration::from_millis(100)).await; - - assert!(finished.load(Ordering::SeqCst)); - assert!(token1.is_shutting_down()); - assert!(token2.is_shutting_down()); - }; - - tokio::join!(stopper, stoppee); - } - - #[tokio::test] - async fn double_shutdown_causes_no_error() { - let token1 = create_shutdown_token(); - let token2 = create_shutdown_token().child_token(); - - token1.shutdown(); - token1.shutdown(); - token2.shutdown(); - token2.shutdown(); - - assert!(token1.is_shutting_down()); - assert!(token2.is_shutting_down()); - } -} diff --git a/src/signal_handling.rs b/src/signal_handling.rs index 61982fd..5a1b5c6 100644 --- a/src/signal_handling.rs +++ b/src/signal_handling.rs @@ -7,8 +7,8 @@ async fn wait_for_signal_impl() { let mut signal_interrupt = signal(SignalKind::interrupt()).unwrap(); tokio::select! { - _ = signal_terminate.recv() => log::debug!("Received SIGTERM."), - _ = signal_interrupt.recv() => log::debug!("Received SIGINT."), + _ = signal_terminate.recv() => tracing::debug!("Received SIGTERM."), + _ = signal_interrupt.recv() => tracing::debug!("Received SIGINT."), }; } @@ -18,12 +18,12 @@ async fn wait_for_signal_impl() { use tokio::signal::ctrl_c; ctrl_c().await.unwrap(); - log::debug!("Received SIGINT."); + tracing::debug!("Received SIGINT."); } /// Registers Ctrl+C and SIGTERM handlers to cause a program shutdown. /// Further, registers a custom panic handler to also initiate a shutdown. /// Otherwise, a multi-threaded system would deadlock on panik. -pub async fn wait_for_signal() { +pub(crate) async fn wait_for_signal() { wait_for_signal_impl().await } diff --git a/src/subsystem/data.rs b/src/subsystem/data.rs deleted file mode 100644 index e7d341b..0000000 --- a/src/subsystem/data.rs +++ /dev/null @@ -1,246 +0,0 @@ -use std::sync::Arc; -use std::sync::Weak; -use tokio::sync::MutexGuard; - -use async_recursion::async_recursion; -use futures::future::join; -use futures::future::join_all; -use std::sync::Mutex; -use tokio_util::sync::CancellationToken; - -use super::NestedSubsystem; -use super::PartialShutdownError; -use super::SubsystemData; -use super::SubsystemDescriptor; -use super::SubsystemIdentifier; -use crate::errors::SubsystemError; -use crate::exit_state::prettify_exit_states; -use crate::exit_state::{join_shutdown_results, ShutdownResults, SubprocessExitState}; -use crate::runner::SubsystemRunner; -use crate::shutdown_token::ShutdownToken; -use crate::utils::ShutdownGuard; -use crate::ErrTypeTraits; - -impl SubsystemData { - pub fn new( - name: &str, - global_shutdown_token: ShutdownToken, - group_shutdown_token: ShutdownToken, - local_shutdown_token: ShutdownToken, - cancellation_token: CancellationToken, - shutdown_guard: Weak, - ) -> Self { - Self { - name: name.to_string(), - subsystems: Mutex::new(Some(Vec::new())), - global_shutdown_token, - group_shutdown_token, - local_shutdown_token, - cancellation_token, - shutdown_subsystems: tokio::sync::Mutex::new(Vec::new()), - shutdown_guard, - } - } - - /// Registers a new subsystem in self.subsystems. - /// - /// If a shutdown is already running, self.subsystems will be 'None', - /// and the newly spawned subsystem will be cancelled. - pub fn add_subsystem( - &self, - subsystem: Arc>, - subsystem_runner: SubsystemRunner, - ) -> SubsystemIdentifier { - let id = SubsystemIdentifier::create(); - match self.subsystems.lock().unwrap().as_mut() { - Some(subsystems) => { - subsystems.push(SubsystemDescriptor { - id: id.clone(), - subsystem_runner, - data: subsystem, - }); - } - None => { - log::error!("Unable to add subsystem, parent subsystem already shutting down!"); - subsystem_runner.abort(); - } - } - id - } - - /// Moves all subsystem descriptors to the self.shutdown_subsystem vector. - /// This indicates to the subsystem that it should no longer be possible to - /// spawn new nested subsystems. - /// - /// This is achieved by writing 'None' to self.subsystems. - /// - /// Preventing new nested subsystems to be registered is important to avoid - /// a race condition where the subsystem could spawn a nested subsystem by calling - /// [`SubsystemHandle.start`] during cleanup, leaking the new nested subsystem. - /// - /// (The place where adding new subsystems will fail is in [`SubsystemData.add_subsystem`]) - async fn prepare_shutdown(&self) -> MutexGuard<'_, Vec>> { - let mut shutdown_subsystems = self.shutdown_subsystems.lock().await; - let mut subsystems = self.subsystems.lock().unwrap(); - if let Some(e) = subsystems.take() { - shutdown_subsystems.extend(e.into_iter()) - }; - shutdown_subsystems - } - - /// Recursively goes through all given subsystems, awaits their join handles, - /// and collects their exit states. - /// - /// Returns the collected subsystem exit states. - /// - /// This function can handle cancellation. - #[async_recursion] - async fn perform_shutdown_on_subsystems( - subsystems: &mut [SubsystemDescriptor], - ) -> ShutdownResults { - let mut subsystem_runners = vec![]; - let mut subsystem_data = vec![]; - for SubsystemDescriptor { - id: _, - subsystem_runner, - data, - } in subsystems.iter_mut() - { - subsystem_runners.push((data.name.clone(), subsystem_runner)); - subsystem_data.push(data); - } - let joinhandles_finished = join_all( - subsystem_runners - .iter_mut() - .map( - |(name, subsystem_runner)| async move { (name, subsystem_runner.join().await) }, - ), - ); - let subsystems_finished = join_all( - subsystem_data - .iter_mut() - .map(|data| data.perform_shutdown()), - ); - - let (results_direct, results_recursive) = join( - async { - let joinhandles_finished = joinhandles_finished.await; - - joinhandles_finished - .into_iter() - .map(|(name, result)| { - SubprocessExitState::::new( - name, - match &result { - Ok(()) => "OK", - Err(SubsystemError::Cancelled(_)) => "Cancelled", - Err(SubsystemError::Failed(_, _)) => "Failed", - Err(SubsystemError::Panicked(_)) => "Panicked", - }, - result, - ) - }) - .collect() - }, - subsystems_finished, - ) - .await; - - join_shutdown_results(results_direct, results_recursive) - } - - /// Recursively goes through all subsystems, awaits their join handles, - /// and collects their exit states. - /// - /// Returns the collected subsystem exit states. - /// - /// This function can handle cancellation. - pub async fn perform_shutdown(&self) -> ShutdownResults { - let mut subsystems = self.prepare_shutdown().await; - - SubsystemData::perform_shutdown_on_subsystems(&mut subsystems).await - } - - pub fn cancel_all_subsystems(&self) { - self.cancellation_token.cancel(); - } - - pub async fn perform_partial_shutdown( - &self, - subsystem_handle: NestedSubsystem, - ) -> Result<(), PartialShutdownError> { - let subsystem = { - let mut subsystems_mutex = self.subsystems.lock().unwrap(); - let subsystems = subsystems_mutex - .as_mut() - .ok_or(PartialShutdownError::AlreadyShuttingDown)?; - let position = subsystems - .iter() - .position(|elem| elem.id == subsystem_handle.id) - .ok_or(PartialShutdownError::SubsystemNotFound)?; - subsystems.swap_remove(position) - }; - - // Initiate shutdown - subsystem.data.local_shutdown_token.shutdown(); - - // Wait for shutdown to finish - let mut subsystem_vec = vec![subsystem]; - let exit_states = SubsystemData::perform_shutdown_on_subsystems(&mut subsystem_vec).await; - - // Prettify exit states - let formatted_exit_states = prettify_exit_states(&exit_states); - - // Collect failed subsystems - let failed_subsystems = exit_states - .into_iter() - .filter_map(|exit_state| match exit_state.raw_result { - Ok(()) => None, - Err(e) => Some(e), - }) - .collect::>(); - - // Print subsystem exit states - if failed_subsystems.is_empty() { - log::debug!("Partial shutdown successful. Subsystem states:"); - } else { - log::debug!("Some subsystems during partial shutdown failed. Subsystem states:"); - }; - for formatted_exit_state in formatted_exit_states { - log::debug!(" {}", formatted_exit_state); - } - - if failed_subsystems.is_empty() { - Ok(()) - } else { - Err(PartialShutdownError::SubsystemsFailed(failed_subsystems)) - } - } -} - -#[cfg(test)] -mod tests { - use crate::{shutdown_token::create_shutdown_token, BoxedError}; - - use super::*; - - #[tokio::test] - async fn prepare_shutdown_does_not_crash_when_called_twice() { - let shutdown_token = create_shutdown_token(); - let shutdown_guard = Arc::new(ShutdownGuard::new(shutdown_token.clone())); - - let data = SubsystemData::::new( - "MySubsys", - shutdown_token.clone(), - shutdown_token.clone(), - shutdown_token.clone(), - CancellationToken::new(), - Arc::downgrade(&shutdown_guard), - ); - - drop(data.prepare_shutdown().await); - drop(data.prepare_shutdown().await); - - assert!(data.subsystems.lock().unwrap().is_none()); - } -} diff --git a/src/subsystem/error_collector.rs b/src/subsystem/error_collector.rs new file mode 100644 index 0000000..93208b0 --- /dev/null +++ b/src/subsystem/error_collector.rs @@ -0,0 +1,43 @@ +use std::sync::Arc; + +use tokio::sync::mpsc; + +use crate::{errors::SubsystemError, ErrTypeTraits}; + +pub(crate) enum ErrorCollector { + Collecting(mpsc::UnboundedReceiver>), + Finished(Arc<[SubsystemError]>), +} + +impl ErrorCollector { + pub(crate) fn new(receiver: mpsc::UnboundedReceiver>) -> Self { + Self::Collecting(receiver) + } + + pub(crate) fn finish(&mut self) -> Arc<[SubsystemError]> { + match self { + ErrorCollector::Collecting(receiver) => { + let mut errors = vec![]; + receiver.close(); + while let Ok(e) = receiver.try_recv() { + errors.push(e); + } + let errors = errors.into_boxed_slice().into(); + *self = ErrorCollector::Finished(Arc::clone(&errors)); + errors + } + ErrorCollector::Finished(errors) => Arc::clone(errors), + } + } +} + +impl Drop for ErrorCollector { + fn drop(&mut self) { + if let Self::Collecting(receiver) = self { + receiver.close(); + while let Ok(e) = receiver.try_recv() { + tracing::warn!("An error got dropped: {e:?}"); + } + } + } +} diff --git a/src/subsystem/handle.rs b/src/subsystem/handle.rs deleted file mode 100644 index cf310b1..0000000 --- a/src/subsystem/handle.rs +++ /dev/null @@ -1,330 +0,0 @@ -use std::future::Future; -use std::sync::Arc; - -use super::NestedSubsystem; -use super::SubsystemData; -use super::SubsystemHandle; -use crate::errors::PartialShutdownError; -use crate::runner::SubsystemRunner; -use crate::ErrTypeTraits; -use crate::ShutdownToken; - -use crate::utils::get_subsystem_name; -#[cfg(doc)] -use crate::Toplevel; - -impl SubsystemHandle { - #[doc(hidden)] - pub fn new(data: Arc>) -> Self { - Self { data } - } - - /// Starts a nested subsystem, analogous to [`Toplevel::start`]. - /// - /// Once called, the subsystem will be started immediately, similar to [`tokio::spawn`]. - /// - /// # Arguments - /// - /// * `name` - The name of the subsystem - /// * `subsystem` - The subsystem to be started - /// - /// # Returns - /// - /// A [`NestedSubsystem`] that can be used to perform a partial shutdown - /// on the created submodule. - /// - /// # Examples - /// - /// ``` - /// use miette::Result; - /// use tokio_graceful_shutdown::SubsystemHandle; - /// - /// async fn nested_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// subsys.on_shutdown_requested().await; - /// Ok(()) - /// } - /// - /// async fn my_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// // start a nested subsystem - /// subsys.start("Nested", nested_subsystem); - /// - /// subsys.on_shutdown_requested().await; - /// Ok(()) - /// } - /// ``` - /// - pub fn start(&self, name: &str, subsystem: Subsys) -> NestedSubsystem - where - Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, - Fut: 'static + Future> + Send, - Err: Into, - { - let name = get_subsystem_name(&self.data.name, name); - - // When we are inside a subsystem, shutdown_guard cannot have gotten dropped, because - // the SubsystemRunner of the current subsystem keeps it alive. - let shutdown_guard = self - .data - .shutdown_guard - .upgrade() - .expect("'start()' called from outside a subsystem"); - - // Create subsystem data structure - let new_subsystem = Arc::new(SubsystemData::new( - &name, - self.global_shutdown_token().clone(), - self.group_shutdown_token().clone(), - self.local_shutdown_token().child_token(), - self.data.cancellation_token.child_token(), - self.data.shutdown_guard.clone(), - )); - - // Create handle - let subsystem_handle = SubsystemHandle::new(new_subsystem.clone()); - - // Shutdown token - let shutdown_token = subsystem_handle.group_shutdown_token().clone(); - - // Future - let subsystem_future = async { subsystem(subsystem_handle).await.map_err(|e| e.into()) }; - - // Spawn new task - let subsystem_runner = SubsystemRunner::new( - name, - shutdown_token, - new_subsystem.local_shutdown_token.child_token(), - new_subsystem.cancellation_token.child_token(), - subsystem_future, - shutdown_guard, - ); - - // Store subsystem data - let id = self.data.add_subsystem(new_subsystem, subsystem_runner); - - NestedSubsystem { id } - } - - /// Wait for the shutdown mode to be triggered. - /// - /// Once the shutdown mode is entered, all existing calls to this - /// method will be released and future calls to this method will - /// return immediately. - /// - /// This is the primary method of subsystems to react to - /// the shutdown requests. Most often, it will be used in `tokio::select` - /// statements to cancel other code as soon as the shutdown is requested. - /// - /// # Examples - /// - /// ``` - /// use miette::Result; - /// use tokio::time::{sleep, Duration}; - /// use tokio_graceful_shutdown::SubsystemHandle; - /// - /// async fn countdown() { - /// for i in (1..10).rev() { - /// log::info!("Countdown: {}", i); - /// sleep(Duration::from_millis(1000)).await; - /// } - /// } - /// - /// async fn countdown_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// log::info!("Starting countdown ..."); - /// - /// // This cancels the countdown as soon as shutdown - /// // mode was entered - /// tokio::select! { - /// _ = subsys.on_shutdown_requested() => { - /// log::info!("Countdown cancelled."); - /// }, - /// _ = countdown() => { - /// log::info!("Countdown finished."); - /// } - /// }; - /// - /// Ok(()) - /// } - /// ``` - pub async fn on_shutdown_requested(&self) { - self.data.local_shutdown_token.wait_for_shutdown().await - } - - /// Returns whether a shutdown should be performed now. - /// - /// This method is provided for subsystems that need to query the shutdown - /// request state repeatedly. - /// - /// This can be useful in scenarios where a subsystem depends on the graceful - /// shutdown of its nested coroutines before it can run final cleanup steps itself. - /// - /// # Examples - /// - /// ``` - /// use miette::Result; - /// use tokio::time::{sleep, Duration}; - /// use tokio_graceful_shutdown::SubsystemHandle; - /// - /// async fn uncancellable_action(subsys: &SubsystemHandle) { - /// tokio::select! { - /// // Execute an action. A dummy `sleep` in this case. - /// _ = sleep(Duration::from_millis(1000)) => { - /// log::info!("Action finished."); - /// } - /// // Perform a shutdown if requested - /// _ = subsys.on_shutdown_requested() => { - /// log::info!("Action aborted."); - /// }, - /// } - /// } - /// - /// async fn my_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// log::info!("Starting subsystem ..."); - /// - /// // We cannot do a `tokio::select` with `on_shutdown_requested` - /// // here, because a shutdown would cancel the action without giving - /// // it the chance to react first. - /// while !subsys.is_shutdown_requested() { - /// uncancellable_action(&subsys).await; - /// } - /// - /// log::info!("Subsystem stopped."); - /// - /// Ok(()) - /// } - /// ``` - pub fn is_shutdown_requested(&self) -> bool { - self.data.local_shutdown_token.is_shutting_down() - } - - /// Triggers a shutdown. - /// - /// This version only propagates up to the next [Toplevel] object. - /// To initiate a shutdown for the entire program, see [request_global_shutdown()](SubsystemHandle::request_global_shutdown). - /// - /// # Examples - /// - /// ``` - /// use miette::Result; - /// use tokio::time::{sleep, Duration}; - /// use tokio_graceful_shutdown::SubsystemHandle; - /// - /// async fn stop_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// // This subsystem wait for one second and then stops the program. - /// sleep(Duration::from_millis(1000)).await; - /// - /// // Shut down the parent `Toplevel` object - /// subsys.request_shutdown(); - /// - /// Ok(()) - /// } - /// ``` - pub fn request_shutdown(&self) { - self.data.group_shutdown_token.shutdown() - } - - /// Triggers the shutdown of the entire program. - /// - /// # Examples - /// - /// ``` - /// use miette::Result; - /// use tokio::time::{sleep, Duration}; - /// use tokio_graceful_shutdown::SubsystemHandle; - /// - /// async fn stop_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// // This subsystem wait for one second and then stops the program. - /// sleep(Duration::from_millis(1000)).await; - /// - /// // Shut down all parent `Toplevel` objects. - /// subsys.request_global_shutdown(); - /// - /// Ok(()) - /// } - /// ``` - pub fn request_global_shutdown(&self) { - self.data.global_shutdown_token.shutdown() - } - - /// Preforms a partial shutdown of the given nested subsystem. - /// - /// # Arguments - /// - /// * `subsystem` - The nested subsystem that should be shut down - /// - /// # Returns - /// - /// A [`PartialShutdownError`] on failure. - /// - /// # Examples - /// - /// ``` - /// use miette::Result; - /// use tokio::time::{sleep, Duration}; - /// use tokio_graceful_shutdown::SubsystemHandle; - /// - /// async fn nested_subsystem(subsys: SubsystemHandle) -> Result<()> { - /// // This subsystem does nothing but wait for the shutdown to happen - /// subsys.on_shutdown_requested().await; - /// Ok(()) - /// } - /// - /// async fn subsystem(subsys: SubsystemHandle) -> Result<()> { - /// // This subsystem waits for one second and then performs a partial shutdown - /// - /// // Spawn nested subsystem - /// let nested = subsys.start("nested", nested_subsystem); - /// - /// // Wait for a second - /// sleep(Duration::from_millis(1000)).await; - /// - /// // Perform a partial shutdown of the nested subsystem - /// subsys.perform_partial_shutdown(nested).await?; - /// - /// Ok(()) - /// } - /// ``` - pub async fn perform_partial_shutdown( - &self, - subsystem: NestedSubsystem, - ) -> Result<(), PartialShutdownError> { - self.data.perform_partial_shutdown(subsystem).await - } - - /// Provides access to the process-wide parent shutdown token. - /// - /// This function is usually not required and is there - /// to provide lower-level access for specific corner cases. - #[doc(hidden)] - pub fn global_shutdown_token(&self) -> &ShutdownToken { - &self.data.global_shutdown_token - } - - /// Provides access to the group local shutdown token. - /// - /// This token shuts down the parent [Toplevel] object. - /// - /// This function is usually not required and is there - /// to provide lower-level access for specific corner cases. - #[doc(hidden)] - pub fn group_shutdown_token(&self) -> &ShutdownToken { - &self.data.group_shutdown_token - } - - /// Provides access to the subsystem local shutdown token. - /// - /// This function is usually not required and is there - /// to provide lower-level access for specific corner cases. - #[doc(hidden)] - pub fn local_shutdown_token(&self) -> &ShutdownToken { - &self.data.local_shutdown_token - } - - /// The name of the subsystem. - /// - /// This function is usually not required and is there - /// to provide lower-level access for specific corner cases. - #[doc(hidden)] - pub fn name(&self) -> &str { - &self.data.name - } -} diff --git a/src/subsystem/identifier.rs b/src/subsystem/identifier.rs deleted file mode 100644 index 0ab5f89..0000000 --- a/src/subsystem/identifier.rs +++ /dev/null @@ -1,36 +0,0 @@ -use std::sync::atomic::{AtomicUsize, Ordering}; - -#[derive(PartialEq, Eq, Debug, Clone)] -pub struct SubsystemIdentifier { - id: usize, -} - -static NEXT_ID: AtomicUsize = AtomicUsize::new(1); - -impl SubsystemIdentifier { - pub fn create() -> Self { - Self { - id: NEXT_ID.fetch_add(1, Ordering::SeqCst), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn equals_with_itself() { - let identifier1 = SubsystemIdentifier::create(); - #[allow(clippy::redundant_clone)] - let identifier2 = identifier1.clone(); - assert_eq!(identifier1, identifier2); - } - - #[test] - fn does_not_equal_with_others() { - let identifier1 = SubsystemIdentifier::create(); - let identifier2 = SubsystemIdentifier::create(); - assert_ne!(identifier1, identifier2); - } -} diff --git a/src/subsystem/mod.rs b/src/subsystem/mod.rs index 3f6101f..7073c1e 100644 --- a/src/subsystem/mod.rs +++ b/src/subsystem/mod.rs @@ -1,58 +1,37 @@ -mod data; -mod handle; -mod identifier; +mod error_collector; +mod nested_subsystem; +mod subsystem_builder; +mod subsystem_handle; -use std::sync::Arc; -use std::sync::Mutex; -use std::sync::Weak; +use std::sync::{Arc, Mutex}; -use tokio_util::sync::CancellationToken; - -use crate::errors::PartialShutdownError; -use crate::runner::SubsystemRunner; -use crate::shutdown_token::ShutdownToken; -use crate::utils::ShutdownGuard; -use crate::ErrTypeTraits; +pub use subsystem_builder::SubsystemBuilder; +pub use subsystem_handle::SubsystemHandle; -use self::identifier::SubsystemIdentifier; +pub(crate) use subsystem_handle::root_handle; -/// The data stored per subsystem, like name or nested subsystems -pub struct SubsystemData { - name: String, - subsystems: Mutex>>>, - shutdown_subsystems: tokio::sync::Mutex>>, - local_shutdown_token: ShutdownToken, - group_shutdown_token: ShutdownToken, - global_shutdown_token: ShutdownToken, - cancellation_token: CancellationToken, - shutdown_guard: Weak, -} +use crate::{utils::JoinerTokenRef, ErrTypeTraits, ErrorAction}; -/// The handle given to each subsystem through which the subsystem can interact with this crate. -pub struct SubsystemHandle { - data: Arc>, -} -// Implement `Clone` manually because the compiler cannot derive `Clone -// from Generics that don't implement `Clone`. -// (https://stackoverflow.com/questions/72150623/) -impl Clone for SubsystemHandle { - fn clone(&self) -> Self { - Self { - data: self.data.clone(), - } - } -} +use atomic::Atomic; +use tokio_util::sync::CancellationToken; -/// A running subsystem. Can be used to stop the subsystem or get its return value. -struct SubsystemDescriptor { - id: SubsystemIdentifier, - data: Arc>, - subsystem_runner: SubsystemRunner, +/// A nested subsystem. +/// +/// Can be used to control the subsystem or wait for it to finish. +/// +/// Dropping this value does not perform any action - the subsystem +/// will be neither cancelled, shut down or detached. +/// +/// For more information, look through the examples directory in +/// the source code. +pub struct NestedSubsystem { + joiner: JoinerTokenRef, + cancellation_token: CancellationToken, + errors: Mutex>, + error_actions: Arc, } -/// A nested subsystem. Can be used to perform a partial shutdown. -/// -/// For more information, see [`SubsystemHandle::start()`] and [`SubsystemHandle::perform_partial_shutdown()`]. -pub struct NestedSubsystem { - id: SubsystemIdentifier, +pub(crate) struct ErrorActions { + pub(crate) on_failure: Atomic, + pub(crate) on_panic: Atomic, } diff --git a/src/subsystem/nested_subsystem.rs b/src/subsystem/nested_subsystem.rs new file mode 100644 index 0000000..cfc5378 --- /dev/null +++ b/src/subsystem/nested_subsystem.rs @@ -0,0 +1,85 @@ +use std::sync::atomic::Ordering; + +use crate::{errors::SubsystemJoinError, ErrTypeTraits, ErrorAction}; + +use super::NestedSubsystem; + +impl NestedSubsystem { + /// Wait for the subsystem to be finished. + /// + /// If its failure/panic action is set to [`ErrorAction::CatchAndLocalShutdown`], + /// this function will return the list of errors caught by the subsystem. + /// + /// # Returns + /// + /// A [`SubsystemJoinError`] on failure. + /// + /// # Examples + /// + /// ``` + /// use miette::Result; + /// use tokio::time::{sleep, Duration}; + /// use tokio_graceful_shutdown::{ErrorAction, SubsystemBuilder, SubsystemHandle}; + /// + /// async fn nested_subsystem(subsys: SubsystemHandle) -> Result<()> { + /// // This subsystem does nothing but wait for the shutdown to happen + /// subsys.on_shutdown_requested().await; + /// Ok(()) + /// } + /// + /// async fn subsystem(subsys: SubsystemHandle) -> Result<()> { + /// // This subsystem waits for one second and then performs a partial shutdown + /// + /// // Spawn nested subsystem. + /// // Make sure to catch errors, so that they are properly + /// // returned at `.join()`. + /// let nested = subsys.start( + /// SubsystemBuilder::new("nested", nested_subsystem) + /// .on_failure(ErrorAction::CatchAndLocalShutdown) + /// .on_panic(ErrorAction::CatchAndLocalShutdown) + /// ); + /// + /// // Wait for a second + /// sleep(Duration::from_millis(1000)).await; + /// + /// // Perform a partial shutdown of the nested subsystem + /// nested.initiate_shutdown(); + /// nested.join().await?; + /// + /// Ok(()) + /// } + /// ``` + pub async fn join(&self) -> Result<(), SubsystemJoinError> { + self.joiner.join().await; + + let errors = self.errors.lock().unwrap().finish(); + if errors.is_empty() { + Ok(()) + } else { + Err(SubsystemJoinError::SubsystemsFailed(errors)) + } + } + + /// Signals the subsystem and all of its children to shut down. + pub fn initiate_shutdown(&self) { + self.cancellation_token.cancel() + } + + /// Changes the way this subsystem should react to failures, + /// meaning if it or one of its children returns an `Err` value. + /// + /// For more information, see [ErrorAction]. + pub fn change_failure_action(&self, action: ErrorAction) { + self.error_actions + .on_failure + .store(action, Ordering::Relaxed); + } + + /// Changes the way this subsystem should react if it or one + /// of its children panic. + /// + /// For more information, see [ErrorAction]. + pub fn change_panic_action(&self, action: ErrorAction) { + self.error_actions.on_panic.store(action, Ordering::Relaxed); + } +} diff --git a/src/subsystem/subsystem_builder.rs b/src/subsystem/subsystem_builder.rs new file mode 100644 index 0000000..e2d0e73 --- /dev/null +++ b/src/subsystem/subsystem_builder.rs @@ -0,0 +1,68 @@ +use std::{borrow::Cow, future::Future, marker::PhantomData}; + +use crate::{ErrTypeTraits, ErrorAction, SubsystemHandle}; + +/// Configures a subsystem before it gets spawned through +/// [`SubsystemHandle::start`]. +pub struct SubsystemBuilder<'a, ErrType, Err, Fut, Subsys> +where + ErrType: ErrTypeTraits, + Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, + Fut: 'static + Future> + Send, + Err: Into, +{ + pub(crate) name: Cow<'a, str>, + pub(crate) subsystem: Subsys, + pub(crate) failure_action: ErrorAction, + pub(crate) panic_action: ErrorAction, + #[allow(clippy::type_complexity)] + _phantom: PhantomData (Fut, ErrType, Err)>, +} + +impl<'a, ErrType, Err, Fut, Subsys> SubsystemBuilder<'a, ErrType, Err, Fut, Subsys> +where + ErrType: ErrTypeTraits, + Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, + Fut: 'static + Future> + Send, + Err: Into, +{ + /// Creates a new SubsystemBuilder from a given subsystem + /// function. + /// + /// # Arguments + /// + /// * `name` - The name of the subsystem. Primarily to identify the + /// subsystem in error messages. + /// * `subsystem` - The subsystem function that the subsystem will execute. + pub fn new(name: impl Into>, subsystem: Subsys) -> Self { + Self { + name: name.into(), + subsystem, + failure_action: ErrorAction::Forward, + panic_action: ErrorAction::Forward, + _phantom: Default::default(), + } + } + + /// Sets the way this subsystem should react to failures, + /// meaning if it or one of its children return an `Err` value. + /// + /// The default is [`ErrorAction::Forward`]. + /// + /// For more information, see [`ErrorAction`]. + pub fn on_failure(mut self, action: ErrorAction) -> Self { + self.failure_action = action; + self + } + + /// Sets the way this subsystem should react if it or one + /// of its children panic. + /// + /// The default is [`ErrorAction::Forward`]. + /// + /// For more information, see [`ErrorAction`]. + pub fn on_panic(mut self, action: ErrorAction) -> Self { + self.panic_action = action; + self + } +} diff --git a/src/subsystem/subsystem_handle.rs b/src/subsystem/subsystem_handle.rs new file mode 100644 index 0000000..ede0fbf --- /dev/null +++ b/src/subsystem/subsystem_handle.rs @@ -0,0 +1,430 @@ +use std::{ + future::Future, + mem::ManuallyDrop, + sync::{atomic::Ordering, Arc, Mutex}, +}; + +use atomic::Atomic; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; + +use crate::{ + errors::SubsystemError, + runner::{AliveGuard, SubsystemRunner}, + utils::{remote_drop_collection::RemotelyDroppableItems, JoinerToken}, + BoxedError, ErrTypeTraits, ErrorAction, NestedSubsystem, SubsystemBuilder, +}; + +use super::{error_collector::ErrorCollector, ErrorActions}; + +struct Inner { + name: Arc, + cancellation_token: CancellationToken, + toplevel_cancellation_token: CancellationToken, + joiner_token: JoinerToken, + children: RemotelyDroppableItems, +} + +/// The handle given to each subsystem through which the subsystem can interact with this crate. +pub struct SubsystemHandle { + inner: ManuallyDrop>, + // When dropped, redirect Self into this channel. + // Required as a workaround for https://stackoverflow.com/questions/77172947/async-lifetime-issues-of-pass-by-reference-parameters. + drop_redirect: Option>>, +} + +pub(crate) struct WeakSubsystemHandle { + pub(crate) joiner_token: JoinerToken, + // Children are stored here to keep them alive + _children: RemotelyDroppableItems, +} + +impl SubsystemHandle { + /// Start a nested subsystem. + /// + /// Once called, the subsystem will be started immediately, similar to [`tokio::spawn`]. + /// + /// # Arguments + /// + /// * `builder` - The [`SubsystemBuilder`] that contains all the information + /// about the subsystem that should be spawned. + /// + /// # Returns + /// + /// A [`NestedSubsystem`] that can be used to control or join the subsystem. + /// + /// # Examples + /// + /// ``` + /// use miette::Result; + /// use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle}; + /// + /// async fn nested_subsystem(subsys: SubsystemHandle) -> Result<()> { + /// subsys.on_shutdown_requested().await; + /// Ok(()) + /// } + /// + /// async fn my_subsystem(subsys: SubsystemHandle) -> Result<()> { + /// // start a nested subsystem + /// subsys.start(SubsystemBuilder::new("Nested", nested_subsystem)); + /// + /// subsys.on_shutdown_requested().await; + /// Ok(()) + /// } + /// ``` + pub fn start( + &self, + builder: SubsystemBuilder, + ) -> NestedSubsystem + where + Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, + Fut: 'static + Future> + Send, + Err: Into, + { + self.start_with_abs_name( + Arc::from(format!("{}/{}", self.inner.name, builder.name)), + builder.subsystem, + ErrorActions { + on_failure: Atomic::new(builder.failure_action), + on_panic: Atomic::new(builder.panic_action), + }, + ) + } + + pub(crate) fn start_with_abs_name( + &self, + name: Arc, + subsystem: Subsys, + error_actions: ErrorActions, + ) -> NestedSubsystem + where + Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, + Fut: 'static + Future> + Send, + Err: Into, + { + let alive_guard = AliveGuard::new(); + + let (error_sender, errors) = mpsc::unbounded_channel(); + + let cancellation_token = self.inner.cancellation_token.child_token(); + + let error_actions = Arc::new(error_actions); + + let (joiner_token, joiner_token_ref) = self.inner.joiner_token.child_token({ + let cancellation_token = cancellation_token.clone(); + let error_actions = Arc::clone(&error_actions); + move |e| { + let error_action = match &e { + SubsystemError::Failed(_, _) => { + error_actions.on_failure.load(Ordering::Relaxed) + } + SubsystemError::Panicked(_) => error_actions.on_panic.load(Ordering::Relaxed), + }; + + match error_action { + ErrorAction::Forward => Some(e), + ErrorAction::CatchAndLocalShutdown => { + if let Err(mpsc::error::SendError(e)) = error_sender.send(e) { + tracing::warn!("An error got dropped: {e:?}"); + }; + cancellation_token.cancel(); + None + } + } + } + }); + + let child_handle = SubsystemHandle { + inner: ManuallyDrop::new(Inner { + name: Arc::clone(&name), + cancellation_token: cancellation_token.clone(), + toplevel_cancellation_token: self.inner.toplevel_cancellation_token.clone(), + joiner_token, + children: RemotelyDroppableItems::new(), + }), + drop_redirect: None, + }; + + let runner = SubsystemRunner::new(name, subsystem, child_handle, alive_guard.clone()); + + // Shenanigans to juggle child ownership + // + // RACE CONDITION SAFETY: + // If the subsystem ends before `on_finished` was able to be called, nothing bad happens. + // alive_guard will keep the guard alive and the callback will only be called inside of + // the guard's drop() implementation. + let child_dropper = self.inner.children.insert(runner); + alive_guard.on_finished(|| { + drop(child_dropper); + }); + + NestedSubsystem { + joiner: joiner_token_ref, + cancellation_token, + errors: Mutex::new(ErrorCollector::new(errors)), + error_actions, + } + } + + /// Waits until all the children of this subsystem are finished. + pub async fn wait_for_children(&mut self) { + self.inner.joiner_token.join_children().await + } + + // For internal use only - should never be used by users. + // Required as a short-lived second reference inside of `runner`. + pub(crate) fn delayed_clone(&mut self) -> oneshot::Receiver> { + let (sender, receiver) = oneshot::channel(); + + let previous = self.drop_redirect.replace(sender); + assert!(previous.is_none()); + + receiver + } + + /// Wait for the shutdown mode to be triggered. + /// + /// Once the shutdown mode is entered, all existing calls to this + /// method will be released and future calls to this method will + /// return immediately. + /// + /// This is the primary method of subsystems to react to + /// the shutdown requests. Most often, it will be used in [`tokio::select`] + /// statements to cancel other code as soon as the shutdown is requested. + /// + /// # Examples + /// + /// ``` + /// use miette::Result; + /// use tokio::time::{sleep, Duration}; + /// use tokio_graceful_shutdown::SubsystemHandle; + /// + /// async fn countdown() { + /// for i in (1..10).rev() { + /// tracing::info!("Countdown: {}", i); + /// sleep(Duration::from_millis(1000)).await; + /// } + /// } + /// + /// async fn countdown_subsystem(subsys: SubsystemHandle) -> Result<()> { + /// tracing::info!("Starting countdown ..."); + /// + /// // This cancels the countdown as soon as shutdown + /// // mode was entered + /// tokio::select! { + /// _ = subsys.on_shutdown_requested() => { + /// tracing::info!("Countdown cancelled."); + /// }, + /// _ = countdown() => { + /// tracing::info!("Countdown finished."); + /// } + /// }; + /// + /// Ok(()) + /// } + /// ``` + pub async fn on_shutdown_requested(&self) { + self.inner.cancellation_token.cancelled().await + } + + /// Returns whether a shutdown should be performed now. + /// + /// This method is provided for subsystems that need to query the shutdown + /// request state repeatedly. + /// + /// This can be useful in scenarios where a subsystem depends on the graceful + /// shutdown of its nested coroutines before it can run final cleanup steps itself. + /// + /// # Examples + /// + /// ``` + /// use miette::Result; + /// use tokio::time::{sleep, Duration}; + /// use tokio_graceful_shutdown::SubsystemHandle; + /// + /// async fn uncancellable_action(subsys: &SubsystemHandle) { + /// tokio::select! { + /// // Execute an action. A dummy `sleep` in this case. + /// _ = sleep(Duration::from_millis(1000)) => { + /// tracing::info!("Action finished."); + /// } + /// // Perform a shutdown if requested + /// _ = subsys.on_shutdown_requested() => { + /// tracing::info!("Action aborted."); + /// }, + /// } + /// } + /// + /// async fn my_subsystem(subsys: SubsystemHandle) -> Result<()> { + /// tracing::info!("Starting subsystem ..."); + /// + /// // We cannot do a `tokio::select` with `on_shutdown_requested` + /// // here, because a shutdown would cancel the action without giving + /// // it the chance to react first. + /// while !subsys.is_shutdown_requested() { + /// uncancellable_action(&subsys).await; + /// } + /// + /// tracing::info!("Subsystem stopped."); + /// + /// Ok(()) + /// } + /// ``` + pub fn is_shutdown_requested(&self) -> bool { + self.inner.cancellation_token.is_cancelled() + } + + /// Triggers a shutdown of the entire subsystem tree. + /// + /// # Examples + /// + /// ``` + /// use miette::Result; + /// use tokio::time::{sleep, Duration}; + /// use tokio_graceful_shutdown::SubsystemHandle; + /// + /// async fn stop_subsystem(subsys: SubsystemHandle) -> Result<()> { + /// // This subsystem wait for one second and then stops the program. + /// sleep(Duration::from_millis(1000)).await; + /// + /// // Shut down the entire subsystem tree + /// subsys.request_shutdown(); + /// + /// Ok(()) + /// } + /// ``` + pub fn request_shutdown(&self) { + self.inner.toplevel_cancellation_token.cancel(); + } + + /// Triggers a shutdown of the current subsystem and all + /// of its children. + pub fn request_local_shutdown(&self) { + self.inner.cancellation_token.cancel(); + } + + pub(crate) fn get_cancellation_token(&self) -> &CancellationToken { + &self.inner.cancellation_token + } +} + +impl Drop for SubsystemHandle { + fn drop(&mut self) { + // SAFETY: This is how ManuallyDrop is meant to be used. + // `self.inner` won't ever be used again because `self` will be gone after this + // function is finished. + // This takes the `self.inner` object and makes it droppable again. + // + // This workaround is required to take ownership for the `self.drop_redirect` channel. + let inner = unsafe { ManuallyDrop::take(&mut self.inner) }; + + if let Some(redirect) = self.drop_redirect.take() { + let redirected_self = WeakSubsystemHandle { + joiner_token: inner.joiner_token, + _children: inner.children, + }; + + // ignore error; an error would indicate that there is no receiver. + // in that case, do nothing. + let _ = redirect.send(redirected_self); + } + } +} + +pub(crate) fn root_handle( + on_error: impl Fn(SubsystemError) + Sync + Send + 'static, +) -> SubsystemHandle { + let cancellation_token = CancellationToken::new(); + + SubsystemHandle { + inner: ManuallyDrop::new(Inner { + name: Arc::from(""), + cancellation_token: cancellation_token.clone(), + toplevel_cancellation_token: cancellation_token.clone(), + joiner_token: JoinerToken::new(move |e| { + on_error(e); + cancellation_token.cancel(); + None + }) + .0, + children: RemotelyDroppableItems::new(), + }), + drop_redirect: None, + } +} + +#[cfg(test)] +mod tests { + + use tokio::time::{sleep, timeout, Duration}; + + use super::*; + use crate::subsystem::SubsystemBuilder; + + #[tokio::test] + async fn recursive_cancellation() { + let root_handle = root_handle::(|_| {}); + + let (drop_sender, mut drop_receiver) = tokio::sync::mpsc::channel::<()>(1); + + root_handle.start(SubsystemBuilder::new("", |_| async move { + drop_sender.send(()).await.unwrap(); + std::future::pending::>().await + })); + + // Make sure we are executing the subsystem + let recv_result = timeout(Duration::from_millis(100), drop_receiver.recv()) + .await + .unwrap(); + assert!(recv_result.is_some()); + + drop(root_handle); + + // Make sure the subsystem got cancelled + let recv_result = timeout(Duration::from_millis(100), drop_receiver.recv()) + .await + .unwrap(); + assert!(recv_result.is_none()); + } + + #[tokio::test] + async fn recursive_cancellation_2() { + let root_handle = root_handle(|_| {}); + + let (drop_sender, mut drop_receiver) = tokio::sync::mpsc::channel::<()>(1); + + let subsys2 = |_| async move { + drop_sender.send(()).await.unwrap(); + std::future::pending::>().await + }; + + let subsys = |x: SubsystemHandle| async move { + x.start(SubsystemBuilder::new("", subsys2)); + + Result::<(), BoxedError>::Ok(()) + }; + + root_handle.start(SubsystemBuilder::new("", subsys)); + + // Make sure we are executing the subsystem + let recv_result = timeout(Duration::from_millis(100), drop_receiver.recv()) + .await + .unwrap(); + assert!(recv_result.is_some()); + + // Make sure the grandchild is still running + sleep(Duration::from_millis(100)).await; + assert!(matches!( + drop_receiver.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + + drop(root_handle); + + // Make sure the subsystem got cancelled + let recv_result = timeout(Duration::from_millis(100), drop_receiver.recv()) + .await + .unwrap(); + assert!(recv_result.is_none()); + } +} diff --git a/src/toplevel.rs b/src/toplevel.rs index b6a3a16..bcd52a0 100644 --- a/src/toplevel.rs +++ b/src/toplevel.rs @@ -1,34 +1,27 @@ -use std::future::Future; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; -use std::sync::Arc; -use std::time::Duration; +use std::{future::Future, sync::Arc, time::Duration}; +use atomic::Atomic; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; -use crate::errors::GracefulShutdownError; -use crate::exit_state::prettify_exit_states; -use crate::shutdown_token::create_shutdown_token; -use crate::signal_handling::wait_for_signal; -use crate::utils::get_subsystem_name; -use crate::utils::wait_forever; -use crate::utils::ShutdownGuard; -use crate::ErrTypeTraits; -use crate::{ShutdownToken, SubsystemHandle}; +use crate::{ + errors::{GracefulShutdownError, SubsystemError}, + signal_handling::wait_for_signal, + subsystem::{self, ErrorActions}, + BoxedError, ErrTypeTraits, ErrorAction, NestedSubsystem, SubsystemHandle, +}; -use super::subsystem::SubsystemData; - -/// Acts as the base for the subsystem tree and forms the entry point for +/// Acts as the root of the subsystem tree and forms the entry point for /// any interaction with this crate. /// -/// Every project that uses this crate has to create a Toplevel object somewhere. +/// Every project that uses this crate has to create a [`Toplevel`] object somewhere. /// /// # Examples /// /// ``` /// use miette::Result; /// use tokio::time::Duration; -/// use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; +/// use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; /// /// async fn my_subsystem(subsys: SubsystemHandle) -> Result<()> { /// subsys.request_shutdown(); @@ -37,124 +30,72 @@ use super::subsystem::SubsystemData; /// /// #[tokio::main] /// async fn main() -> Result<()> { -/// // Create toplevel -/// Toplevel::new() -/// .start("MySubsystem", my_subsystem) -/// .catch_signals() -/// .handle_shutdown_requests(Duration::from_millis(1000)) -/// .await -/// .map_err(Into::into) +/// Toplevel::new(|s| async move { +/// s.start(SubsystemBuilder::new("MySubsystem", my_subsystem)); +/// }) +/// .catch_signals() +/// .handle_shutdown_requests(Duration::from_millis(1000)) +/// .await +/// .map_err(Into::into) /// } /// ``` /// #[must_use = "This toplevel must be consumed by calling `handle_shutdown_requests` on it."] -pub struct Toplevel { - subsys_data: Arc>, - subsys_handle: SubsystemHandle, - shutdown_guard: Option>, +pub struct Toplevel { + root_handle: SubsystemHandle, + toplevel_subsys: NestedSubsystem, + errors: mpsc::UnboundedReceiver>, } impl Toplevel { /// Creates a new Toplevel object. /// /// The Toplevel object is the base for everything else in this crate. - #[allow(clippy::new_without_default)] - pub fn new() -> Self { - // On the top-level, the global and local shutdown token are identical - let global_shutdown_token = create_shutdown_token(); - let group_shutdown_token = global_shutdown_token.clone(); - let local_shutdown_token = group_shutdown_token.clone(); - let cancellation_token = CancellationToken::new(); - let shutdown_guard = Arc::new(ShutdownGuard::new(group_shutdown_token.clone())); - - let subsys_data = Arc::new(SubsystemData::new( - "", - global_shutdown_token, - group_shutdown_token, - local_shutdown_token, - cancellation_token, - Arc::downgrade(&shutdown_guard), - )); - let subsys_handle = SubsystemHandle::new(subsys_data.clone()); - Self { - subsys_data, - subsys_handle, - shutdown_guard: Some(shutdown_guard), - } - } - - /// Creates a new nested Toplevel object. - /// - /// This method is identical to `.new()`, except that the returned [Toplevel] object - /// will receive shutdown requests from the given [SubsystemHandle] object. - /// - /// Any errors caused by subsystems inside the new [Toplevel] object will cause - /// the [Toplevel] object to initiate a shutdown, but will not propagate up to the - /// [SubsystemHandle] object. - /// - /// # Arguments - /// - /// * `parent` - The subsystemhandle that the [Toplevel] object will receive shutdown - /// requests from - /// * `name` - The name of the nested toplevel object. Can be `""`. - pub fn nested(parent: &SubsystemHandle, name: &str) -> Self { - // Take shutdown tokesn from parent - let global_shutdown_token = parent.global_shutdown_token().clone(); - let group_shutdown_token = parent.local_shutdown_token().child_token(); - let local_shutdown_token = group_shutdown_token.clone(); - let cancellation_token = CancellationToken::new(); - let shutdown_guard = Arc::new(ShutdownGuard::new(group_shutdown_token.clone())); - - let name = get_subsystem_name(parent.name(), name); - - let subsys_data = Arc::new(SubsystemData::new( - &name, - global_shutdown_token, - group_shutdown_token, - local_shutdown_token, - cancellation_token, - Arc::downgrade(&shutdown_guard), - )); - let subsys_handle = SubsystemHandle::new(subsys_data.clone()); - Self { - subsys_data, - subsys_handle, - shutdown_guard: Some(shutdown_guard), - } - } - - /// Starts a new subsystem. - /// - /// Once called, the subsystem will be started immediately, similar to [`tokio::spawn`]. - /// - /// # Subsystem - /// - /// The functionality of the subsystem is represented by the 'subsystem' argument. - /// It has to be provided either as an asynchronous function or an asynchronous closure. - /// - /// It gets provided with a [`SubsystemHandle`] object which can be used to interact with this crate. - /// - /// ## Returns - /// - /// When the subsystem returns `Ok(())` it is assumed that the subsystem was stopped intentionally and no further - /// actions are performed. - /// - /// When the subsystem returns an `Err`, it is assumed that the subsystem failed and a program shutdown gets initiated. /// /// # Arguments /// - /// * `name` - The name of the subsystem - /// * `subsystem` - The subsystem to be started - /// - pub fn start(self, name: &str, subsystem: Subsys) -> Self + /// * `subsystem` - The subsystem that should be spawned as the root node. + /// Usually the job of this subsystem is to spawn further subsystems. + #[allow(clippy::new_without_default)] + pub fn new(subsystem: Subsys) -> Self where Subsys: 'static + FnOnce(SubsystemHandle) -> Fut + Send, - Fut: 'static + Future> + Send, - Err: Into, + Fut: 'static + Future + Send, { - self.subsys_handle.start(name, subsystem); + let (error_sender, errors) = mpsc::unbounded_channel(); + + let root_handle = subsystem::root_handle(move |e| { + match &e { + SubsystemError::Panicked(name) => { + tracing::error!("Uncaught panic from subsytem '{name}'.") + } + SubsystemError::Failed(name, e) => { + tracing::error!("Uncaught error from subsystem '{name}': {e}",) + } + }; + + if let Err(mpsc::error::SendError(e)) = error_sender.send(e) { + tracing::warn!("An error got dropped: {e:?}"); + }; + }); - self + let toplevel_subsys = root_handle.start_with_abs_name( + Arc::from(""), + move |s| async move { + subsystem(s).await; + Result::<(), ErrType>::Ok(()) + }, + ErrorActions { + on_failure: Atomic::new(ErrorAction::Forward), + on_panic: Atomic::new(ErrorAction::Forward), + }, + ); + + Self { + root_handle, + toplevel_subsys, + errors, + } } /// Registers signal handlers to initiate a program shutdown when certain operating system @@ -175,47 +116,16 @@ impl Toplevel { /// Especially the caveats from [tokio::signal::unix::Signal] are important for Unix targets. /// pub fn catch_signals(self) -> Self { - let shutdown_token = self.subsys_handle.group_shutdown_token().clone(); + let shutdown_token = self.root_handle.get_cancellation_token().clone(); tokio::spawn(async move { wait_for_signal().await; - shutdown_token.shutdown(); + shutdown_token.cancel(); }); self } - /// Wait for all subsystems to finish. - /// Then return and print all of their exit codes. - async fn attempt_clean_shutdown(&self) -> Result<(), GracefulShutdownError> { - let exit_states = self.subsys_data.perform_shutdown().await; - - // Prettify exit states - let formatted_exit_states = prettify_exit_states(&exit_states); - - // Collect failed subsystems - let failed_subsystems = exit_states - .into_iter() - .filter_map(|exit_state| exit_state.raw_result.err()) - .collect::>(); - - // Print subsystem exit states - if failed_subsystems.is_empty() { - log::debug!("Shutdown successful. Subsystem states:"); - } else { - log::debug!("Some subsystems failed. Subsystem states:"); - }; - for formatted_exit_state in formatted_exit_states { - log::debug!(" {}", formatted_exit_state); - } - - if failed_subsystems.is_empty() { - Ok(()) - } else { - Err(GracefulShutdownError::SubsystemsFailed(failed_subsystems)) - } - } - /// Performs a clean program shutdown, once a shutdown is requested or all subsystems have /// finished. /// @@ -240,48 +150,61 @@ impl Toplevel { mut self, shutdown_timeout: Duration, ) -> Result<(), GracefulShutdownError> { - // Remove the shutdown guard we hold ourselves, to enable auto-shutdown triggering - // when all subsystems are finished - self.shutdown_guard.take(); - - self.subsys_handle.on_shutdown_requested().await; - - let timeout_occurred = AtomicBool::new(false); - - let cancel_on_timeout = async { - // Wait for the timeout to happen - tokio::time::sleep(shutdown_timeout).await; - log::error!("Shutdown timed out. Attempting to cleanup stale subsystems ..."); - timeout_occurred.store(true, Ordering::SeqCst); - self.subsys_data.cancel_all_subsystems(); - // Await forever, because we don't want to cancel the attempt_clean_shutdown. - // Resolving this arm of the tokio::select would cancel the other side. - wait_forever().await; - }; - - let result = tokio::select! { - _ = cancel_on_timeout => unreachable!(), - result = self.attempt_clean_shutdown() => result + let collect_errors = move || { + let mut errors = vec![]; + self.errors.close(); + while let Ok(e) = self.errors.try_recv() { + errors.push(e); + } + drop(self.errors); + errors.into_boxed_slice() }; - // Overwrite return value with "ShutdownTimeout" if a timeout occurred - if timeout_occurred.load(Ordering::SeqCst) { - Err(GracefulShutdownError::ShutdownTimeout( - result.err().map_or(vec![], |e| e.into_subsystem_errors()), - )) - } else { - result + tokio::select!( + _ = self.toplevel_subsys.join() => { + tracing::info!("All subsystems finished."); + + // Not really necessary, but for good measure. + self.root_handle.request_shutdown(); + + let errors = collect_errors(); + let result = if errors.is_empty() { + Ok(()) + } else { + Err(GracefulShutdownError::SubsystemsFailed(errors)) + }; + return result; + }, + _ = self.root_handle.on_shutdown_requested() => { + tracing::info!("Shutting down ..."); + } + ); + + match tokio::time::timeout(shutdown_timeout, self.toplevel_subsys.join()).await { + Ok(Ok(())) => { + let errors = collect_errors(); + if errors.is_empty() { + tracing::info!("Shutdown finished."); + Ok(()) + } else { + tracing::warn!("Shutdown finished with errors."); + Err(GracefulShutdownError::SubsystemsFailed(errors)) + } + } + Ok(Err(_)) => { + // This can't happen because the toplevel subsys doesn't catch any errors; it only forwards them. + unreachable!(); + } + Err(_) => { + tracing::error!("Shutdown timed out!"); + Err(GracefulShutdownError::ShutdownTimeout(collect_errors())) + } } } #[doc(hidden)] - pub fn get_shutdown_token(&self) -> &ShutdownToken { - self.subsys_handle.local_shutdown_token() - } -} - -impl Drop for Toplevel { - fn drop(&mut self) { - self.subsys_data.cancel_all_subsystems(); + // Only for unit tests; not intended for public use + pub fn _get_shutdown_token(&self) -> &CancellationToken { + self.root_handle.get_cancellation_token() } } diff --git a/src/utils/joiner_token.rs b/src/utils/joiner_token.rs new file mode 100644 index 0000000..e792d8e --- /dev/null +++ b/src/utils/joiner_token.rs @@ -0,0 +1,412 @@ +use std::{fmt::Debug, sync::Arc}; + +use tokio::sync::watch; + +use crate::{errors::SubsystemError, ErrTypeTraits}; + +struct Inner { + counter: watch::Sender<(bool, u32)>, + parent: Option>>, + on_error: Box) -> Option> + Sync + Send>, +} + +/// A token that keeps reference of its existance and its children. +pub(crate) struct JoinerToken { + inner: Arc>, +} + +/// A reference version that does not keep the content alive; purely for +/// joining the subtree. +pub(crate) struct JoinerTokenRef { + counter: watch::Receiver<(bool, u32)>, +} + +impl Debug for JoinerToken { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "JoinerToken(children = {})", + self.inner.counter.borrow().1 + ) + } +} + +impl Debug for JoinerTokenRef { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let counter = self.counter.borrow(); + write!( + f, + "JoinerTokenRef(alive = {}, children = {})", + counter.0, counter.1 + ) + } +} + +impl JoinerToken { + /// Creates a new joiner token. + /// + /// The `on_error` callback will receive errors/panics and has to decide + /// how to handle them. It can also not handle them and instead pass them on. + /// If it returns `Some`, the error will get passed on to its parent. + pub(crate) fn new( + on_error: impl Fn(SubsystemError) -> Option> + + Sync + + Send + + 'static, + ) -> (Self, JoinerTokenRef) { + let inner = Arc::new(Inner { + counter: watch::channel((true, 0)).0, + parent: None, + on_error: Box::new(on_error), + }); + + let weak_ref = JoinerTokenRef { + counter: inner.counter.subscribe(), + }; + + (Self { inner }, weak_ref) + } + + // Requires `mut` access to prevent children from being spawned + // while waiting + pub(crate) async fn join_children(&mut self) { + let mut subscriber = self.inner.counter.subscribe(); + + // Ignore errors; if the channel got closed, that definitely means + // no more children exist. + let _ = subscriber + .wait_for(|(_alive, children)| *children == 0) + .await; + } + + pub(crate) fn child_token( + &self, + on_error: impl Fn(SubsystemError) -> Option> + + Sync + + Send + + 'static, + ) -> (Self, JoinerTokenRef) { + let mut maybe_parent = Some(&self.inner); + while let Some(parent) = maybe_parent { + parent + .counter + .send_modify(|(_alive, children)| *children += 1); + maybe_parent = parent.parent.as_ref(); + } + + let inner = Arc::new(Inner { + counter: watch::channel((true, 0)).0, + parent: Some(Arc::clone(&self.inner)), + on_error: Box::new(on_error), + }); + + let weak_ref = JoinerTokenRef { + counter: inner.counter.subscribe(), + }; + + (Self { inner }, weak_ref) + } + + #[cfg(test)] + pub(crate) fn count(&self) -> u32 { + self.inner.counter.borrow().1 + } + + pub(crate) fn raise_failure(&self, stop_reason: SubsystemError) { + let mut maybe_stop_reason = Some(stop_reason); + + let mut maybe_parent = Some(&self.inner); + while let Some(parent) = maybe_parent { + if let Some(stop_reason) = maybe_stop_reason { + maybe_stop_reason = (parent.on_error)(stop_reason); + } else { + break; + } + + maybe_parent = parent.parent.as_ref(); + } + + if let Some(stop_reason) = maybe_stop_reason { + tracing::warn!("Unhandled stop reason: {:?}", stop_reason); + } + } + + pub(crate) fn downgrade(self) -> JoinerTokenRef { + JoinerTokenRef { + counter: self.inner.counter.subscribe(), + } + } +} + +impl JoinerTokenRef { + pub(crate) async fn join(&self) { + // Ignore errors; if the channel got closed, that definitely means + // the token and all its children got dropped. + let _ = self + .counter + .clone() + .wait_for(|&(alive, children)| !alive && children == 0) + .await; + } + + #[cfg(test)] + pub(crate) fn count(&self) -> u32 { + self.counter.borrow().1 + } + + #[cfg(test)] + pub(crate) fn alive(&self) -> bool { + self.counter.borrow().0 + } +} + +impl Drop for JoinerToken { + fn drop(&mut self) { + self.inner + .counter + .send_modify(|(alive, _children)| *alive = false); + + let mut maybe_parent = self.inner.parent.as_ref(); + while let Some(parent) = maybe_parent { + parent + .counter + .send_modify(|(_alive, children)| *children -= 1); + maybe_parent = parent.parent.as_ref(); + } + } +} + +#[cfg(test)] +mod tests { + use tokio::time::{sleep, timeout, Duration}; + use tracing_test::traced_test; + + use crate::BoxedError; + + use super::*; + + #[test] + #[traced_test] + fn counters() { + let (root, _) = JoinerToken::::new(|_| None); + assert_eq!(0, root.count()); + + let (child1, _) = root.child_token(|_| None); + assert_eq!(1, root.count()); + assert_eq!(0, child1.count()); + + let (child2, _) = child1.child_token(|_| None); + assert_eq!(2, root.count()); + assert_eq!(1, child1.count()); + assert_eq!(0, child2.count()); + + let (child3, _) = child1.child_token(|_| None); + assert_eq!(3, root.count()); + assert_eq!(2, child1.count()); + assert_eq!(0, child2.count()); + assert_eq!(0, child3.count()); + + drop(child1); + assert_eq!(2, root.count()); + assert_eq!(0, child2.count()); + assert_eq!(0, child3.count()); + + drop(child2); + assert_eq!(1, root.count()); + assert_eq!(0, child3.count()); + + drop(child3); + assert_eq!(0, root.count()); + } + + #[test] + #[traced_test] + fn counters_weak() { + let (root, weak_root) = JoinerToken::::new(|_| None); + assert_eq!(0, weak_root.count()); + assert!(weak_root.alive()); + + let (child1, weak_child1) = root.child_token(|_| None); + assert_eq!(1, weak_root.count()); + assert!(weak_root.alive()); + assert_eq!(0, weak_child1.count()); + assert!(weak_child1.alive()); + + let (child2, weak_child2) = child1.child_token(|_| None); + assert_eq!(2, weak_root.count()); + assert!(weak_root.alive()); + assert_eq!(1, weak_child1.count()); + assert!(weak_child1.alive()); + assert_eq!(0, weak_child2.count()); + assert!(weak_child2.alive()); + + let (child3, weak_child3) = child1.child_token(|_| None); + assert_eq!(3, weak_root.count()); + assert!(weak_root.alive()); + assert_eq!(2, weak_child1.count()); + assert!(weak_child1.alive()); + assert_eq!(0, weak_child2.count()); + assert!(weak_child2.alive()); + assert_eq!(0, weak_child3.count()); + assert!(weak_child3.alive()); + + drop(child1); + assert_eq!(2, weak_root.count()); + assert!(weak_root.alive()); + assert_eq!(2, weak_child1.count()); + assert!(!weak_child1.alive()); + assert_eq!(0, weak_child2.count()); + assert!(weak_child2.alive()); + assert_eq!(0, weak_child3.count()); + assert!(weak_child3.alive()); + + drop(child2); + assert_eq!(1, weak_root.count()); + assert!(weak_root.alive()); + assert_eq!(1, weak_child1.count()); + assert!(!weak_child1.alive()); + assert_eq!(0, weak_child2.count()); + assert!(!weak_child2.alive()); + assert_eq!(0, weak_child3.count()); + assert!(weak_child3.alive()); + + drop(child3); + assert_eq!(0, weak_root.count()); + assert!(weak_root.alive()); + assert_eq!(0, weak_child1.count()); + assert!(!weak_child1.alive()); + assert_eq!(0, weak_child2.count()); + assert!(!weak_child2.alive()); + assert_eq!(0, weak_child3.count()); + assert!(!weak_child3.alive()); + + drop(root); + assert_eq!(0, weak_root.count()); + assert!(!weak_root.alive()); + assert_eq!(0, weak_child1.count()); + assert!(!weak_child1.alive()); + assert_eq!(0, weak_child2.count()); + assert!(!weak_child2.alive()); + assert_eq!(0, weak_child3.count()); + assert!(!weak_child3.alive()); + } + + #[tokio::test] + #[traced_test] + async fn join() { + let (superroot, _) = JoinerToken::::new(|_| None); + + let (mut root, _) = superroot.child_token(|_| None); + + let (child1, _) = root.child_token(|_| None); + let (child2, _) = child1.child_token(|_| None); + let (child3, _) = child1.child_token(|_| None); + + let (set_finished, mut finished) = tokio::sync::oneshot::channel(); + tokio::join!( + async { + timeout(Duration::from_millis(500), root.join_children()) + .await + .unwrap(); + set_finished.send(root.count()).unwrap(); + }, + async { + sleep(Duration::from_millis(50)).await; + assert!(finished.try_recv().is_err()); + + drop(child1); + sleep(Duration::from_millis(50)).await; + assert!(finished.try_recv().is_err()); + + drop(child2); + sleep(Duration::from_millis(50)).await; + assert!(finished.try_recv().is_err()); + + drop(child3); + sleep(Duration::from_millis(50)).await; + let count = timeout(Duration::from_millis(50), finished) + .await + .unwrap() + .unwrap(); + assert_eq!(count, 0); + } + ); + } + + #[tokio::test] + #[traced_test] + async fn join_through_ref() { + let (root, joiner) = JoinerToken::::new(|_| None); + + let (child1, _) = root.child_token(|_| None); + let (child2, _) = child1.child_token(|_| None); + + let (set_finished, mut finished) = tokio::sync::oneshot::channel(); + tokio::join!( + async { + timeout(Duration::from_millis(500), joiner.join()) + .await + .unwrap(); + set_finished.send(()).unwrap(); + }, + async { + sleep(Duration::from_millis(50)).await; + assert!(finished.try_recv().is_err()); + + drop(child1); + sleep(Duration::from_millis(50)).await; + assert!(finished.try_recv().is_err()); + + drop(root); + sleep(Duration::from_millis(50)).await; + assert!(finished.try_recv().is_err()); + + drop(child2); + sleep(Duration::from_millis(50)).await; + timeout(Duration::from_millis(50), finished) + .await + .unwrap() + .unwrap(); + } + ); + } + + #[test] + fn debug_print() { + let (root, _) = JoinerToken::::new(|_| None); + assert_eq!(format!("{:?}", root), "JoinerToken(children = 0)"); + + let (child1, _) = root.child_token(|_| None); + assert_eq!(format!("{:?}", root), "JoinerToken(children = 1)"); + + let (_child2, _) = child1.child_token(|_| None); + assert_eq!(format!("{:?}", root), "JoinerToken(children = 2)"); + } + + #[test] + fn debug_print_ref() { + let (root, root_ref) = JoinerToken::::new(|_| None); + assert_eq!( + format!("{:?}", root_ref), + "JoinerTokenRef(alive = true, children = 0)" + ); + + let (child1, _) = root.child_token(|_| None); + assert_eq!( + format!("{:?}", root_ref), + "JoinerTokenRef(alive = true, children = 1)" + ); + + drop(root); + assert_eq!( + format!("{:?}", root_ref), + "JoinerTokenRef(alive = false, children = 1)" + ); + + drop(child1); + assert_eq!( + format!("{:?}", root_ref), + "JoinerTokenRef(alive = false, children = 0)" + ); + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 718550a..c96851c 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,14 +1,5 @@ -mod wait_forever; -pub use wait_forever::wait_forever; +mod joiner_token; +pub(crate) use joiner_token::JoinerToken; +pub(crate) use joiner_token::JoinerTokenRef; -mod shutdown_guard; -pub use shutdown_guard::ShutdownGuard; - -pub fn get_subsystem_name(parent_name: &str, name: &str) -> String { - match (parent_name, name) { - ("", "") => "".to_string(), - (l, "") => l.to_string(), - ("", r) => r.to_string(), - (l, r) => l.to_string() + "/" + r, - } -} +pub(crate) mod remote_drop_collection; diff --git a/src/utils/remote_drop_collection.rs b/src/utils/remote_drop_collection.rs new file mode 100644 index 0000000..56cfc5a --- /dev/null +++ b/src/utils/remote_drop_collection.rs @@ -0,0 +1,157 @@ +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, Weak, +}; + +struct RemotelyDroppableItem { + _item: T, + offset: Arc, +} + +/// A vector that owns a bunch of objects. +/// Every object is connected to a guard token. +/// Once the token is dropped, the object gets dropped as well. +/// +/// Note that the token does not keep the object alive, it is only responsible +/// for triggering a drop. +/// +/// The important part here is that the token is sendable to other context/threads, +/// so it's basically a 'remote drop guard' concept. +pub(crate) struct RemotelyDroppableItems { + items: Arc>>>, +} + +impl RemotelyDroppableItems { + pub(crate) fn new() -> Self { + Self { + items: Default::default(), + } + } + + pub(crate) fn insert(&self, item: T) -> RemoteDrop { + let mut items = self.items.lock().unwrap(); + + let offset = Arc::new(AtomicUsize::new(items.len())); + let weak_offset = Arc::downgrade(&offset); + + items.push(RemotelyDroppableItem { + _item: item, + offset, + }); + + RemoteDrop { + data: Arc::downgrade(&self.items), + offset: weak_offset, + } + } +} + +/// Drops its referenced item when dropped +pub(crate) struct RemoteDrop { + // Both weak. + // If data is gone, then our item collection dropped. + data: Weak>>>, + // If offset is gone, then the item itself got removed + // while the dropguard still exists. + offset: Weak, +} + +impl Drop for RemoteDrop { + fn drop(&mut self) { + if let Some(data) = self.data.upgrade() { + // Important: lock first, then read the offset. + let mut data = data.lock().unwrap(); + + if let Some(offset) = self.offset.upgrade() { + let offset = offset.load(Ordering::Acquire); + + if let Some(last_item) = data.pop() { + if offset != data.len() { + // There must have been at least two items, and we are not at the end. + // So swap first before dropping. + + last_item.offset.store(offset, Ordering::Release); + data[offset] = last_item; + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + + use super::*; + use crate::{utils::JoinerToken, BoxedError}; + + #[test] + fn insert_and_drop() { + let items = RemotelyDroppableItems::new(); + + let (count1, _) = JoinerToken::::new(|_| None); + let (count2, _) = JoinerToken::::new(|_| None); + + assert_eq!(0, count1.count()); + assert_eq!(0, count2.count()); + + let _token1 = items.insert(count1.child_token(|_| None)); + assert_eq!(1, count1.count()); + assert_eq!(0, count2.count()); + + let _token2 = items.insert(count2.child_token(|_| None)); + assert_eq!(1, count1.count()); + assert_eq!(1, count2.count()); + + drop(items); + assert_eq!(0, count1.count()); + assert_eq!(0, count2.count()); + } + + #[test] + fn drop_token() { + let items = RemotelyDroppableItems::new(); + + let (count1, _) = JoinerToken::::new(|_| None); + let (count2, _) = JoinerToken::::new(|_| None); + let (count3, _) = JoinerToken::::new(|_| None); + let (count4, _) = JoinerToken::::new(|_| None); + + let token1 = items.insert(count1.child_token(|_| None)); + let token2 = items.insert(count2.child_token(|_| None)); + let token3 = items.insert(count3.child_token(|_| None)); + let token4 = items.insert(count4.child_token(|_| None)); + assert_eq!(1, count1.count()); + assert_eq!(1, count2.count()); + assert_eq!(1, count3.count()); + assert_eq!(1, count4.count()); + + // Last item + drop(token4); + assert_eq!(1, count1.count()); + assert_eq!(1, count2.count()); + assert_eq!(1, count3.count()); + assert_eq!(0, count4.count()); + + // Middle item + drop(token2); + assert_eq!(1, count1.count()); + assert_eq!(0, count2.count()); + assert_eq!(1, count3.count()); + assert_eq!(0, count4.count()); + + // First item + drop(token1); + assert_eq!(0, count1.count()); + assert_eq!(0, count2.count()); + assert_eq!(1, count3.count()); + assert_eq!(0, count4.count()); + + // Only item + drop(token3); + assert_eq!(0, count1.count()); + assert_eq!(0, count2.count()); + assert_eq!(0, count3.count()); + assert_eq!(0, count4.count()); + } +} diff --git a/src/utils/shutdown_guard.rs b/src/utils/shutdown_guard.rs deleted file mode 100644 index a7adb4d..0000000 --- a/src/utils/shutdown_guard.rs +++ /dev/null @@ -1,16 +0,0 @@ -use crate::ShutdownToken; - -/// Triggers the ShutdownToken when dropped -pub struct ShutdownGuard(ShutdownToken); - -impl ShutdownGuard { - pub fn new(token: ShutdownToken) -> Self { - Self(token) - } -} - -impl Drop for ShutdownGuard { - fn drop(&mut self) { - self.0.shutdown() - } -} diff --git a/src/utils/wait_forever.rs b/src/utils/wait_forever.rs deleted file mode 100644 index 8eecfb9..0000000 --- a/src/utils/wait_forever.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub async fn wait_forever() -> ! { - loop { - std::future::pending::<()>().await; - } -} diff --git a/tests/cancel_on_shutdown.rs b/tests/cancel_on_shutdown.rs index 9931dde..43dbbb7 100644 --- a/tests/cancel_on_shutdown.rs +++ b/tests/cancel_on_shutdown.rs @@ -1,8 +1,8 @@ use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{errors::CancelledByShutdown, FutureExt, SubsystemHandle, Toplevel}; - -pub mod common; -use common::setup; +use tokio_graceful_shutdown::{ + errors::CancelledByShutdown, FutureExt, SubsystemBuilder, SubsystemHandle, Toplevel, +}; +use tracing_test::traced_test; use std::error::Error; @@ -11,9 +11,8 @@ type BoxedError = Box; type BoxedResult = Result<(), BoxedError>; #[tokio::test] +#[traced_test] async fn cancel_on_shutdown_propagates_result() { - setup(); - let subsystem1 = |subsys: SubsystemHandle| async move { let compute_value = async { sleep(Duration::from_millis(10)).await; @@ -40,40 +39,39 @@ async fn cancel_on_shutdown_propagates_result() { BoxedResult::Ok(()) }; - let result = Toplevel::::new() - .start("subsys1", subsystem1) - .start("subsys2", subsystem2) - .handle_shutdown_requests(Duration::from_millis(200)) - .await; + let result = Toplevel::::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys1", subsystem1)); + s.start(SubsystemBuilder::new("subsys2", subsystem2)); + }) + .handle_shutdown_requests(Duration::from_millis(200)) + .await; assert!(result.is_ok()); } #[tokio::test] +#[traced_test] async fn cancel_on_shutdown_cancels_on_shutdown() { - setup(); - let subsystem = |subsys: SubsystemHandle| async move { - async fn compute_value(subsys: SubsystemHandle) -> i32 { + async fn compute_value(subsys: &SubsystemHandle) -> i32 { sleep(Duration::from_millis(100)).await; subsys.request_shutdown(); sleep(Duration::from_millis(100)).await; 42 } - let value = compute_value(subsys.clone()) - .cancel_on_shutdown(&subsys) - .await; + let value = compute_value(&subsys).cancel_on_shutdown(&subsys).await; assert!(matches!(value, Err(CancelledByShutdown))); BoxedResult::Ok(()) }; - let result = Toplevel::::new() - .start("subsys", subsystem) - .handle_shutdown_requests(Duration::from_millis(200)) - .await; + let result = Toplevel::::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }) + .handle_shutdown_requests(Duration::from_millis(200)) + .await; assert!(result.is_ok()); } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 45f8ddf..90b97a9 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -1,13 +1,3 @@ -use std::sync::Once; +mod event; -pub mod event; - -static INIT: Once = Once::new(); - -/// Setup function that is only run once, even if called multiple times. -pub fn setup() { - INIT.call_once(|| { - // Init logging - env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("off")).init(); - }); -} +pub use event::Event; diff --git a/tests/integration_test.rs b/tests/integration_test.rs index c177f5e..463d463 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -1,13 +1,13 @@ use anyhow::anyhow; use tokio::time::{sleep, timeout, Duration}; use tokio_graceful_shutdown::{ - errors::{GracefulShutdownError, PartialShutdownError, SubsystemError}, - IntoSubsystem, SubsystemHandle, Toplevel, + errors::{GracefulShutdownError, SubsystemError, SubsystemJoinError}, + ErrorAction, IntoSubsystem, SubsystemBuilder, SubsystemHandle, Toplevel, }; +use tracing_test::traced_test; pub mod common; -use common::event::Event; -use common::setup; +use common::Event; use std::error::Error; @@ -16,36 +16,30 @@ type BoxedError = Box; type BoxedResult = Result<(), BoxedError>; #[tokio::test] +#[traced_test] async fn normal_shutdown() { - setup(); - let subsystem = |subsys: SubsystemHandle| async move { subsys.on_shutdown_requested().await; sleep(Duration::from_millis(200)).await; BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); - tokio::join!( - async { - sleep(Duration::from_millis(100)).await; - shutdown_token.shutdown(); - }, - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(400)) - .await; - assert!(result.is_ok()); - }, - ); + sleep(Duration::from_millis(100)).await; + s.request_shutdown(); + }); + + let result = toplevel + .handle_shutdown_requests(Duration::from_millis(400)) + .await; + assert!(result.is_ok()); } #[tokio::test] +#[traced_test] async fn use_subsystem_struct() { - setup(); - struct MySubsystem; #[async_trait::async_trait] @@ -57,58 +51,51 @@ async fn use_subsystem_struct() { } } - let toplevel = Toplevel::new().start("subsys", MySubsystem {}.into_subsystem()); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new( + "subsys", + MySubsystem {}.into_subsystem(), + )); - tokio::join!( - async { - sleep(Duration::from_millis(100)).await; - shutdown_token.shutdown(); - }, - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(400)) - .await; - assert!(result.is_ok()); - }, - ); + sleep(Duration::from_millis(100)).await; + s.request_shutdown(); + }); + + let result = toplevel + .handle_shutdown_requests(Duration::from_millis(400)) + .await; + assert!(result.is_ok()); } #[tokio::test] +#[traced_test] async fn shutdown_timeout_causes_error() { - setup(); - let subsystem = |subsys: SubsystemHandle| async move { subsys.on_shutdown_requested().await; sleep(Duration::from_millis(400)).await; BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); - tokio::join!( - async { - sleep(Duration::from_millis(100)).await; - shutdown_token.shutdown(); - }, - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(200)) - .await; - assert!(result.is_err()); - assert!(matches!( - result, - Err(GracefulShutdownError::ShutdownTimeout(_)) - )) - }, - ); + sleep(Duration::from_millis(100)).await; + s.request_shutdown(); + }); + + let result = toplevel + .handle_shutdown_requests(Duration::from_millis(200)) + .await; + assert!(result.is_err()); + assert!(matches!( + result, + Err(GracefulShutdownError::ShutdownTimeout(_)) + )); } #[tokio::test] +#[traced_test] async fn subsystem_finishes_with_success() { - setup(); - let subsystem = |_| async { BoxedResult::Ok(()) }; let subsystem2 = |subsys: SubsystemHandle| async move { subsys.on_shutdown_requested().await; @@ -117,10 +104,11 @@ async fn subsystem_finishes_with_success() { let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::::new() - .start("subsys", subsystem) - .start("subsys2", subsystem2); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + s.start(SubsystemBuilder::new("subsys2", subsystem2)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -135,7 +123,7 @@ async fn subsystem_finishes_with_success() { sleep(Duration::from_millis(200)).await; // Assert Ok(()) doesn't cause a shutdown assert!(!toplevel_finished.get()); - shutdown_token.shutdown(); + shutdown_token.cancel(); sleep(Duration::from_millis(200)).await; // Assert toplevel sucessfully gets stopped, nothing hangs assert!(toplevel_finished.get()); @@ -144,9 +132,8 @@ async fn subsystem_finishes_with_success() { } #[tokio::test] +#[traced_test] async fn subsystem_finishes_with_error() { - setup(); - let subsystem = |_| async { Err(anyhow!("Error!")) }; let subsystem2 = |subsys: SubsystemHandle| async move { subsys.on_shutdown_requested().await; @@ -155,10 +142,11 @@ async fn subsystem_finishes_with_error() { let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::::new() - .start("subsys", subsystem) - .start("subsys2", subsystem2); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + s.start(SubsystemBuilder::new("subsys2", subsystem2)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -173,15 +161,14 @@ async fn subsystem_finishes_with_error() { sleep(Duration::from_millis(200)).await; // Assert Err(()) causes a shutdown assert!(toplevel_finished.get()); - assert!(shutdown_token.is_shutting_down()); + assert!(shutdown_token.is_cancelled()); }, ); } #[tokio::test] +#[traced_test] async fn subsystem_receives_shutdown() { - setup(); - let (subsys_finished, set_subsys_finished) = Event::create(); let subsys = |subsys: SubsystemHandle| async move { @@ -190,14 +177,16 @@ async fn subsystem_receives_shutdown() { BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().start("subsys", subsys); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsys)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); let result = tokio::spawn(toplevel.handle_shutdown_requests(Duration::from_millis(100))); sleep(Duration::from_millis(100)).await; assert!(!subsys_finished.get()); - shutdown_token.shutdown(); + shutdown_token.cancel(); timeout(Duration::from_millis(100), subsys_finished.wait()) .await .unwrap(); @@ -211,9 +200,8 @@ async fn subsystem_receives_shutdown() { } #[tokio::test] +#[traced_test] async fn nested_subsystem_receives_shutdown() { - setup(); - let (subsys_finished, set_subsys_finished) = Event::create(); let nested_subsystem = |subsys: SubsystemHandle| async move { @@ -223,19 +211,21 @@ async fn nested_subsystem_receives_shutdown() { }; let subsystem = |subsys: SubsystemHandle| async move { - subsys.start("nested", nested_subsystem); + subsys.start(SubsystemBuilder::new("nested", nested_subsystem)); subsys.on_shutdown_requested().await; BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); let result = tokio::spawn(toplevel.handle_shutdown_requests(Duration::from_millis(100))); sleep(Duration::from_millis(100)).await; assert!(!subsys_finished.get()); - shutdown_token.shutdown(); + shutdown_token.cancel(); timeout(Duration::from_millis(100), subsys_finished.wait()) .await .unwrap(); @@ -249,21 +239,22 @@ async fn nested_subsystem_receives_shutdown() { } #[tokio::test] +#[traced_test] async fn nested_subsystem_error_propagates() { - setup(); - let nested_subsystem = |_subsys: SubsystemHandle| async move { Err(anyhow!("Error!")) }; let subsystem = move |subsys: SubsystemHandle| async move { - subsys.start("nested", nested_subsystem); + subsys.start(SubsystemBuilder::new("nested", nested_subsystem)); subsys.on_shutdown_requested().await; BoxedResult::Ok(()) }; let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -278,29 +269,30 @@ async fn nested_subsystem_error_propagates() { sleep(Duration::from_millis(200)).await; // Assert Err(()) causes a shutdown assert!(toplevel_finished.get()); - assert!(shutdown_token.is_shutting_down()); + assert!(shutdown_token.is_cancelled()); }, ); } #[tokio::test] +#[traced_test] async fn panic_gets_handled_correctly() { - setup(); - let nested_subsystem = |_subsys: SubsystemHandle| async move { panic!("Error!"); }; let subsystem = move |subsys: SubsystemHandle| async move { - subsys.start::("nested", nested_subsystem); + subsys.start::(SubsystemBuilder::new("nested", nested_subsystem)); subsys.on_shutdown_requested().await; BoxedResult::Ok(()) }; let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -315,15 +307,14 @@ async fn panic_gets_handled_correctly() { sleep(Duration::from_millis(200)).await; // Assert panic causes a shutdown assert!(toplevel_finished.get()); - assert!(shutdown_token.is_shutting_down()); + assert!(shutdown_token.is_cancelled()); }, ); } #[tokio::test] +#[traced_test] async fn subsystem_can_request_shutdown() { - setup(); - let (subsystem_should_stop, stop_subsystem) = Event::create(); let (subsys_finished, set_subsys_finished) = Event::create(); @@ -338,8 +329,10 @@ async fn subsystem_can_request_shutdown() { let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -355,7 +348,7 @@ async fn subsystem_can_request_shutdown() { sleep(Duration::from_millis(200)).await; assert!(!toplevel_finished.get()); assert!(!subsys_finished.get()); - assert!(!shutdown_token.is_shutting_down()); + assert!(!shutdown_token.is_cancelled()); stop_subsystem(); sleep(Duration::from_millis(200)).await; @@ -363,15 +356,14 @@ async fn subsystem_can_request_shutdown() { // Assert request_shutdown() causes a shutdown assert!(toplevel_finished.get()); assert!(subsys_finished.get()); - assert!(shutdown_token.is_shutting_down()); + assert!(shutdown_token.is_cancelled()); }, ); } #[tokio::test] +#[traced_test] async fn shutdown_timeout_causes_cancellation() { - setup(); - let (subsys_finished, set_subsys_finished) = Event::create(); let subsystem = |subsys: SubsystemHandle| async move { @@ -383,8 +375,10 @@ async fn shutdown_timeout_causes_cancellation() { let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -400,9 +394,9 @@ async fn shutdown_timeout_causes_cancellation() { sleep(Duration::from_millis(200)).await; assert!(!toplevel_finished.get()); assert!(!subsys_finished.get()); - assert!(!shutdown_token.is_shutting_down()); + assert!(!shutdown_token.is_cancelled()); - shutdown_token.shutdown(); + shutdown_token.cancel(); timeout(Duration::from_millis(300), toplevel_finished.wait()) .await .unwrap(); @@ -419,14 +413,14 @@ async fn shutdown_timeout_causes_cancellation() { } #[tokio::test] +#[traced_test] async fn spawning_task_during_shutdown_causes_task_to_be_cancelled() { - setup(); - let (subsys_finished, set_subsys_finished) = Event::create(); let (nested_finished, set_nested_finished) = Event::create(); - let nested = |_: SubsystemHandle| async move { + let nested = |subsys: SubsystemHandle| async move { sleep(Duration::from_millis(100)).await; + subsys.on_shutdown_requested().await; set_nested_finished(); BoxedResult::Ok(()) }; @@ -434,15 +428,17 @@ async fn spawning_task_during_shutdown_causes_task_to_be_cancelled() { let subsystem = move |subsys: SubsystemHandle| async move { subsys.on_shutdown_requested().await; sleep(Duration::from_millis(100)).await; - subsys.start("Nested", nested); + subsys.start(SubsystemBuilder::new("Nested", nested)); set_subsys_finished(); BoxedResult::Ok(()) }; let (toplevel_finished, set_toplevel_finished) = Event::create(); - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { @@ -458,30 +454,22 @@ async fn spawning_task_during_shutdown_causes_task_to_be_cancelled() { sleep(Duration::from_millis(200)).await; assert!(!toplevel_finished.get()); assert!(!subsys_finished.get()); - assert!(!shutdown_token.is_shutting_down()); + assert!(!shutdown_token.is_cancelled()); assert!(!nested_finished.get()); - shutdown_token.shutdown(); - timeout(Duration::from_millis(200), toplevel_finished.wait()) + shutdown_token.cancel(); + timeout(Duration::from_millis(300), toplevel_finished.wait()) .await .unwrap(); - // Assert that subsystem did not get past spawning the task, as spawning a task while shutting - // down causes a panic. assert!(subsys_finished.get()); - assert!(!nested_finished.get()); - - // Assert nested was canceled and didn't continue running in the background - sleep(Duration::from_millis(500)).await; - assert!(!nested_finished.get()); + assert!(nested_finished.get()); }, ); } #[tokio::test(flavor = "multi_thread", worker_threads = 3)] async fn double_panic_does_not_stop_graceful_shutdown() { - setup(); - let (subsys_finished, set_subsys_finished) = Event::create(); let subsys3 = |subsys: SubsystemHandle| async move { @@ -497,38 +485,41 @@ async fn double_panic_does_not_stop_graceful_shutdown() { }; let subsys1 = move |subsys: SubsystemHandle| async move { - subsys.start::("Subsys2", subsys2); - subsys.start::("Subsys3", subsys3); + subsys.start::(SubsystemBuilder::new("Subsys2", subsys2)); + subsys.start::(SubsystemBuilder::new("Subsys3", subsys3)); subsys.on_shutdown_requested().await; sleep(Duration::from_millis(100)).await; panic!("Subsystem1 panicked!") }; - let result = Toplevel::new() - .start::("subsys", subsys1) - .handle_shutdown_requests(Duration::from_millis(500)) - .await; + let result = Toplevel::new(|s| async move { + s.start::(SubsystemBuilder::new("subsys", subsys1)); + }) + .handle_shutdown_requests(Duration::from_millis(500)) + .await; assert!(result.is_err()); assert!(subsys_finished.get()); } #[tokio::test] +#[traced_test] async fn destroying_toplevel_cancels_subsystems() { - setup(); - let (subsys_started, set_subsys_started) = Event::create(); let (subsys_finished, set_subsys_finished) = Event::create(); let subsys1 = move |_subsys: SubsystemHandle| async move { set_subsys_started(); - sleep(Duration::from_millis(100)).await; + sleep(Duration::from_millis(200)).await; set_subsys_finished(); BoxedResult::Ok(()) }; { - let _result = Toplevel::new().start("subsys", subsys1); + let _result = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsys1)); + }); + sleep(Duration::from_millis(100)).await; } sleep(Duration::from_millis(300)).await; @@ -537,22 +528,22 @@ async fn destroying_toplevel_cancels_subsystems() { } #[tokio::test] +#[traced_test] async fn shutdown_triggers_if_all_tasks_ended() { - setup(); - let nested_subsys = move |_subsys: SubsystemHandle| async move { BoxedResult::Ok(()) }; let subsys = move |subsys: SubsystemHandle| async move { - subsys.start("nested", nested_subsys); + subsys.start(SubsystemBuilder::new("nested", nested_subsys)); BoxedResult::Ok(()) }; tokio::time::timeout( Duration::from_millis(100), - Toplevel::new() - .start("subsys1", subsys) - .start("subsys2", subsys) - .handle_shutdown_requests(Duration::from_millis(100)), + Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys1", subsys)); + s.start(SubsystemBuilder::new("subsys2", subsys)); + }) + .handle_shutdown_requests(Duration::from_millis(100)), ) .await .unwrap() @@ -560,12 +551,12 @@ async fn shutdown_triggers_if_all_tasks_ended() { } #[tokio::test] +#[traced_test] async fn shutdown_triggers_if_no_task_exists() { - setup(); - tokio::time::timeout( Duration::from_millis(100), - Toplevel::::new().handle_shutdown_requests(Duration::from_millis(100)), + Toplevel::::new(|_| async {}) + .handle_shutdown_requests(Duration::from_millis(100)), ) .await .unwrap() @@ -573,9 +564,8 @@ async fn shutdown_triggers_if_no_task_exists() { } #[tokio::test] +#[traced_test] async fn destroying_toplevel_cancels_nested_toplevel_subsystems() { - setup(); - let (subsys_started, set_subsys_started) = Event::create(); let (subsys_finished, set_subsys_finished) = Event::create(); @@ -587,14 +577,18 @@ async fn destroying_toplevel_cancels_nested_toplevel_subsystems() { }; let subsys1 = move |_subsys: SubsystemHandle| async move { - Toplevel::new() - .start("subsys2", subsys2) - .handle_shutdown_requests(Duration::from_millis(100)) - .await + Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys2", subsys2)); + }) + .handle_shutdown_requests(Duration::from_millis(100)) + .await }; { - let _result = Toplevel::new().start("subsys", subsys1); + let _result = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsys1)); + }); + sleep(Duration::from_millis(50)).await; } sleep(Duration::from_millis(300)).await; @@ -603,9 +597,8 @@ async fn destroying_toplevel_cancels_nested_toplevel_subsystems() { } #[tokio::test] +#[traced_test] async fn partial_shutdown_request_stops_nested_subsystems() { - setup(); - let (subsys1_started, set_subsys1_started) = Event::create(); let (subsys1_finished, set_subsys1_finished) = Event::create(); let (subsys2_started, set_subsys2_started) = Event::create(); @@ -622,7 +615,7 @@ async fn partial_shutdown_request_stops_nested_subsystems() { }; let subsys2 = move |subsys: SubsystemHandle| async move { set_subsys2_started(); - subsys.start("subsys3", subsys3); + subsys.start(SubsystemBuilder::new("subsys3", subsys3)); subsys.on_shutdown_requested().await; set_subsys2_finished(); BoxedResult::Ok(()) @@ -630,25 +623,26 @@ async fn partial_shutdown_request_stops_nested_subsystems() { let subsys1 = move |subsys: SubsystemHandle| async move { set_subsys1_started(); - let nested_subsys = subsys.start("subsys2", subsys2); + let nested_subsys = subsys.start(SubsystemBuilder::new("subsys2", subsys2)); sleep(Duration::from_millis(200)).await; - subsys - .perform_partial_shutdown(nested_subsys) - .await - .unwrap(); + nested_subsys.change_failure_action(ErrorAction::CatchAndLocalShutdown); + nested_subsys.change_panic_action(ErrorAction::CatchAndLocalShutdown); + nested_subsys.initiate_shutdown(); + nested_subsys.join().await.unwrap(); set_subsys1_shutdown_performed(); subsys.on_shutdown_requested().await; set_subsys1_finished(); BoxedResult::Ok(()) }; - let toplevel = Toplevel::new(); - let shutdown_token = toplevel.get_shutdown_token().clone(); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsys1)); + }); + let shutdown_token = toplevel._get_shutdown_token().clone(); tokio::join!( async { let result = toplevel - .start("subsys", subsys1) .handle_shutdown_requests(Duration::from_millis(500)) .await; assert!(result.is_ok()); @@ -662,15 +656,14 @@ async fn partial_shutdown_request_stops_nested_subsystems() { assert!(subsys2_finished.get()); assert!(subsys3_finished.get()); assert!(subsys1_shutdown_performed.get()); - shutdown_token.shutdown(); + shutdown_token.cancel(); } ); } #[tokio::test] +#[traced_test] async fn partial_shutdown_panic_gets_propagated_correctly() { - setup(); - let (nested_started, set_nested_started) = Event::create(); let (nested_finished, set_nested_finished) = Event::create(); @@ -682,34 +675,39 @@ async fn partial_shutdown_panic_gets_propagated_correctly() { }; let subsys1 = move |subsys: SubsystemHandle| async move { - let handle = subsys.start::("nested", nested_subsys); + let handle = subsys.start::( + SubsystemBuilder::new("nested", nested_subsys) + .on_failure(ErrorAction::CatchAndLocalShutdown) + .on_panic(ErrorAction::CatchAndLocalShutdown), + ); sleep(Duration::from_millis(100)).await; - let result = subsys.perform_partial_shutdown(handle).await; + handle.initiate_shutdown(); + let result = handle.join().await; assert!(matches!( result.err(), - Some(PartialShutdownError::SubsystemsFailed(_)) + Some(SubsystemJoinError::SubsystemsFailed(_)) )); assert!(nested_started.get()); assert!(nested_finished.get()); - assert!(!subsys.local_shutdown_token().is_shutting_down()); + assert!(!subsys.is_shutdown_requested()); subsys.request_shutdown(); BoxedResult::Ok(()) }; - let result = Toplevel::new() - .start("subsys", subsys1) - .handle_shutdown_requests(Duration::from_millis(500)) - .await; + let result = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsys1)); + }) + .handle_shutdown_requests(Duration::from_millis(500)) + .await; assert!(result.is_ok()); } #[tokio::test] +#[traced_test] async fn partial_shutdown_error_gets_propagated_correctly() { - setup(); - let (nested_started, set_nested_started) = Event::create(); let (nested_finished, set_nested_finished) = Event::create(); @@ -721,189 +719,39 @@ async fn partial_shutdown_error_gets_propagated_correctly() { }; let subsys1 = move |subsys: SubsystemHandle| async move { - let handle = subsys.start("nested", nested_subsys); + let handle = subsys.start( + SubsystemBuilder::new("nested", nested_subsys) + .on_failure(ErrorAction::CatchAndLocalShutdown) + .on_panic(ErrorAction::CatchAndLocalShutdown), + ); sleep(Duration::from_millis(100)).await; - let result = subsys.perform_partial_shutdown(handle).await; + handle.initiate_shutdown(); + let result = handle.join().await; assert!(matches!( result.err(), - Some(PartialShutdownError::SubsystemsFailed(_)) + Some(SubsystemJoinError::SubsystemsFailed(_)) )); assert!(nested_started.get()); assert!(nested_finished.get()); - assert!(!subsys.local_shutdown_token().is_shutting_down()); - - subsys.request_shutdown(); - BoxedResult::Ok(()) - }; - - let result = Toplevel::new() - .start("subsys", subsys1) - .handle_shutdown_requests(Duration::from_millis(500)) - .await; - - assert!(result.is_ok()); -} - -#[tokio::test] -async fn partial_shutdown_during_program_shutdown_causes_error() { - setup(); - - let (nested_started, set_nested_started) = Event::create(); - let (nested_finished, set_nested_finished) = Event::create(); - - let nested_subsys = move |subsys: SubsystemHandle| async move { - set_nested_started(); - subsys.on_shutdown_requested().await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsys1 = move |subsys: SubsystemHandle| async move { - let handle = subsys.start("nested", nested_subsys); - sleep(Duration::from_millis(100)).await; + assert!(!subsys.is_shutdown_requested()); subsys.request_shutdown(); - sleep(Duration::from_millis(100)).await; - let result = subsys.perform_partial_shutdown(handle).await; - - assert!(matches!( - result.err(), - Some(PartialShutdownError::AlreadyShuttingDown) - )); - - sleep(Duration::from_millis(100)).await; - - assert!(nested_started.get()); - assert!(nested_finished.get()); - - BoxedResult::Ok(()) - }; - - let result = Toplevel::new() - .start("subsys", subsys1) - .handle_shutdown_requests(Duration::from_millis(500)) - .await; - - assert!(result.is_ok()); -} - -#[tokio::test] -async fn partial_shutdown_on_wrong_parent_causes_error() { - setup(); - - let (nested_started, set_nested_started) = Event::create(); - let (nested_finished, set_nested_finished) = Event::create(); - - let nested_subsys = move |subsys: SubsystemHandle| async move { - set_nested_started(); - subsys.on_shutdown_requested().await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsys1 = move |subsys: SubsystemHandle| async move { - let handle = subsys.start("nested", nested_subsys); - - sleep(Duration::from_millis(100)).await; - - let wrong_parent = |child_subsys: SubsystemHandle| async move { - sleep(Duration::from_millis(100)).await; - let result = child_subsys.perform_partial_shutdown(handle).await; - assert!(matches!( - result.err(), - Some(PartialShutdownError::SubsystemNotFound) - )); - - child_subsys.request_shutdown(); - sleep(Duration::from_millis(100)).await; - - assert!(nested_started.get()); - assert!(nested_finished.get()); - - BoxedResult::Ok(()) - }; - - subsys.start("wrong_parent", wrong_parent); - subsys.on_shutdown_requested().await; - BoxedResult::Ok(()) }; - let result = Toplevel::new() - .start("subsys", subsys1) - .handle_shutdown_requests(Duration::from_millis(500)) - .await; + let result = Toplevel::new(|s| async move { + s.start(SubsystemBuilder::new("subsys", subsys1)); + }) + .handle_shutdown_requests(Duration::from_millis(500)) + .await; assert!(result.is_ok()); } #[tokio::test] -async fn cloned_handles_can_spawn_nested_subsystems() { - setup(); - - let (toplevel_finished, set_toplevel_finished) = Event::create(); - let (subsys_finished, set_subsys_finished) = Event::create(); - let (nested1_finished, set_nested1_finished) = Event::create(); - let (nested2_finished, set_nested2_finished) = Event::create(); - - let nested_subsystem1 = |subsys: SubsystemHandle| async move { - subsys.on_shutdown_requested().await; - set_nested1_finished(); - BoxedResult::Ok(()) - }; - - let nested_subsystem2 = |subsys: SubsystemHandle| async move { - subsys.on_shutdown_requested().await; - set_nested2_finished(); - BoxedResult::Ok(()) - }; - - let subsystem = move |subsys: SubsystemHandle| async move { - let subsys_clone = subsys.clone(); - subsys.start("nested1", nested_subsystem1); - subsys_clone.start("nested2", nested_subsystem2); - subsys_clone.on_shutdown_requested().await; - set_subsys_finished(); - BoxedResult::Ok(()) - }; - - let toplevel = Toplevel::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); - - tokio::join!( - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(200)) - .await; - set_toplevel_finished(); - // Assert panic causes Error propagation to Toplevel - assert!(result.is_ok()); - }, - async { - // Assert that subsystems don't shut down prematurely - sleep(Duration::from_millis(100)).await; - assert!(!subsys_finished.get()); - assert!(!nested1_finished.get()); - assert!(!nested2_finished.get()); - assert!(!toplevel_finished.get()); - - shutdown_token.shutdown(); - sleep(Duration::from_millis(100)).await; - // Assert subsystems did shut down properly - assert!(subsys_finished.get()); - assert!(nested1_finished.get()); - assert!(nested2_finished.get()); - assert!(toplevel_finished.get()); - assert!(shutdown_token.is_shutting_down()); - }, - ); -} - -#[tokio::test] +#[traced_test] async fn subsystem_errors_get_propagated_to_user() { - setup(); - let nested_subsystem1 = |_: SubsystemHandle| async { sleep(Duration::from_millis(100)).await; panic!("Subsystem panicked!"); @@ -915,15 +763,17 @@ async fn subsystem_errors_get_propagated_to_user() { }; let subsystem = move |subsys: SubsystemHandle| async move { - subsys.start::("nested1", nested_subsystem1); - subsys.start("nested2", nested_subsystem2); + subsys.start::(SubsystemBuilder::new("nested1", nested_subsystem1)); + subsys.start(SubsystemBuilder::new("nested2", nested_subsystem2)); sleep(Duration::from_millis(100)).await; subsys.request_shutdown(); BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().start("subsys", subsystem); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); let result = toplevel .handle_shutdown_requests(Duration::from_millis(200)) .await; @@ -933,30 +783,29 @@ async fn subsystem_errors_get_propagated_to_user() { errors.sort_by_key(|el| el.name().to_string()); - let mut iter = errors.into_iter(); + let mut iter = errors.iter(); let el = iter.next().unwrap(); assert!(matches!(el, SubsystemError::Panicked(_))); - assert_eq!("subsys/nested1", el.name()); + assert_eq!("/subsys/nested1", el.name()); let el = iter.next().unwrap(); if let SubsystemError::Failed(name, e) = &el { - assert_eq!("subsys/nested2", name); + assert_eq!("/subsys/nested2", name.as_ref()); assert_eq!("MyGreatError", format!("{}", e)); } else { panic!("Incorrect error type!"); } assert!(matches!(el, SubsystemError::Failed(_, _))); - assert_eq!("subsys/nested2", el.name()); + assert_eq!("/subsys/nested2", el.name()); } else { panic!("Incorrect return value!"); } } #[tokio::test] +#[traced_test] async fn subsystem_errors_get_propagated_to_user_when_timeout() { - setup(); - let nested_subsystem1 = |_: SubsystemHandle| async { sleep(Duration::from_millis(100)).await; panic!("Subsystem panicked!"); @@ -973,53 +822,52 @@ async fn subsystem_errors_get_propagated_to_user_when_timeout() { }; let subsystem = move |subsys: SubsystemHandle| async move { - subsys.start::("nested1", nested_subsystem1); - subsys.start("nested2", nested_subsystem2); - subsys.start::("nested3", nested_subsystem3); + subsys.start::(SubsystemBuilder::new("nested1", nested_subsystem1)); + subsys.start(SubsystemBuilder::new("nested2", nested_subsystem2)); + subsys.start::(SubsystemBuilder::new("nested3", nested_subsystem3)); sleep(Duration::from_millis(100)).await; subsys.request_shutdown(); BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().start("subsys", subsystem); + let toplevel = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + }); let result = toplevel .handle_shutdown_requests(Duration::from_millis(200)) .await; if let Err(GracefulShutdownError::ShutdownTimeout(mut errors)) = result { - assert_eq!(3, errors.len()); + assert_eq!(2, errors.len()); errors.sort_by_key(|el| el.name().to_string()); - let mut iter = errors.into_iter(); + let mut iter = errors.iter(); let el = iter.next().unwrap(); assert!(matches!(el, SubsystemError::Panicked(_))); - assert_eq!("subsys/nested1", el.name()); + assert_eq!("/subsys/nested1", el.name()); let el = iter.next().unwrap(); if let SubsystemError::Failed(name, e) = &el { - assert_eq!("subsys/nested2", name); + assert_eq!("/subsys/nested2", name.as_ref()); assert_eq!("MyGreatError", format!("{}", e)); } else { panic!("Incorrect error type!"); } assert!(matches!(el, SubsystemError::Failed(_, _))); - assert_eq!("subsys/nested2", el.name()); + assert_eq!("/subsys/nested2", el.name()); - let el = iter.next().unwrap(); - assert!(matches!(el, SubsystemError::Cancelled(_))); - assert_eq!("subsys/nested3", el.name()); + assert!(iter.next().is_none()); } else { panic!("Incorrect return value!"); } } #[tokio::test] +#[traced_test] async fn is_shutdown_requested_works_as_intended() { - setup(); - let subsys1 = move |subsys: SubsystemHandle| async move { assert!(!subsys.is_shutdown_requested()); subsys.request_shutdown(); @@ -1027,20 +875,21 @@ async fn is_shutdown_requested_works_as_intended() { BoxedResult::Ok(()) }; - Toplevel::new() - .start("subsys", subsys1) - .handle_shutdown_requests(Duration::from_millis(100)) - .await - .unwrap(); + Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsys1)); + }) + .handle_shutdown_requests(Duration::from_millis(100)) + .await + .unwrap(); } #[cfg(unix)] #[tokio::test] +#[traced_test] async fn shutdown_through_signal() { use nix::sys::signal::{self, Signal}; use nix::unistd::Pid; - - setup(); + use tokio_graceful_shutdown::FutureExt; let subsystem = |subsys: SubsystemHandle| async move { subsys.on_shutdown_requested().await; @@ -1048,7 +897,6 @@ async fn shutdown_through_signal() { BoxedResult::Ok(()) }; - let toplevel = Toplevel::new().catch_signals(); tokio::join!( async { sleep(Duration::from_millis(100)).await; @@ -1057,10 +905,17 @@ async fn shutdown_through_signal() { signal::kill(Pid::this(), Signal::SIGINT).unwrap(); }, async { - let result = toplevel - .start("subsys", subsystem) - .handle_shutdown_requests(Duration::from_millis(400)) - .await; + let result = Toplevel::new(move |s| async move { + s.start(SubsystemBuilder::new("subsys", subsystem)); + assert!(sleep(Duration::from_millis(1000)) + .cancel_on_shutdown(&s) + .await + .is_err()); + assert!(s.is_shutdown_requested()); + }) + .catch_signals() + .handle_shutdown_requests(Duration::from_millis(400)) + .await; assert!(result.is_ok()); }, ); diff --git a/tests/nested_toplevel.rs b/tests/nested_toplevel.rs deleted file mode 100644 index 11b28c1..0000000 --- a/tests/nested_toplevel.rs +++ /dev/null @@ -1,318 +0,0 @@ -use tokio::time::{sleep, Duration}; -use tokio_graceful_shutdown::{SubsystemHandle, Toplevel}; - -pub mod common; -use common::event::Event; -use common::setup; - -use std::error::Error; - -/// Wrapper function to simplify lambdas -type BoxedError = Box; -type BoxedResult = Result<(), BoxedError>; - -#[tokio::test] -async fn nested_toplevel_shuts_down_when_requested() { - setup(); - - let (nested_finished, set_nested_finished) = Event::create(); - let (toplevel_finished, set_toplevel_finished) = Event::create(); - - let nested_subsystem = |subsys: SubsystemHandle| async move { - subsys.on_shutdown_requested().await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsystem = |subsys: SubsystemHandle| async move { - let nested_toplevel = Toplevel::nested(&subsys, "NestedToplevel"); - nested_toplevel - .start("nested", nested_subsystem) - .handle_shutdown_requests(Duration::from_millis(100)) - .await?; - BoxedResult::Ok(()) - }; - - let toplevel = Toplevel::::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); - - tokio::join!( - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - set_toplevel_finished(); - // Assert Ok(()) returncode properly propagates to Toplevel - assert!(result.is_ok()); - }, - async { - sleep(Duration::from_millis(200)).await; - // Assert Ok(()) doesn't cause a shutdown - assert!(!toplevel_finished.get()); - assert!(!nested_finished.get()); - shutdown_token.shutdown(); - sleep(Duration::from_millis(200)).await; - // Assert toplevel sucessfully gets stopped, nothing hangs - assert!(toplevel_finished.get()); - assert!(nested_finished.get()); - }, - ); -} - -#[tokio::test] -async fn nested_toplevel_errors_do_not_get_propagated_up() { - setup(); - - let (nested_finished, set_nested_finished) = Event::create(); - let (subsys_finished, set_subsys_finished) = Event::create(); - let (toplevel_finished, set_toplevel_finished) = Event::create(); - - let nested_error_subsystem = |_subsys: SubsystemHandle| async move { - sleep(Duration::from_millis(200)).await; - BoxedResult::Err("Error from nested subsystem".into()) - }; - let nested_panic_subsystem = |_subsys: SubsystemHandle| async move { - sleep(Duration::from_millis(200)).await; - panic!("Panic from nested subsystem"); - }; - - let nested_subsystem = |subsys: SubsystemHandle| async move { - subsys.on_shutdown_requested().await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsystem = move |subsys: SubsystemHandle| async move { - let nested_toplevel = Toplevel::nested(&subsys, "NestedToplevel"); - let result = nested_toplevel - .start("nested", nested_subsystem) - .start::("nested_panic", nested_panic_subsystem) - .start("nested_error", nested_error_subsystem) - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - assert!(result.is_err()); - set_subsys_finished(); - BoxedResult::Ok(()) - }; - - let toplevel = Toplevel::::new().start("subsys", subsystem); - - tokio::join!( - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - set_toplevel_finished(); - // Assert Ok(()) returncode properly propagates to Toplevel - assert!(result.is_ok()); - }, - async { - sleep(Duration::from_millis(100)).await; - // Assert Ok(()) doesn't cause a shutdown - assert!(!toplevel_finished.get()); - assert!(!nested_finished.get()); - assert!(!subsys_finished.get()); - sleep(Duration::from_millis(200)).await; - // Assert toplevel sucessfully gets stopped, nothing hangs - assert!(toplevel_finished.get()); - assert!(nested_finished.get()); - assert!(subsys_finished.get()); - }, - ); -} - -#[tokio::test] -async fn nested_toplevel_local_shutdown_does_not_get_propagated_up() { - setup(); - - let (nested_finished, set_nested_finished) = Event::create(); - let (nested_toplevel_finished, set_nested_toplevel_finished) = Event::create(); - let (subsys_finished, set_subsys_finished) = Event::create(); - let (toplevel_finished, set_toplevel_finished) = Event::create(); - - let nested_shutdown_subsystem = |subsys: SubsystemHandle| async move { - sleep(Duration::from_millis(200)).await; - subsys.request_shutdown(); - BoxedResult::Ok(()) - }; - - let nested_subsystem = |subsys: SubsystemHandle| async move { - subsys.on_shutdown_requested().await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsystem = move |subsys: SubsystemHandle| async move { - let nested_toplevel = Toplevel::nested(&subsys, "NestedToplevel"); - let result = nested_toplevel - .start("nested", nested_subsystem) - .start("nested_shutdown", nested_shutdown_subsystem) - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - assert!(result.is_ok()); - set_nested_toplevel_finished(); - subsys.on_shutdown_requested().await; - set_subsys_finished(); - BoxedResult::Ok(()) - }; - - let toplevel = Toplevel::::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); - - tokio::join!( - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - set_toplevel_finished(); - // Assert Ok(()) returncode properly propagates to Toplevel - assert!(result.is_ok()); - }, - async { - sleep(Duration::from_millis(100)).await; - // Assert Ok(()) doesn't cause a shutdown - assert!(!toplevel_finished.get()); - assert!(!nested_finished.get()); - assert!(!nested_toplevel_finished.get()); - assert!(!subsys_finished.get()); - sleep(Duration::from_millis(200)).await; - // Assert toplevel sucessfully gets stopped, nothing hangs - assert!(!toplevel_finished.get()); - assert!(nested_finished.get()); - assert!(nested_toplevel_finished.get()); - assert!(!subsys_finished.get()); - shutdown_token.shutdown(); - sleep(Duration::from_millis(200)).await; - assert!(toplevel_finished.get()); - assert!(nested_finished.get()); - assert!(nested_toplevel_finished.get()); - assert!(subsys_finished.get()); - }, - ); -} - -#[tokio::test] -async fn nested_toplevel_global_shutdown_does_get_propagated_up() { - setup(); - - let (nested_finished, set_nested_finished) = Event::create(); - let (nested_toplevel_finished, set_nested_toplevel_finished) = Event::create(); - let (subsys_finished, set_subsys_finished) = Event::create(); - let (toplevel_finished, set_toplevel_finished) = Event::create(); - - let nested_shutdown_subsystem = |subsys: SubsystemHandle| async move { - sleep(Duration::from_millis(200)).await; - subsys.request_global_shutdown(); - BoxedResult::Ok(()) - }; - - let nested_subsystem = |subsys: SubsystemHandle| async move { - subsys.on_shutdown_requested().await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsystem = move |subsys: SubsystemHandle| async move { - let nested_toplevel = Toplevel::nested(&subsys, "NestedToplevel"); - let result = nested_toplevel - .start("nested", nested_subsystem) - .start("nested_shutdown", nested_shutdown_subsystem) - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - assert!(result.is_ok()); - set_nested_toplevel_finished(); - subsys.on_shutdown_requested().await; - set_subsys_finished(); - BoxedResult::Ok(()) - }; - - let toplevel = Toplevel::::new().start("subsys", subsystem); - - tokio::join!( - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - set_toplevel_finished(); - // Assert Ok(()) returncode properly propagates to Toplevel - assert!(result.is_ok()); - }, - async { - sleep(Duration::from_millis(100)).await; - // Assert Ok(()) doesn't cause a shutdown - assert!(!toplevel_finished.get()); - assert!(!nested_finished.get()); - assert!(!nested_toplevel_finished.get()); - assert!(!subsys_finished.get()); - sleep(Duration::from_millis(200)).await; - // Assert toplevel sucessfully gets stopped, nothing hangs - assert!(toplevel_finished.get()); - assert!(nested_finished.get()); - assert!(nested_toplevel_finished.get()); - assert!(subsys_finished.get()); - }, - ); -} - -#[tokio::test] -async fn nested_toplevel_shuts_down_when_subsytems_are_finished() { - setup(); - - let (nested_finished, set_nested_finished) = Event::create(); - let (nested_toplevel_finished, set_nested_toplevel_finished) = Event::create(); - let (subsys_finished, set_subsys_finished) = Event::create(); - let (toplevel_finished, set_toplevel_finished) = Event::create(); - - let nested_subsystem = |_subsys: SubsystemHandle| async move { - sleep(Duration::from_millis(200)).await; - set_nested_finished(); - BoxedResult::Ok(()) - }; - - let subsystem = move |subsys: SubsystemHandle| async move { - let nested_toplevel = Toplevel::nested(&subsys, "NestedToplevel"); - let result = nested_toplevel - .start("nested", nested_subsystem) - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - assert!(result.is_ok()); - set_nested_toplevel_finished(); - subsys.on_shutdown_requested().await; - set_subsys_finished(); - BoxedResult::Ok(()) - }; - - let toplevel = Toplevel::::new().start("subsys", subsystem); - let shutdown_token = toplevel.get_shutdown_token().clone(); - - tokio::join!( - async { - let result = toplevel - .handle_shutdown_requests(Duration::from_millis(100)) - .await; - set_toplevel_finished(); - // Assert Ok(()) returncode properly propagates to Toplevel - assert!(result.is_ok()); - }, - async { - sleep(Duration::from_millis(100)).await; - // Assert Ok(()) doesn't cause a shutdown - assert!(!toplevel_finished.get()); - assert!(!nested_finished.get()); - assert!(!nested_toplevel_finished.get()); - assert!(!subsys_finished.get()); - sleep(Duration::from_millis(200)).await; - // Assert toplevel sucessfully gets stopped, nothing hangs - assert!(!toplevel_finished.get()); - assert!(nested_finished.get()); - assert!(nested_toplevel_finished.get()); - assert!(!subsys_finished.get()); - shutdown_token.shutdown(); - sleep(Duration::from_millis(200)).await; - assert!(toplevel_finished.get()); - assert!(nested_finished.get()); - assert!(nested_toplevel_finished.get()); - assert!(subsys_finished.get()); - }, - ); -}