├── .editorconfig ├── .github └── workflows │ └── rust.yml ├── .gitignore ├── CHANGELOG.md ├── Cargo.toml ├── LICENSE-APACHE ├── LICENSE-BSD ├── README.md ├── README.tpl ├── examples └── tcp-echo-server.rs ├── rustfmt.toml ├── src ├── lib.rs ├── shutdown_complete.rs ├── shutdown_signal.rs ├── waker_list.rs ├── wrap_cancel.rs ├── wrap_delay_shutdown.rs └── wrap_trigger_shutdown.rs └── tests └── shutdown.rs /.editorconfig: -------------------------------------------------------------------------------- 1 | root = true 2 | 3 | [*] 4 | indent_style = tab 5 | indent_size = 4 6 | end_of_line = lf 7 | trim_trailing_whitespace = true 8 | insert_final_newline = true 9 | -------------------------------------------------------------------------------- /.github/workflows/rust.yml: -------------------------------------------------------------------------------- 1 | name: tests 2 | on: 3 | push: { branches: ["main"] } 4 | pull_request: { branches: ["*"] } 5 | 6 | jobs: 7 | build_and_test: 8 | name: Build and test 9 | runs-on: ubuntu-latest 10 | steps: 11 | - name: Checkout code 12 | uses: actions/checkout@master 13 | - uses: actions/cache@v2 14 | with: 15 | path: | 16 | ~/.cargo/registry 17 | ~/.cargo/git 18 | target 19 | key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }} 20 | restore-keys: ${{ runner.os }}-cargo 21 | - name: Build 22 | uses: actions-rs/cargo@v1 23 | with: 24 | command: build 25 | args: --workspace --all-targets --all-features --color=always 26 | - name: Test 27 | uses: actions-rs/cargo@v1 28 | with: 29 | command: test 30 | args: --workspace --all-targets --all-features --color=always 31 | 32 | clippy: 33 | name: Clippy 34 | runs-on: ubuntu-latest 35 | steps: 36 | - name: Checkout code 37 | uses: actions/checkout@master 38 | - uses: actions/cache@v2 39 | with: 40 | path: | 41 | ~/.cargo/registry 42 | ~/.cargo/git 43 | target 44 | key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }} 45 | restore-keys: ${{ runner.os }}-cargo 46 | - name: Clippy 47 | uses: actions-rs-plus/clippy-check@v2.1.1 48 | with: 49 | args: --workspace --all-targets --all-features 50 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | /Cargo.lock 3 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Version 0.2.2 - 2024-03-22 2 | * Fix bug where the list of wakers to trigger on shutdown or shutdown completion could grow indefinitely. 3 | 4 | # Version 0.2.1 - 2023-10-08 5 | * Fix `ShutdownManager::wait_shutdown_complete()` never completing if called when no shutdown was triggered yet and no delay tokens exist. 6 | 7 | # Version 0.2.0 - 2023-09-26: 8 | * Rename `Shutdown` struct to `ShutdownManager`. 9 | * Rename `ShutdownManager` methods: 10 | * `shutdown()` is now called `trigger_shutdown()` 11 | * `shutdown_started()` is now called `is_shutdown_triggered()`. 12 | * `shutdown_completed()` is now called `is_shutdown_completed()`. 13 | * `wrap_vital()` is now called `wrap_trigger_shutdown()`. 14 | * `wrap_wait()` is now called `wrap_delay_shutdown()`. 15 | * Rename `VitalToken` to `TriggerShutdownToken`. 16 | * Rename `TriggerShutdownToken::wrap_vital()` to `wrap_future()`. 17 | * Rename `DelayShutdownToken::wrap_wait()` to `wrap_future()`. 18 | * All types now take a generic parameter `T: Clone` for the shutdown reason. 19 | * Add a parameter for the shutdown reason in `ShutdownManager` methods `trigger_shutdown()`, `wrap_trigger_shutdown()` and `trigger_shutdown_token()`. 20 | * Add `ShutdownManager::shutdown_reason()` to retrieve the shutdown reason. 21 | * Return the shutdown reason from `ShutdownManager::wait_shutdown_triggered()` and `ShutdownManager::wait_shutdown_complete()`. 22 | * Add `shutdown_reason` field to `ShutdownAlreadyCompleted` struct. 23 | * Return a `ShutdownAlreadyStarted` error when calling `trigger_shutdown()` multiple times. 24 | * Change the output type of `WrapCancel` futures from `Option` to a `Result` with the shutdown reason as error. 25 | 26 | # Version 0.1.4 - 2024-03-22 27 | * Fix shutdown_complete() completing before shutdown even started. 28 | * Fix bug where the list of wakers to trigger on shutdown or shutdown completion could grow indefinitely. 29 | 30 | 31 | # Version 0.1.3 - 2023-08-14 32 | * Mark all future wrappers as `#[must_use]`. 33 | 34 | # Version 0.1.2 - 2021-10-30 35 | * Update README. 36 | 37 | # Version 0.1.1 - 2021-10-30 38 | * Improve TCP echo server example. 39 | 40 | # Version 0.1.0 - 2021-10-28 41 | * Add an example with a tokio-based TCP echo server. 42 | 43 | # Version 0.1.0-beta2 - 2021-10-04 44 | * Regenerate README.md from library documentation. 45 | 46 | # Version 0.1.0-beta1 - 2021-10-04 47 | * Change `DelayShutdownToken::wrap_wait()` to consume the token. 48 | * Do not consume the `Shutdown` object in `Shutdown::wrap_vital()`. 49 | * Rename `wait_shutdown()` to `wait_shutdown_triggered()`. 50 | 51 | # Version 0.1.0-alpha4 - 2021-09-29 52 | * Fix `shutdown` and `shutdown_complete` notification. 53 | 54 | # Version 0.1.0-alpha3 - 2021-09-22 55 | * Add missing `Clone` impl for `Shutdown`. 56 | 57 | # Version 0.1.0-alpha2 - 2021-09-22 58 | * Fix crate name in README. 59 | 60 | # Version 0.1.0-alpha1 - 2021-09-22 61 | * Initial release. 62 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "async-shutdown" 3 | description = "one-stop solution for async graceful shutdown" 4 | version = "0.2.2" 5 | license = "BSD-2-Clause OR Apache-2.0" 6 | repository = "https://github.com/de-vri-es/async-shutdown-rs" 7 | readme = "README.md" 8 | documentation = "https://docs.rs/async-shutdown" 9 | 10 | keywords = ["async", "shutdown", "graceful", "graceful-shutdown"] 11 | categories = ["asynchronous"] 12 | 13 | edition = "2018" 14 | 15 | [dev-dependencies] 16 | assert2 = "0.3.4" 17 | tokio = { version = "1.12.0", features = ["io-util", "macros", "net", "rt-multi-thread", "signal", "time"] } 18 | futures = "0.3.17" 19 | -------------------------------------------------------------------------------- /LICENSE-APACHE: -------------------------------------------------------------------------------- 1 | Copyright 2021, Maarten de Vries 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | -------------------------------------------------------------------------------- /LICENSE-BSD: -------------------------------------------------------------------------------- 1 | Copyright 2021, Maarten de Vries 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions are met: 5 | 6 | 1. Redistributions of source code must retain the above copyright notice, this 7 | list of conditions and the following disclaimer. 8 | 9 | 2. Redistributions in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 14 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 15 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 16 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 17 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 18 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 19 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 20 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 21 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 22 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # async-shutdown 2 | 3 | Runtime agnostic one-stop solution for graceful shutdown in asynchronous code. 4 | 5 | This crate addresses two separate but related problems regarding graceful shutdown: 6 | * You have to be able to stop running futures when a shutdown signal is given. 7 | * You have to be able to wait for futures to finish potential clean-up. 8 | * You want to know why the shutdown was triggered (for example to set your process exit code). 9 | 10 | All of these problems are handled by the [`ShutdownManager`] struct. 11 | 12 | ## Stopping running futures 13 | You can get a future to wait for the shutdown signal with [`ShutdownManager::wait_shutdown_triggered()`]. 14 | In this case you must write your async code to react to the shutdown signal appropriately. 15 | 16 | Alternatively, you can wrap a future to be cancelled (by being dropped) when the shutdown is triggered with [`ShutdownManager::wrap_cancel()`]. 17 | This doesn't require the wrapped future to know anything about the shutdown signal, 18 | but it also doesn't allow the future to run custom shutdown code. 19 | 20 | To trigger the shutdown signal, simply call [`ShutdownManager::trigger_shutdown(reason)`][`ShutdownManager::trigger_shutdown()`]. 21 | The shutdown reason can be any type, as long as it implements [`Clone`]. 22 | If you want to pass a non-[`Clone`] object or an object that is expensive to clone, you can wrap it in an [`Arc`]. 23 | 24 | ## Waiting for futures to complete. 25 | You may also want to wait for some futures to complete before actually shutting down instead of just dropping them. 26 | This might be important to cleanly shutdown and prevent data loss. 27 | You can do that with [`ShutdownManager::wait_shutdown_complete()`]. 28 | That function returns a future that only completes when the shutdown is "completed". 29 | 30 | You must also prevent the shutdown from completing too early by calling [`ShutdownManager::delay_shutdown_token()`] or [`ShutdownManager::wrap_delay_shutdown()`]. 31 | The [`ShutdownManager::delay_shutdown_token()`] function gives you a [`DelayShutdownToken`] which prevents the shutdown from completing. 32 | To allow the shutdown to finish, simply drop the token. 33 | Alternatively, [`ShutdownManager::wrap_delay_shutdown()`] wraps an existing future, 34 | and will prevent the shutdown from completing until the future either completes or is dropped. 35 | 36 | Note that you can only delay the shutdown completion if it has not completed already. 37 | If the shutdown is already complete those functions will return an error. 38 | 39 | You can also use a token to wrap a future with [`DelayShutdownToken::wrap_future()`]. 40 | If you already have a token, this allows you to wrap a future without having to worry that the shutdown might already be completed. 41 | 42 | ## Automatically triggering shutdowns 43 | You can also trigger a shutdown automatically using a [`TriggerShutdownToken`]. 44 | Call [`ShutdownManager::trigger_shutdown_token()`] to obtain the token. 45 | When the token is dropped, a shutdown is triggered. 46 | 47 | You can use [`ShutdownManager::wrap_trigger_shutdown()`] or [`TriggerShutdownToken::wrap_future()`] to wrap a future. 48 | When the wrapped future completes (or when it is dropped) it will trigger a shutdown. 49 | This can be used as a convenient way to trigger a shutdown when a vital task stops. 50 | 51 | ## Futures versus Tasks 52 | Be careful when using `JoinHandles` as if they're a regular future. 53 | Depending on your async runtime, when you drop a `JoinHandle` this doesn't normally cause the task to stop. 54 | It may simply detach the join handle from the task, meaning that your task is still running. 55 | If you're not careful, this could still cause data loss on shutdown. 56 | As a rule of thumb, you should usually wrap futures *before* you spawn them on a new task. 57 | 58 | ## Example 59 | 60 | This example is a tokio-based TCP echo server. 61 | It simply echos everything it receives from a peer back to that same peer, 62 | and it uses this crate for graceful shutdown. 63 | 64 | This example is also available in the repository as under the name [`tcp-echo-server`] if you want to run it locally. 65 | 66 | [`tcp-echo-server`]: https://github.com/de-vri-es/async-shutdown-rs/blob/main/examples/tcp-echo-server.rs 67 | 68 | ```rust 69 | use async_shutdown::ShutdownManager; 70 | use std::net::SocketAddr; 71 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 72 | use tokio::net::{TcpListener, TcpStream}; 73 | 74 | #[tokio::main] 75 | async fn main() { 76 | // Create a new shutdown object. 77 | // We will clone it into all tasks that need it. 78 | let shutdown = ShutdownManager::new(); 79 | 80 | // Spawn a task to wait for CTRL+C and trigger a shutdown. 81 | tokio::spawn({ 82 | let shutdown = shutdown.clone(); 83 | async move { 84 | if let Err(e) = tokio::signal::ctrl_c().await { 85 | eprintln!("Failed to wait for CTRL+C: {}", e); 86 | std::process::exit(1); 87 | } else { 88 | eprintln!("\nReceived interrupt signal. Shutting down server..."); 89 | shutdown.trigger_shutdown(0).ok(); 90 | } 91 | } 92 | }); 93 | 94 | // Run the server and set a non-zero exit code if we had an error. 95 | let exit_code = match run_server(shutdown.clone(), "[::]:9372").await { 96 | Ok(()) => { 97 | shutdown.trigger_shutdown(0).ok(); 98 | }, 99 | Err(e) => { 100 | eprintln!("Server task finished with an error: {}", e); 101 | shutdown.trigger_shutdown(1).ok(); 102 | }, 103 | }; 104 | 105 | // Wait for clients to run their cleanup code, then exit. 106 | // Without this, background tasks could be killed before they can run their cleanup code. 107 | let exit_code = shutdown.wait_shutdown_complete().await; 108 | 109 | std::process::exit(exit_code); 110 | } 111 | 112 | async fn run_server(shutdown: ShutdownManager, bind_address: &str) -> std::io::Result<()> { 113 | let server = TcpListener::bind(&bind_address).await?; 114 | eprintln!("Server listening on {}", bind_address); 115 | 116 | // Simply use `wrap_cancel` for everything, since we do not need clean-up for the listening socket. 117 | // See `handle_client` for a case where a future is given the time to perform logging after the shutdown was triggered. 118 | while let Ok(connection) = shutdown.wrap_cancel(server.accept()).await { 119 | let (stream, address) = connection?; 120 | tokio::spawn(handle_client(shutdown.clone(), stream, address)); 121 | } 122 | 123 | Ok(()) 124 | } 125 | 126 | async fn handle_client(shutdown: ShutdownManager, mut stream: TcpStream, address: SocketAddr) { 127 | eprintln!("Accepted new connection from {}", address); 128 | 129 | // Make sure the shutdown doesn't complete until the delay token is dropped. 130 | // 131 | // Getting the token will fail if the shutdown has already started, 132 | // in which case we just log a message and return. 133 | // 134 | // If you already have a future that should be allowed to complete, 135 | // you can also use `shutdown.wrap_delay_shutdown(...)`. 136 | // Here it is easier to use a token though. 137 | let _delay_token = match shutdown.delay_shutdown_token() { 138 | Ok(token) => token, 139 | Err(_) => { 140 | eprintln!("Shutdown already started, closing connection with {}", address); 141 | return; 142 | } 143 | }; 144 | 145 | // Now run the echo loop, but cancel it when the shutdown is triggered. 146 | match shutdown.wrap_cancel(echo_loop(&mut stream)).await { 147 | Ok(Err(e)) => eprintln!("Error in connection {}: {}", address, e), 148 | Ok(Ok(())) => eprintln!("Connection closed by {}", address), 149 | Err(_exit_code) => eprintln!("Shutdown triggered, closing connection with {}", address), 150 | } 151 | 152 | // The delay token will be dropped here, allowing the shutdown to complete. 153 | } 154 | 155 | async fn echo_loop(stream: &mut TcpStream) -> std::io::Result<()> { 156 | // Echo everything we receive back to the peer in a loop. 157 | let mut buffer = vec![0; 512]; 158 | loop { 159 | let read = stream.read(&mut buffer).await?; 160 | if read == 0 { 161 | break; 162 | } 163 | stream.write(&buffer[..read]).await?; 164 | } 165 | 166 | Ok(()) 167 | } 168 | ``` 169 | 170 | [`ShutdownManager`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html 171 | [`ShutdownManager::wait_shutdown_triggered()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wait_shutdown_triggered 172 | [`ShutdownManager::wrap_cancel()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wrap_cancel 173 | [`ShutdownManager::trigger_shutdown()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.trigger_shutdown 174 | [`Clone`]: https://doc.rust-lang.org/stable/std/clone/trait.Clone.html 175 | [`Arc`]: https://doc.rust-lang.org/stable/std/sync/struct.Arc.html 176 | [`ShutdownManager::wait_shutdown_complete()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wait_shutdown_complete 177 | [`ShutdownManager::delay_shutdown_token()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.delay_shutdown_token 178 | [`ShutdownManager::wrap_delay_shutdown()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wrap_delay_shutdown 179 | [`DelayShutdownToken`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.DelayShutdownToken.html 180 | [`DelayShutdownToken::wrap_future()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.DelayShutdownToken.html#method.wrap_future 181 | [`TriggerShutdownToken`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.TriggerShutdownToken.html 182 | [`ShutdownManager::trigger_shutdown_token()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.trigger_shutdown_token 183 | [`ShutdownManager::wrap_trigger_shutdown()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wrap_trigger_shutdown 184 | [`TriggerShutdownToken::wrap_future()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.TriggerShutdownToken.html#method.wrap_future 185 | -------------------------------------------------------------------------------- /README.tpl: -------------------------------------------------------------------------------- 1 | # {{crate}} 2 | 3 | {{readme}} 4 | 5 | [`ShutdownManager`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html 6 | [`ShutdownManager::wait_shutdown_triggered()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wait_shutdown_triggered 7 | [`ShutdownManager::wrap_cancel()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wrap_cancel 8 | [`ShutdownManager::trigger_shutdown()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.trigger_shutdown 9 | [`Clone`]: https://doc.rust-lang.org/stable/std/clone/trait.Clone.html 10 | [`Arc`]: https://doc.rust-lang.org/stable/std/sync/struct.Arc.html 11 | [`ShutdownManager::wait_shutdown_complete()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wait_shutdown_complete 12 | [`ShutdownManager::delay_shutdown_token()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.delay_shutdown_token 13 | [`ShutdownManager::wrap_delay_shutdown()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wrap_delay_shutdown 14 | [`DelayShutdownToken`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.DelayShutdownToken.html 15 | [`DelayShutdownToken::wrap_future()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.DelayShutdownToken.html#method.wrap_future 16 | [`TriggerShutdownToken`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.TriggerShutdownToken.html 17 | [`ShutdownManager::trigger_shutdown_token()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.trigger_shutdown_token 18 | [`ShutdownManager::wrap_trigger_shutdown()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.ShutdownManager.html#method.wrap_trigger_shutdown 19 | [`TriggerShutdownToken::wrap_future()`]: https://docs.rs/async-shutdown/latest/async_shutdown/struct.TriggerShutdownToken.html#method.wrap_future 20 | -------------------------------------------------------------------------------- /examples/tcp-echo-server.rs: -------------------------------------------------------------------------------- 1 | use async_shutdown::ShutdownManager; 2 | use std::net::SocketAddr; 3 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 4 | use tokio::net::{TcpListener, TcpStream}; 5 | 6 | // This example is a tokio-based TCP echo server. 7 | // It simply echos everything it receives from a peer back to that same peer, 8 | // and it uses this crate for graceful shutdown. 9 | 10 | #[tokio::main] 11 | async fn main() { 12 | // Create a new shutdown object. 13 | // We will clone it into all tasks that need it. 14 | let shutdown = ShutdownManager::new(); 15 | 16 | // Spawn a task to wait for CTRL+C and trigger a shutdown. 17 | tokio::spawn({ 18 | let shutdown = shutdown.clone(); 19 | async move { 20 | if let Err(e) = tokio::signal::ctrl_c().await { 21 | eprintln!("Failed to wait for CTRL+C: {}", e); 22 | std::process::exit(1); 23 | } else { 24 | eprintln!("\nReceived interrupt signal. Shutting down server..."); 25 | shutdown.trigger_shutdown(0).ok(); 26 | } 27 | } 28 | }); 29 | 30 | // Run the server and set a non-zero exit code if we had an error. 31 | match run_server(shutdown.clone(), "[::]:9372").await { 32 | Ok(()) => { 33 | shutdown.trigger_shutdown(0).ok(); 34 | }, 35 | Err(e) => { 36 | eprintln!("Server task finished with an error: {}", e); 37 | shutdown.trigger_shutdown(1).ok(); 38 | }, 39 | }; 40 | 41 | // Wait for clients to run their cleanup code, then exit. 42 | // Without this, background tasks could be killed before they can run their cleanup code. 43 | let exit_code = shutdown.wait_shutdown_complete().await; 44 | 45 | std::process::exit(exit_code); 46 | } 47 | 48 | async fn run_server(shutdown: ShutdownManager, bind_address: &str) -> std::io::Result<()> { 49 | let server = TcpListener::bind(&bind_address).await?; 50 | eprintln!("Server listening on {}", bind_address); 51 | 52 | // Simply use `wrap_cancel` for everything, since we do not need clean-up for the listening socket. 53 | // See `handle_client` for a case where a future is given the time to perform logging after the shutdown was triggered. 54 | while let Ok(connection) = shutdown.wrap_cancel(server.accept()).await { 55 | let (stream, address) = connection?; 56 | tokio::spawn(handle_client(shutdown.clone(), stream, address)); 57 | } 58 | 59 | Ok(()) 60 | } 61 | 62 | async fn handle_client(shutdown: ShutdownManager, mut stream: TcpStream, address: SocketAddr) { 63 | eprintln!("Accepted new connection from {}", address); 64 | 65 | // Make sure the shutdown doesn't complete until the delay token is dropped. 66 | // 67 | // Getting the token will fail if the shutdown has already started, 68 | // in which case we just log a message and return. 69 | // 70 | // If you already have a future that should be allowed to complete, 71 | // you can also use `shutdown.wrap_wait(...)`. 72 | // Here it is easier to use a token though. 73 | let _delay_token = match shutdown.delay_shutdown_token() { 74 | Ok(token) => token, 75 | Err(_) => { 76 | eprintln!("Shutdown already started, closing connection with {}", address); 77 | return; 78 | } 79 | }; 80 | 81 | // Now run the echo loop, but cancel it when the shutdown is triggered. 82 | match shutdown.wrap_cancel(echo_loop(&mut stream)).await { 83 | Ok(Err(e)) => eprintln!("Error in connection {}: {}", address, e), 84 | Ok(Ok(())) => eprintln!("Connection closed by {}", address), 85 | Err(_exit_code) => eprintln!("Shutdown triggered, closing connection with {}", address), 86 | } 87 | 88 | // The delay token will be dropped here, allowing the shutdown to complete. 89 | } 90 | 91 | async fn echo_loop(stream: &mut TcpStream) -> std::io::Result<()> { 92 | // Echo everything we receive back to the peer in a loop. 93 | let mut buffer = vec![0; 512]; 94 | loop { 95 | let read = stream.read(&mut buffer).await?; 96 | if read == 0 { 97 | break; 98 | } 99 | stream.write_all(&buffer[..read]).await?; 100 | } 101 | 102 | Ok(()) 103 | } 104 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | hard_tabs = true 2 | tab_spaces = 4 3 | max_width = 120 4 | imports_layout = "HorizontalVertical" 5 | match_block_trailing_comma = true 6 | overflow_delimited_expr = true 7 | reorder_impl_items = true 8 | unstable_features = true 9 | use_field_init_shorthand = true 10 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! Runtime agnostic one-stop solution for graceful shutdown in asynchronous code. 2 | //! 3 | //! This crate addresses two separate but related problems regarding graceful shutdown: 4 | //! * You have to be able to stop running futures when a shutdown signal is given. 5 | //! * You have to be able to wait for futures to finish potential clean-up. 6 | //! * You want to know why the shutdown was triggered (for example to set your process exit code). 7 | //! 8 | //! All of these problems are handled by the [`ShutdownManager`] struct. 9 | //! 10 | //! # Stopping running futures 11 | //! You can get a future to wait for the shutdown signal with [`ShutdownManager::wait_shutdown_triggered()`]. 12 | //! In this case you must write your async code to react to the shutdown signal appropriately. 13 | //! 14 | //! Alternatively, you can wrap a future to be cancelled (by being dropped) when the shutdown is triggered with [`ShutdownManager::wrap_cancel()`]. 15 | //! This doesn't require the wrapped future to know anything about the shutdown signal, 16 | //! but it also doesn't allow the future to run custom shutdown code. 17 | //! 18 | //! To trigger the shutdown signal, simply call [`ShutdownManager::trigger_shutdown(reason)`][`ShutdownManager::trigger_shutdown()`]. 19 | //! The shutdown reason can be any type, as long as it implements [`Clone`]. 20 | //! If you want to pass a non-[`Clone`] object or an object that is expensive to clone, you can wrap it in an [`Arc`]. 21 | //! 22 | //! # Waiting for futures to complete. 23 | //! You may also want to wait for some futures to complete before actually shutting down instead of just dropping them. 24 | //! This might be important to cleanly shutdown and prevent data loss. 25 | //! You can do that with [`ShutdownManager::wait_shutdown_complete()`]. 26 | //! That function returns a future that only completes when the shutdown is "completed". 27 | //! 28 | //! You must also prevent the shutdown from completing too early by calling [`ShutdownManager::delay_shutdown_token()`] or [`ShutdownManager::wrap_delay_shutdown()`]. 29 | //! The [`ShutdownManager::delay_shutdown_token()`] function gives you a [`DelayShutdownToken`] which prevents the shutdown from completing. 30 | //! To allow the shutdown to finish, simply drop the token. 31 | //! Alternatively, [`ShutdownManager::wrap_delay_shutdown()`] wraps an existing future, 32 | //! and will prevent the shutdown from completing until the future either completes or is dropped. 33 | //! 34 | //! Note that you can only delay the shutdown completion if it has not completed already. 35 | //! If the shutdown is already complete those functions will return an error. 36 | //! 37 | //! You can also use a token to wrap a future with [`DelayShutdownToken::wrap_future()`]. 38 | //! If you already have a token, this allows you to wrap a future without having to worry that the shutdown might already be completed. 39 | //! 40 | //! # Automatically triggering shutdowns 41 | //! You can also trigger a shutdown automatically using a [`TriggerShutdownToken`]. 42 | //! Call [`ShutdownManager::trigger_shutdown_token()`] to obtain the token. 43 | //! When the token is dropped, a shutdown is triggered. 44 | //! 45 | //! You can use [`ShutdownManager::wrap_trigger_shutdown()`] or [`TriggerShutdownToken::wrap_future()`] to wrap a future. 46 | //! When the wrapped future completes (or when it is dropped) it will trigger a shutdown. 47 | //! This can be used as a convenient way to trigger a shutdown when a vital task stops. 48 | //! 49 | //! # Futures versus Tasks 50 | //! Be careful when using `JoinHandles` as if they're a regular future. 51 | //! Depending on your async runtime, when you drop a `JoinHandle` this doesn't normally cause the task to stop. 52 | //! It may simply detach the join handle from the task, meaning that your task is still running. 53 | //! If you're not careful, this could still cause data loss on shutdown. 54 | //! As a rule of thumb, you should usually wrap futures *before* you spawn them on a new task. 55 | //! 56 | //! # Example 57 | //! 58 | //! This example is a tokio-based TCP echo server. 59 | //! It simply echos everything it receives from a peer back to that same peer, 60 | //! and it uses this crate for graceful shutdown. 61 | //! 62 | //! This example is also available in the repository as under the name [`tcp-echo-server`] if you want to run it locally. 63 | //! 64 | //! [`tcp-echo-server`]: https://github.com/de-vri-es/async-shutdown-rs/blob/main/examples/tcp-echo-server.rs 65 | //! 66 | //! ```no_run 67 | //! use async_shutdown::ShutdownManager; 68 | //! use std::net::SocketAddr; 69 | //! use tokio::io::{AsyncReadExt, AsyncWriteExt}; 70 | //! use tokio::net::{TcpListener, TcpStream}; 71 | //! 72 | //! #[tokio::main] 73 | //! async fn main() { 74 | //! // Create a new shutdown object. 75 | //! // We will clone it into all tasks that need it. 76 | //! let shutdown = ShutdownManager::new(); 77 | //! 78 | //! // Spawn a task to wait for CTRL+C and trigger a shutdown. 79 | //! tokio::spawn({ 80 | //! let shutdown = shutdown.clone(); 81 | //! async move { 82 | //! if let Err(e) = tokio::signal::ctrl_c().await { 83 | //! eprintln!("Failed to wait for CTRL+C: {}", e); 84 | //! std::process::exit(1); 85 | //! } else { 86 | //! eprintln!("\nReceived interrupt signal. Shutting down server..."); 87 | //! shutdown.trigger_shutdown(0).ok(); 88 | //! } 89 | //! } 90 | //! }); 91 | //! 92 | //! // Run the server and set a non-zero exit code if we had an error. 93 | //! let exit_code = match run_server(shutdown.clone(), "[::]:9372").await { 94 | //! Ok(()) => { 95 | //! shutdown.trigger_shutdown(0).ok(); 96 | //! }, 97 | //! Err(e) => { 98 | //! eprintln!("Server task finished with an error: {}", e); 99 | //! shutdown.trigger_shutdown(1).ok(); 100 | //! }, 101 | //! }; 102 | //! 103 | //! // Wait for clients to run their cleanup code, then exit. 104 | //! // Without this, background tasks could be killed before they can run their cleanup code. 105 | //! let exit_code = shutdown.wait_shutdown_complete().await; 106 | //! 107 | //! std::process::exit(exit_code); 108 | //! } 109 | //! 110 | //! async fn run_server(shutdown: ShutdownManager, bind_address: &str) -> std::io::Result<()> { 111 | //! let server = TcpListener::bind(&bind_address).await?; 112 | //! eprintln!("Server listening on {}", bind_address); 113 | //! 114 | //! // Simply use `wrap_cancel` for everything, since we do not need clean-up for the listening socket. 115 | //! // See `handle_client` for a case where a future is given the time to perform logging after the shutdown was triggered. 116 | //! while let Ok(connection) = shutdown.wrap_cancel(server.accept()).await { 117 | //! let (stream, address) = connection?; 118 | //! tokio::spawn(handle_client(shutdown.clone(), stream, address)); 119 | //! } 120 | //! 121 | //! Ok(()) 122 | //! } 123 | //! 124 | //! async fn handle_client(shutdown: ShutdownManager, mut stream: TcpStream, address: SocketAddr) { 125 | //! eprintln!("Accepted new connection from {}", address); 126 | //! 127 | //! // Make sure the shutdown doesn't complete until the delay token is dropped. 128 | //! // 129 | //! // Getting the token will fail if the shutdown has already started, 130 | //! // in which case we just log a message and return. 131 | //! // 132 | //! // If you already have a future that should be allowed to complete, 133 | //! // you can also use `shutdown.wrap_delay_shutdown(...)`. 134 | //! // Here it is easier to use a token though. 135 | //! let _delay_token = match shutdown.delay_shutdown_token() { 136 | //! Ok(token) => token, 137 | //! Err(_) => { 138 | //! eprintln!("Shutdown already started, closing connection with {}", address); 139 | //! return; 140 | //! } 141 | //! }; 142 | //! 143 | //! // Now run the echo loop, but cancel it when the shutdown is triggered. 144 | //! match shutdown.wrap_cancel(echo_loop(&mut stream)).await { 145 | //! Ok(Err(e)) => eprintln!("Error in connection {}: {}", address, e), 146 | //! Ok(Ok(())) => eprintln!("Connection closed by {}", address), 147 | //! Err(_exit_code) => eprintln!("Shutdown triggered, closing connection with {}", address), 148 | //! } 149 | //! 150 | //! // The delay token will be dropped here, allowing the shutdown to complete. 151 | //! } 152 | //! 153 | //! async fn echo_loop(stream: &mut TcpStream) -> std::io::Result<()> { 154 | //! // Echo everything we receive back to the peer in a loop. 155 | //! let mut buffer = vec![0; 512]; 156 | //! loop { 157 | //! let read = stream.read(&mut buffer).await?; 158 | //! if read == 0 { 159 | //! break; 160 | //! } 161 | //! stream.write(&buffer[..read]).await?; 162 | //! } 163 | //! 164 | //! Ok(()) 165 | //! } 166 | //! ``` 167 | 168 | #![warn(missing_docs)] 169 | 170 | use std::future::Future; 171 | use std::sync::{Arc, Mutex}; 172 | 173 | mod shutdown_complete; 174 | pub use shutdown_complete::ShutdownComplete; 175 | 176 | mod shutdown_signal; 177 | pub use shutdown_signal::ShutdownSignal; 178 | 179 | mod wrap_cancel; 180 | use waker_list::WakerList; 181 | pub use wrap_cancel::WrapCancel; 182 | 183 | mod wrap_trigger_shutdown; 184 | pub use wrap_trigger_shutdown::WrapTriggerShutdown; 185 | 186 | mod wrap_delay_shutdown; 187 | pub use wrap_delay_shutdown::WrapDelayShutdown; 188 | 189 | mod waker_list; 190 | 191 | /// Shutdown manager for asynchronous tasks and futures. 192 | /// 193 | /// The shutdown manager allows you to: 194 | /// * Signal futures to shutdown or forcibly cancel them (by dropping them). 195 | /// * Wait for futures to perform their clean-up after a shutdown was triggered. 196 | /// * Retrieve the shutdown reason after the shutdown was triggered. 197 | /// 198 | /// The shutdown manager can be cloned and shared with multiple tasks. 199 | /// Each clone uses the same internal state. 200 | #[derive(Clone)] 201 | pub struct ShutdownManager { 202 | inner: Arc>>, 203 | } 204 | 205 | impl ShutdownManager { 206 | /// Create a new shutdown manager. 207 | #[inline] 208 | pub fn new() -> Self { 209 | Self { 210 | inner: Arc::new(Mutex::new(ShutdownManagerInner::new())), 211 | } 212 | } 213 | 214 | /// Check if the shutdown has been triggered. 215 | #[inline] 216 | pub fn is_shutdown_triggered(&self) -> bool { 217 | self.inner.lock().unwrap().shutdown_reason.is_some() 218 | } 219 | 220 | /// Check if the shutdown has completed. 221 | #[inline] 222 | pub fn is_shutdown_completed(&self) -> bool { 223 | let inner = self.inner.lock().unwrap(); 224 | inner.shutdown_reason.is_some() && inner.delay_tokens == 0 225 | } 226 | 227 | /// Get the shutdown reason, if the shutdown has been triggered. 228 | /// 229 | /// Returns [`None`] if the shutdown has not been triggered yet. 230 | #[inline] 231 | pub fn shutdown_reason(&self) -> Option { 232 | self.inner.lock().unwrap().shutdown_reason.clone() 233 | } 234 | 235 | /// Asynchronously wait for the shutdown to be triggered. 236 | /// 237 | /// This returns a future that completes when the shutdown is triggered. 238 | /// The future can be cloned and sent to other threads or tasks freely. 239 | /// 240 | /// If the shutdown is already triggered, the returned future immediately resolves. 241 | /// 242 | /// You can also use `ShutdownSignal::wrap_cancel()` of the returned object 243 | /// to automatically cancel a future when the shutdown signal is received. 244 | /// This is identical to `Self::wrap_cancel()`. 245 | #[inline] 246 | pub fn wait_shutdown_triggered(&self) -> ShutdownSignal { 247 | ShutdownSignal { 248 | inner: self.inner.clone(), 249 | waker_token: None, 250 | } 251 | } 252 | 253 | /// Asynchronously wait for the shutdown to complete. 254 | /// 255 | /// This returns a future that completes when the shutdown is complete. 256 | /// The future can be cloned and sent to other threads or tasks freely. 257 | /// 258 | /// The shutdown is complete when all [`DelayShutdownToken`] are dropped 259 | /// and all [`WrapDelayShutdown`] futures have completed or are dropped. 260 | #[inline] 261 | pub fn wait_shutdown_complete(&self) -> ShutdownComplete { 262 | ShutdownComplete { 263 | inner: self.inner.clone(), 264 | waker_token: None, 265 | } 266 | } 267 | 268 | /// Trigger the shutdown. 269 | /// 270 | /// This will cause all [`ShutdownSignal`] and [`WrapCancel`] futures associated with this shutdown manager to be resolved. 271 | /// 272 | /// The shutdown will not be considered complete until all [`DelayShutdownTokens`][DelayShutdownToken] are dropped. 273 | /// 274 | /// If the shutdown was already started, this function returns an error. 275 | #[inline] 276 | pub fn trigger_shutdown(&self, reason: T) -> Result<(), ShutdownAlreadyStarted> { 277 | self.inner.lock().unwrap().shutdown(reason) 278 | } 279 | 280 | /// Wrap a future so that it is cancelled (dropped) when the shutdown is triggered. 281 | /// 282 | /// The returned future completes with `Err(shutdown_reason)` if the shutdown is triggered, 283 | /// and with `Ok(x)` if the wrapped future completes first. 284 | #[inline] 285 | pub fn wrap_cancel(&self, future: F) -> WrapCancel { 286 | self.wait_shutdown_triggered().wrap_cancel(future) 287 | } 288 | 289 | /// Wrap a future to cause a shutdown when the future completes or when it is dropped. 290 | #[inline] 291 | pub fn wrap_trigger_shutdown(&self, shutdown_reason: T, future: F) -> WrapTriggerShutdown { 292 | self.trigger_shutdown_token(shutdown_reason).wrap_future(future) 293 | } 294 | 295 | /// Wrap a future to delay shutdown completion until the wrapped future completes or until it is dropped. 296 | /// 297 | /// The returned future transparently completes with the value of the wrapped future. 298 | /// However, the shutdown will not be considered complete until the wrapped future completes or is dropped. 299 | /// 300 | /// If the shutdown has already completed, this function returns an error. 301 | #[inline] 302 | pub fn wrap_delay_shutdown(&self, future: F) -> Result, ShutdownAlreadyCompleted> { 303 | Ok(self.delay_shutdown_token()?.wrap_future(future)) 304 | } 305 | 306 | /// Get a token that delays shutdown completion as long as it exists. 307 | /// 308 | /// The manager keeps track of all the tokens it hands out. 309 | /// The tokens can be cloned and sent to different threads and tasks. 310 | /// All tokens (including the clones) must be dropped before the shutdown is considered to be complete. 311 | /// 312 | /// If the shutdown has already completed, this function returns an error. 313 | /// 314 | /// If you want to delay the shutdown until a future completes, 315 | /// consider using [`Self::wrap_delay_shutdown()`] instead. 316 | #[inline] 317 | pub fn delay_shutdown_token(&self) -> Result, ShutdownAlreadyCompleted> { 318 | let mut inner = self.inner.lock().unwrap(); 319 | // Shutdown already completed, can't delay completion anymore. 320 | if inner.delay_tokens == 0 { 321 | if let Some(reason) = &inner.shutdown_reason { 322 | return Err(ShutdownAlreadyCompleted::new(reason.clone())); 323 | } 324 | } 325 | 326 | inner.increase_delay_count(); 327 | Ok(DelayShutdownToken { 328 | inner: self.inner.clone(), 329 | }) 330 | } 331 | 332 | /// Get a token that triggers a shutdown when dropped. 333 | /// 334 | /// When a [`TriggerShutdownToken`] is dropped, the shutdown is triggered automatically. 335 | /// This applies to *any* token. 336 | /// If you clone a token five times and drop one of them, it will trigger a shutdown/ 337 | /// 338 | /// You can also use [`Self::wrap_trigger_shutdown()`] to wrap a future so that a shutdown is triggered 339 | /// when the future completes or if it is dropped. 340 | #[inline] 341 | pub fn trigger_shutdown_token(&self, shutdown_reason: T) -> TriggerShutdownToken { 342 | TriggerShutdownToken { 343 | shutdown_reason: Arc::new(Mutex::new(Some(shutdown_reason))), 344 | inner: self.inner.clone(), 345 | } 346 | } 347 | } 348 | 349 | impl Default for ShutdownManager { 350 | #[inline] 351 | fn default() -> Self { 352 | Self::new() 353 | } 354 | } 355 | 356 | /// Token that delays shutdown completion as long as it exists. 357 | /// 358 | /// The token can be cloned and sent to different threads and tasks freely. 359 | /// 360 | /// All clones must be dropped before the shutdown can complete. 361 | pub struct DelayShutdownToken { 362 | inner: Arc>>, 363 | } 364 | 365 | impl DelayShutdownToken { 366 | /// Wrap a future to delay shutdown completion until the wrapped future completes or until it is dropped. 367 | /// 368 | /// This consumes the token to avoid keeping an unused token around by accident, which would delay shutdown indefinitely. 369 | /// If you wish to use the token multiple times, you can clone it first: 370 | /// ``` 371 | /// # let shutdown = async_shutdown::ShutdownManager::<()>::new(); 372 | /// # let delay_shutdown_token = shutdown.delay_shutdown_token().unwrap(); 373 | /// # let future = async { () }; 374 | /// let future = delay_shutdown_token.clone().wrap_future(future); 375 | /// ``` 376 | /// 377 | /// The returned future transparently completes with the value of the wrapped future. 378 | /// However, the shutdown will not be considered complete until the future completes or is dropped. 379 | #[inline] 380 | pub fn wrap_future(self, future: F) -> WrapDelayShutdown { 381 | WrapDelayShutdown { 382 | delay_token: Some(self), 383 | future, 384 | } 385 | } 386 | } 387 | 388 | impl Clone for DelayShutdownToken { 389 | #[inline] 390 | fn clone(&self) -> Self { 391 | self.inner.lock().unwrap().increase_delay_count(); 392 | DelayShutdownToken { 393 | inner: self.inner.clone(), 394 | } 395 | } 396 | } 397 | 398 | impl Drop for DelayShutdownToken { 399 | #[inline] 400 | fn drop(&mut self) { 401 | self.inner.lock().unwrap().decrease_delay_count(); 402 | } 403 | } 404 | 405 | /// Token that triggers a shutdown when it is dropped. 406 | /// 407 | /// The token can be cloned and sent to different threads and tasks freely. 408 | /// If *one* of the cloned tokens is dropped, a shutdown is triggered. 409 | /// Even if the rest of the clones still exist. 410 | #[derive(Clone)] 411 | pub struct TriggerShutdownToken { 412 | shutdown_reason: Arc>>, 413 | inner: Arc>>, 414 | } 415 | 416 | impl TriggerShutdownToken { 417 | /// Wrap a future to trigger a shutdown when it completes or is dropped. 418 | /// 419 | /// This consumes the token to avoid accidentally dropping the token 420 | /// after wrapping a future and instantly causing a shutdown. 421 | /// 422 | /// If you need to keep the token around, you can clone it first: 423 | /// ``` 424 | /// # let trigger_shutdown_token = async_shutdown::ShutdownManager::new().trigger_shutdown_token(()); 425 | /// # let future = async { () }; 426 | /// let future = trigger_shutdown_token.clone().wrap_future(future); 427 | /// ``` 428 | #[inline] 429 | pub fn wrap_future(self, future: F) -> WrapTriggerShutdown { 430 | WrapTriggerShutdown { 431 | trigger_shutdown_token: Some(self), 432 | future, 433 | } 434 | } 435 | 436 | /// Drop the token without causing a shutdown. 437 | /// 438 | /// This is equivalent to calling [`std::mem::forget()`] on the token. 439 | #[inline] 440 | pub fn forget(self) { 441 | std::mem::forget(self) 442 | } 443 | } 444 | 445 | impl Drop for TriggerShutdownToken { 446 | #[inline] 447 | fn drop(&mut self) { 448 | let mut inner = self.inner.lock().unwrap(); 449 | let reason = self.shutdown_reason.lock().unwrap().take(); 450 | if let Some(reason) = reason { 451 | inner.shutdown(reason).ok(); 452 | } 453 | } 454 | } 455 | 456 | struct ShutdownManagerInner { 457 | /// The shutdown reason. 458 | shutdown_reason: Option, 459 | 460 | /// Number of delay tokens in existence. 461 | /// 462 | /// Must reach 0 before shutdown can complete. 463 | delay_tokens: usize, 464 | 465 | /// Tasks to wake when a shutdown is triggered. 466 | on_shutdown: WakerList, 467 | 468 | /// Tasks to wake when the shutdown is complete. 469 | on_shutdown_complete: WakerList, 470 | } 471 | 472 | impl ShutdownManagerInner { 473 | fn new() -> Self { 474 | Self { 475 | shutdown_reason: None, 476 | delay_tokens: 0, 477 | on_shutdown_complete: WakerList::new(), 478 | on_shutdown: WakerList::new(), 479 | } 480 | } 481 | 482 | fn increase_delay_count(&mut self) { 483 | self.delay_tokens += 1; 484 | } 485 | 486 | fn decrease_delay_count(&mut self) { 487 | self.delay_tokens -= 1; 488 | if self.delay_tokens == 0 { 489 | self.notify_shutdown_complete(); 490 | } 491 | } 492 | 493 | fn shutdown(&mut self, reason: T) -> Result<(), ShutdownAlreadyStarted> { 494 | match &self.shutdown_reason { 495 | Some(original_reason) => { 496 | Err(ShutdownAlreadyStarted::new(original_reason.clone(), reason)) 497 | }, 498 | None => { 499 | self.shutdown_reason = Some(reason); 500 | self.on_shutdown.wake_all(); 501 | if self.delay_tokens == 0 { 502 | self.notify_shutdown_complete() 503 | } 504 | Ok(()) 505 | }, 506 | } 507 | } 508 | 509 | fn notify_shutdown_complete(&mut self) { 510 | self.on_shutdown_complete.wake_all(); 511 | } 512 | } 513 | 514 | /// Error returned when you try to trigger the shutdown multiple times on the same [`ShutdownManager`]. 515 | #[derive(Debug, Clone)] 516 | #[non_exhaustive] 517 | pub struct ShutdownAlreadyStarted { 518 | /// The shutdown reason of the already started shutdown. 519 | pub shutdown_reason: T, 520 | 521 | /// The provided reason that was ignored because the shutdown was already started. 522 | pub ignored_reason: T, 523 | } 524 | 525 | impl ShutdownAlreadyStarted { 526 | pub(crate) const fn new(shutdown_reason: T, ignored_reason:T ) -> Self { 527 | Self { shutdown_reason, ignored_reason } 528 | } 529 | } 530 | 531 | impl std::error::Error for ShutdownAlreadyStarted {} 532 | 533 | impl std::fmt::Display for ShutdownAlreadyStarted { 534 | fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { 535 | write!(f, "shutdown has already started, can not delay shutdown completion") 536 | } 537 | } 538 | 539 | /// Error returned when trying to delay a shutdown that has already completed. 540 | #[derive(Debug)] 541 | #[non_exhaustive] 542 | pub struct ShutdownAlreadyCompleted { 543 | /// The shutdown reason of the already completed shutdown. 544 | pub shutdown_reason: T, 545 | } 546 | 547 | impl ShutdownAlreadyCompleted { 548 | pub(crate) const fn new(shutdown_reason: T) -> Self { 549 | Self { shutdown_reason } 550 | } 551 | } 552 | 553 | impl std::error::Error for ShutdownAlreadyCompleted {} 554 | 555 | impl std::fmt::Display for ShutdownAlreadyCompleted { 556 | fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { 557 | write!(f, "shutdown has already completed, can not delay shutdown completion") 558 | } 559 | } 560 | -------------------------------------------------------------------------------- /src/shutdown_complete.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::sync::{Arc, Mutex}; 4 | use std::task::{Context, Poll}; 5 | 6 | use crate::waker_list::WakerToken; 7 | use crate::ShutdownManagerInner; 8 | 9 | /// Future to wait for a shutdown to complete. 10 | pub struct ShutdownComplete { 11 | pub(crate) inner: Arc>>, 12 | pub(crate) waker_token: Option, 13 | } 14 | 15 | impl Clone for ShutdownComplete { 16 | fn clone(&self) -> Self { 17 | // Clone only the reference to the shutdown manager, not the waker token. 18 | // The waker token is personal to each future. 19 | Self { 20 | inner: self.inner.clone(), 21 | waker_token: None, 22 | } 23 | } 24 | } 25 | 26 | impl Drop for ShutdownComplete { 27 | fn drop(&mut self) { 28 | if let Some(token) = self.waker_token.take() { 29 | let mut inner = self.inner.lock().unwrap(); 30 | inner.on_shutdown_complete.deregister(token); 31 | } 32 | } 33 | } 34 | 35 | impl Future for ShutdownComplete { 36 | type Output = T; 37 | 38 | #[inline] 39 | fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll { 40 | let me = self.get_mut(); 41 | let mut inner = me.inner.lock().unwrap(); 42 | 43 | // We're being polled, so we should deregister the waker (if any). 44 | if let Some(token) = me.waker_token.take() { 45 | inner.on_shutdown_complete.deregister(token); 46 | } 47 | 48 | // Check if the shutdown is completed. 49 | if inner.delay_tokens == 0 { 50 | if let Some(reason) = inner.shutdown_reason.clone() { 51 | return Poll::Ready(reason); 52 | } 53 | } 54 | 55 | // We're not ready, so register the waker to wake us on shutdown completion. 56 | me.waker_token = Some(inner.on_shutdown_complete.register(context.waker().clone())); 57 | 58 | Poll::Pending 59 | } 60 | } 61 | 62 | #[cfg(test)] 63 | mod test { 64 | use assert2::assert; 65 | use std::future::Future; 66 | use std::pin::Pin; 67 | use std::task::Poll; 68 | 69 | /// Wrapper around a future to poll it only once. 70 | struct PollOnce<'a, F>(&'a mut F); 71 | 72 | impl<'a, F: std::marker::Unpin + Future> Future for PollOnce<'a, F> { 73 | type Output = Poll; 74 | 75 | fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { 76 | Poll::Ready(Pin::new(&mut self.get_mut().0).poll(cx)) 77 | } 78 | } 79 | 80 | /// Poll a future once. 81 | async fn poll_once(future: &mut F) -> Poll { 82 | PollOnce(future).await 83 | } 84 | 85 | #[tokio::test] 86 | async fn waker_list_doesnt_grow_infinitely() { 87 | let shutdown = crate::ShutdownManager::<()>::new(); 88 | for i in 0..100_000 { 89 | let mut wait_shutdown_complete = shutdown.wait_shutdown_complete(); 90 | let task = tokio::spawn(async move { 91 | assert!(let Poll::Pending = poll_once(&mut wait_shutdown_complete).await); 92 | }); 93 | assert!(let Ok(()) = task.await, "task = {i}"); 94 | } 95 | 96 | // Since we wait for each task to complete before spawning another, 97 | // the total amount of waker slots used should be only 1. 98 | let inner = shutdown.inner.lock().unwrap(); 99 | assert!(inner.on_shutdown_complete.total_slots() == 1); 100 | assert!(inner.on_shutdown_complete.empty_slots() == 1); 101 | } 102 | 103 | #[tokio::test] 104 | async fn cloning_does_not_clone_waker_token() { 105 | let shutdown = crate::ShutdownManager::<()>::new(); 106 | 107 | let mut signal = shutdown.wait_shutdown_complete(); 108 | assert!(let None = &signal.waker_token); 109 | 110 | assert!(let Poll::Pending = poll_once(&mut signal).await); 111 | assert!(let Some(_) = &signal.waker_token); 112 | 113 | let mut cloned = signal.clone(); 114 | assert!(let None = &cloned.waker_token); 115 | assert!(let Some(_) = &signal.waker_token); 116 | 117 | assert!(let Poll::Pending = poll_once(&mut cloned).await); 118 | assert!(let Some(_) = &cloned.waker_token); 119 | assert!(let Some(_) = &signal.waker_token); 120 | 121 | { 122 | let inner = shutdown.inner.lock().unwrap(); 123 | assert!(inner.on_shutdown_complete.total_slots() == 2); 124 | assert!(inner.on_shutdown_complete.empty_slots() == 0); 125 | } 126 | 127 | { 128 | drop(signal); 129 | let inner = shutdown.inner.lock().unwrap(); 130 | assert!(inner.on_shutdown_complete.empty_slots() == 1); 131 | } 132 | 133 | { 134 | drop(cloned); 135 | let inner = shutdown.inner.lock().unwrap(); 136 | assert!(inner.on_shutdown_complete.empty_slots() == 2); 137 | } 138 | } 139 | } 140 | -------------------------------------------------------------------------------- /src/shutdown_signal.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::sync::{Arc, Mutex}; 4 | use std::task::{Context, Poll}; 5 | 6 | use crate::waker_list::WakerToken; 7 | use crate::{WrapCancel, ShutdownManagerInner}; 8 | 9 | /// A future to wait for a shutdown signal. 10 | /// 11 | /// The future completes when the associated [`ShutdownManager`][crate::ShutdownManager] triggers a shutdown. 12 | /// 13 | /// The shutdown signal can be cloned and sent between threads freely. 14 | pub struct ShutdownSignal { 15 | pub(crate) inner: Arc>>, 16 | pub(crate) waker_token: Option, 17 | } 18 | 19 | impl Clone for ShutdownSignal { 20 | fn clone(&self) -> Self { 21 | // Clone only the reference to the shutdown manager, not the waker token. 22 | // The waker token is personal to each future. 23 | Self { 24 | inner: self.inner.clone(), 25 | waker_token: None, 26 | } 27 | } 28 | } 29 | 30 | impl Drop for ShutdownSignal { 31 | fn drop(&mut self) { 32 | if let Some(token) = self.waker_token.take() { 33 | let mut inner = self.inner.lock().unwrap(); 34 | inner.on_shutdown.deregister(token); 35 | } 36 | } 37 | } 38 | 39 | impl ShutdownSignal { 40 | /// Wrap a future so that it is cancelled when a shutdown is triggered. 41 | /// 42 | /// The returned future completes with `Err(reason)` containing the shutdown reason if a shutdown is triggered, 43 | /// and with `Ok(x)` when the wrapped future completes. 44 | /// 45 | /// The wrapped future is dropped if the shutdown starts before the wrapped future completes. 46 | #[inline] 47 | pub fn wrap_cancel(&self, future: F) -> WrapCancel { 48 | WrapCancel { 49 | shutdown_signal: self.clone(), 50 | future: Ok(future), 51 | } 52 | } 53 | } 54 | 55 | impl Future for ShutdownSignal { 56 | type Output = T; 57 | 58 | #[inline] 59 | fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll { 60 | let me = self.get_mut(); 61 | let mut inner = me.inner.lock().unwrap(); 62 | 63 | // We're being polled, so we should deregister the waker (if any). 64 | if let Some(token) = me.waker_token.take() { 65 | inner.on_shutdown.deregister(token); 66 | } 67 | 68 | if let Some(reason) = inner.shutdown_reason.clone() { 69 | // Shutdown started, so we're ready. 70 | Poll::Ready(reason) 71 | } else { 72 | // We're not ready, so register the waker to wake us on shutdown start. 73 | me.waker_token = Some(inner.on_shutdown.register(context.waker().clone())); 74 | Poll::Pending 75 | } 76 | } 77 | } 78 | 79 | #[cfg(test)] 80 | mod test { 81 | use assert2::assert; 82 | use std::future::Future; 83 | use std::pin::Pin; 84 | use std::task::Poll; 85 | 86 | /// Wrapper around a future to poll it only once. 87 | struct PollOnce<'a, F>(&'a mut F); 88 | 89 | impl<'a, F: std::marker::Unpin + Future> Future for PollOnce<'a, F> { 90 | type Output = Poll; 91 | 92 | fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { 93 | Poll::Ready(Pin::new(&mut self.get_mut().0).poll(cx)) 94 | } 95 | } 96 | 97 | /// Poll a future once. 98 | async fn poll_once(future: &mut F) -> Poll { 99 | PollOnce(future).await 100 | } 101 | 102 | #[tokio::test] 103 | async fn waker_list_doesnt_grow_infinitely() { 104 | let shutdown = crate::ShutdownManager::<()>::new(); 105 | for i in 0..100_000 { 106 | let task = tokio::spawn(shutdown.wrap_cancel(async move { 107 | tokio::task::yield_now().await; 108 | })); 109 | assert!(let Ok(Ok(())) = task.await, "task = {i}"); 110 | } 111 | 112 | // Since we wait for each task to complete before spawning another, 113 | // the total amount of waker slots used should be only 1. 114 | let inner = shutdown.inner.lock().unwrap(); 115 | assert!(inner.on_shutdown.total_slots() == 1); 116 | assert!(inner.on_shutdown.empty_slots() == 1); 117 | } 118 | 119 | #[tokio::test] 120 | async fn cloning_does_not_clone_waker_token() { 121 | let shutdown = crate::ShutdownManager::<()>::new(); 122 | 123 | let mut signal = shutdown.wait_shutdown_triggered(); 124 | assert!(let None = &signal.waker_token); 125 | 126 | assert!(let Poll::Pending = poll_once(&mut signal).await); 127 | assert!(let Some(_) = &signal.waker_token); 128 | 129 | let mut cloned = signal.clone(); 130 | assert!(let None = &cloned.waker_token); 131 | assert!(let Some(_) = &signal.waker_token); 132 | 133 | assert!(let Poll::Pending = poll_once(&mut cloned).await); 134 | assert!(let Some(_) = &cloned.waker_token); 135 | assert!(let Some(_) = &signal.waker_token); 136 | 137 | { 138 | let inner = shutdown.inner.lock().unwrap(); 139 | assert!(inner.on_shutdown.total_slots() == 2); 140 | assert!(inner.on_shutdown.empty_slots() == 0); 141 | } 142 | 143 | { 144 | drop(signal); 145 | let inner = shutdown.inner.lock().unwrap(); 146 | assert!(inner.on_shutdown.empty_slots() == 1); 147 | } 148 | 149 | { 150 | drop(cloned); 151 | let inner = shutdown.inner.lock().unwrap(); 152 | assert!(inner.on_shutdown.empty_slots() == 2); 153 | } 154 | } 155 | } 156 | -------------------------------------------------------------------------------- /src/waker_list.rs: -------------------------------------------------------------------------------- 1 | use std::task::Waker; 2 | 3 | /// A list of wakers. 4 | #[derive(Debug, Default)] 5 | pub struct WakerList { 6 | /// The wakers (with possibly empty slots) 7 | wakers: Vec>, 8 | 9 | /// The empty slots in the list. 10 | empty_slots: Vec, 11 | 12 | /// The current epoch, increased whenever `wake_all` is called. 13 | epoch: usize, 14 | } 15 | 16 | pub struct WakerToken { 17 | epoch: usize, 18 | index: usize, 19 | } 20 | 21 | impl WakerList { 22 | /// Create a new empty list of wakers. 23 | pub fn new() -> Self { 24 | Self::default() 25 | } 26 | 27 | /// Register a waker to be woken up when `wake_all` is called. 28 | /// 29 | /// Returns a token that can be used to unregister the waker again. 30 | pub fn register(&mut self, waker: Waker) -> WakerToken { 31 | if let Some(index) = self.empty_slots.pop() { 32 | debug_assert!(self.wakers[index].is_none()); 33 | self.wakers[index] = Some(waker); 34 | self.token(index) 35 | } else { 36 | self.wakers.push(Some(waker)); 37 | self.token(self.wakers.len() - 1) 38 | } 39 | } 40 | 41 | /// Deregister a waker so it will not be woken up by `wake_all` any more. 42 | /// 43 | /// This should be called when a future that registered the waker is dropped, 44 | /// to prevent the list of wakers growing infinitely large. 45 | /// 46 | /// # Panic 47 | /// May panic now or later if you give this function a token from another [`WakerList`]. 48 | pub fn deregister(&mut self, token: WakerToken) -> Option { 49 | if self.epoch != token.epoch { 50 | None 51 | } else if let Some(waker) = self.wakers[token.index].take() { 52 | self.empty_slots.push(token.index); 53 | Some(waker) 54 | } else { 55 | None 56 | } 57 | } 58 | 59 | /// Wake all wakers, clear the list and increase the epoch. 60 | #[allow(clippy::manual_flatten)] // Ssssh. 61 | pub fn wake_all(&mut self) { 62 | for waker in &mut self.wakers { 63 | if let Some(waker) = waker.take() { 64 | waker.wake() 65 | } 66 | } 67 | self.wakers.clear(); 68 | self.empty_slots.clear(); 69 | self.epoch += 1; 70 | } 71 | 72 | /// Create a token for the current epoch with the given index. 73 | fn token(&self, index: usize) -> WakerToken { 74 | WakerToken { 75 | epoch: self.epoch, 76 | index, 77 | } 78 | } 79 | 80 | /// Get the total number of waker slots. 81 | /// 82 | /// This includes empty slots. 83 | #[cfg(test)] 84 | pub fn total_slots(&self) -> usize { 85 | self.wakers.len() 86 | } 87 | 88 | /// Get the number of empty waker slots. 89 | #[cfg(test)] 90 | pub fn empty_slots(&self) -> usize { 91 | self.empty_slots.len() 92 | } 93 | } 94 | -------------------------------------------------------------------------------- /src/wrap_cancel.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::task::{Context, Poll}; 4 | 5 | use crate::shutdown_signal::ShutdownSignal; 6 | 7 | /// Wrapped future that is automatically cancelled when a shutdown is triggered. 8 | /// 9 | /// If the wrapped future completes before the shutdown is triggered, 10 | /// the output of the original future is yielded as `Ok(value)`. 11 | /// 12 | /// If the shutdown is triggered before the wrapped future completes, 13 | /// the original future is dropped and the shutdown reason is yielded as `Err(shutdown_reason)`. 14 | #[must_use = "futures must be polled to make progress"] 15 | pub struct WrapCancel { 16 | pub(crate) shutdown_signal: ShutdownSignal, 17 | pub(crate) future: Result, 18 | } 19 | 20 | impl Future for WrapCancel { 21 | type Output = Result; 22 | 23 | #[inline] 24 | fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll { 25 | // SAFETY: We never move `future`, so we can not violate the requirements of `F`. 26 | // We do drop it, but that's allowed by `Pin`. 27 | let me = unsafe { self.get_unchecked_mut() }; 28 | 29 | match &mut me.future { 30 | Err(e) => return Poll::Ready(Err(e.clone())), 31 | Ok(future) => { 32 | let future = unsafe { Pin::new_unchecked(future) }; 33 | if let Poll::Ready(value) = future.poll(context) { 34 | return Poll::Ready(Ok(value)); 35 | } 36 | }, 37 | } 38 | 39 | // Otherwise check if the shutdown signal has been given. 40 | let shutdown = Pin::new(&mut me.shutdown_signal).poll(context); 41 | match shutdown { 42 | Poll::Ready(reason) => { 43 | me.future = Err(reason.clone()); 44 | Poll::Ready(Err(reason)) 45 | }, 46 | Poll::Pending => Poll::Pending, 47 | } 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /src/wrap_delay_shutdown.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::task::{Context, Poll}; 4 | 5 | use crate::DelayShutdownToken; 6 | 7 | /// Wrapped future that delays shutdown completion until it completes or until it is droppped. 8 | #[must_use = "futures must be polled to make progress"] 9 | pub struct WrapDelayShutdown { 10 | pub(crate) delay_token: Option>, 11 | pub(crate) future: F, 12 | } 13 | 14 | impl Future for WrapDelayShutdown { 15 | type Output = F::Output; 16 | 17 | #[inline] 18 | fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll { 19 | // SAFETY: We never move `future`, so we can not violate the requirements of `F`. 20 | unsafe { 21 | let me = self.get_unchecked_mut(); 22 | match Pin::new_unchecked(&mut me.future).poll(context) { 23 | Poll::Pending => Poll::Pending, 24 | Poll::Ready(value) => { 25 | me.delay_token = None; 26 | Poll::Ready(value) 27 | }, 28 | } 29 | } 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /src/wrap_trigger_shutdown.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::task::{Context, Poll}; 4 | 5 | use crate::TriggerShutdownToken; 6 | 7 | /// Wrapped future that triggers a shutdown when it completes or when it is dropped. 8 | #[must_use = "futures must be polled to make progress"] 9 | pub struct WrapTriggerShutdown { 10 | pub(crate) trigger_shutdown_token: Option>, 11 | pub(crate) future: F, 12 | } 13 | 14 | impl Future for WrapTriggerShutdown { 15 | type Output = F::Output; 16 | 17 | #[inline] 18 | fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll { 19 | // SAFETY: We never move `future`, so we can not violate the requirements of `F`. 20 | unsafe { 21 | let me = self.get_unchecked_mut(); 22 | match Pin::new_unchecked(&mut me.future).poll(context) { 23 | Poll::Pending => Poll::Pending, 24 | Poll::Ready(value) => { 25 | me.trigger_shutdown_token = None; 26 | Poll::Ready(value) 27 | }, 28 | } 29 | } 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /tests/shutdown.rs: -------------------------------------------------------------------------------- 1 | use assert2::{assert, let_assert}; 2 | use futures::future; 3 | use std::future::Future; 4 | use std::time::Duration; 5 | 6 | use async_shutdown::ShutdownManager; 7 | 8 | #[track_caller] 9 | fn test_timeout(test: impl Future) { 10 | let_assert!(Ok(runtime) = tokio::runtime::Runtime::new(), "failed to initialize tokio runtime"); 11 | runtime.block_on(async move { 12 | let test = tokio::time::timeout(Duration::from_millis(100), test); 13 | assert!(let Ok(()) = test.await, "test timed out"); 14 | }); 15 | } 16 | 17 | #[test] 18 | fn shutdown() { 19 | // Create a `ShutdownManager` manager and instantly trigger a shutdown, 20 | test_timeout(async { 21 | let shutdown = ShutdownManager::new(); 22 | assert!(let Ok(()) = shutdown.trigger_shutdown(1)); 23 | assert!(shutdown.wait_shutdown_triggered().await == 1); 24 | assert!(shutdown.wait_shutdown_complete().await == 1); 25 | }); 26 | 27 | // Trigger the shutdown from a task after a short sleep. 28 | test_timeout(async { 29 | let shutdown = ShutdownManager::new(); 30 | 31 | tokio::spawn({ 32 | let shutdown = shutdown.clone(); 33 | async move { 34 | tokio::time::sleep(Duration::from_millis(20)).await; 35 | assert!(let Ok(()) = shutdown.trigger_shutdown(2)); 36 | } 37 | }); 38 | 39 | assert!(shutdown.wait_shutdown_triggered().await == 2); 40 | assert!(shutdown.wait_shutdown_complete().await == 2); 41 | }); 42 | } 43 | 44 | #[test] 45 | fn shutdown_only_works_once() { 46 | let shutdown = ShutdownManager::new(); 47 | assert!(let Ok(()) = shutdown.trigger_shutdown("first")); 48 | assert!(let Err(async_shutdown::ShutdownAlreadyStarted { 49 | shutdown_reason: "first", 50 | ignored_reason: "second", 51 | .. 52 | }) = shutdown.trigger_shutdown("second")); 53 | } 54 | 55 | #[test] 56 | fn wrap_cancel() { 57 | // Spawn a never-completing future that is automatically dropped on shutdown. 58 | test_timeout(async { 59 | let shutdown = ShutdownManager::new(); 60 | let task = tokio::spawn(shutdown.wrap_cancel(future::pending::<()>())); 61 | assert!(let Ok(()) = shutdown.trigger_shutdown("goodbye!")); 62 | let_assert!(Ok(Err(reason)) = task.await); 63 | assert!(reason == "goodbye!"); 64 | }); 65 | 66 | // Same test, but use a ShutdownSignal to wrap the future. 67 | test_timeout(async { 68 | let shutdown = ShutdownManager::new(); 69 | let wait_shutdown_triggered = shutdown.wait_shutdown_triggered(); 70 | let task = tokio::spawn(wait_shutdown_triggered.wrap_cancel(future::pending::<()>())); 71 | assert!(let Ok(()) = shutdown.trigger_shutdown("wooh")); 72 | let_assert!(Ok(Err(reason)) = task.await); 73 | assert!(reason == "wooh"); 74 | }); 75 | } 76 | 77 | #[test] 78 | fn wrap_cancel_no_shutdown() { 79 | // Spawn an already ready future and verify that it can complete if no shutdown happens. 80 | test_timeout(async { 81 | let shutdown = ShutdownManager::<()>::new(); 82 | let task = tokio::spawn(shutdown.wrap_cancel(future::ready(10))); 83 | assert!(let Ok(Ok(10)) = task.await); 84 | }); 85 | 86 | // Same test, but use a future that first sleeps and then completes. 87 | test_timeout(async { 88 | let shutdown = ShutdownManager::<()>::new(); 89 | let task = tokio::spawn(shutdown.wrap_cancel(async { 90 | tokio::time::sleep(Duration::from_millis(20)).await; 91 | 10u32 92 | })); 93 | assert!(let Ok(Ok(10)) = task.await); 94 | }); 95 | } 96 | 97 | #[test] 98 | fn delay_token() { 99 | // Delay shutdown completion using a delay-shutdown token. 100 | test_timeout(async { 101 | let shutdown = ShutdownManager::new(); 102 | 103 | // Create a token to delay shutdown completion. 104 | let_assert!(Ok(delay) = shutdown.delay_shutdown_token()); 105 | 106 | assert!(let Ok(()) = shutdown.trigger_shutdown(10)); 107 | shutdown.wait_shutdown_triggered().await; 108 | assert!(shutdown.is_shutdown_completed() == false); 109 | 110 | // Move the token into a task where it is dropped after a short sleep. 111 | tokio::spawn(async move { 112 | tokio::time::sleep(Duration::from_millis(20)).await; 113 | drop(delay); 114 | }); 115 | 116 | shutdown.wait_shutdown_complete().await; 117 | }); 118 | } 119 | 120 | #[test] 121 | fn wrap_delay() { 122 | // Spawn a future that delays shutdown as long as it is running. 123 | test_timeout(async { 124 | let shutdown = ShutdownManager::new(); 125 | let_assert!(Ok(delay) = shutdown.delay_shutdown_token()); 126 | assert!(let Ok(()) = shutdown.trigger_shutdown(10)); 127 | 128 | tokio::spawn(delay.wrap_future(async move { 129 | tokio::time::sleep(Duration::from_millis(10)).await; 130 | })); 131 | 132 | shutdown.wait_shutdown_triggered().await; 133 | shutdown.wait_shutdown_complete().await; 134 | }); 135 | } 136 | 137 | #[test] 138 | fn wait_for_shutdown_complete_in_task_without_delay_tokens() { 139 | test_timeout(async { 140 | let shutdown = ShutdownManager::new(); 141 | 142 | tokio::spawn({ 143 | let shutdown = shutdown.clone(); 144 | async move { 145 | tokio::time::sleep(Duration::from_millis(10)).await; 146 | assert!(let Ok(()) = shutdown.trigger_shutdown(10)); 147 | } 148 | }); 149 | 150 | let task = tokio::spawn({ 151 | let shutdown = shutdown.clone(); 152 | async move { 153 | assert!(shutdown.wait_shutdown_complete().await == 10); 154 | } 155 | }); 156 | 157 | assert!(let Ok(()) = task.await); 158 | }); 159 | } 160 | 161 | #[test] 162 | fn delay_token_too_late() { 163 | // Try go get a delay-shutdown token after the shutdown completed. 164 | let shutdown = ShutdownManager::new(); 165 | assert!(let Ok(()) = shutdown.trigger_shutdown(())); 166 | assert!(let Err(async_shutdown::ShutdownAlreadyCompleted { .. }) = shutdown.delay_shutdown_token()); 167 | assert!(let Err(async_shutdown::ShutdownAlreadyCompleted { .. }) = shutdown.wrap_delay_shutdown(future::pending::<()>())); 168 | } 169 | 170 | #[test] 171 | fn vital_token() { 172 | // Trigger a shutdown by dropping a token. 173 | test_timeout(async { 174 | let shutdown = ShutdownManager::new(); 175 | 176 | let trigger_shutdown_token = shutdown.trigger_shutdown_token("stop!"); 177 | drop(trigger_shutdown_token); 178 | 179 | assert!(shutdown.wait_shutdown_triggered().await == "stop!"); 180 | assert!(shutdown.wait_shutdown_complete().await == "stop!"); 181 | }); 182 | 183 | // Same test, but drop the vital token in a task. 184 | test_timeout(async { 185 | let shutdown = ShutdownManager::new(); 186 | 187 | let vital = shutdown.trigger_shutdown_token("done"); 188 | tokio::spawn(async move { 189 | tokio::time::sleep(Duration::from_millis(20)).await; 190 | drop(vital); 191 | }); 192 | 193 | assert!(shutdown.wait_shutdown_triggered().await == "done"); 194 | assert!(shutdown.wait_shutdown_complete().await == "done"); 195 | }); 196 | } 197 | 198 | #[test] 199 | fn wrap_vital() { 200 | // Trigger a shutdown by dropping a vital token from a task after a short sleep. 201 | test_timeout(async { 202 | let shutdown = ShutdownManager::new(); 203 | 204 | tokio::spawn(shutdown.wrap_trigger_shutdown("sleep done", async move { 205 | tokio::time::sleep(Duration::from_millis(20)).await; 206 | })); 207 | 208 | assert!(shutdown.wait_shutdown_triggered().await == "sleep done"); 209 | assert!(shutdown.wait_shutdown_complete().await == "sleep done"); 210 | }); 211 | 212 | // Same test, but now use a future that is instantly ready. 213 | test_timeout(async { 214 | let shutdown = ShutdownManager::new(); 215 | 216 | // Trigger the shutdown by dropping a vital token from a task after a short sleep. 217 | tokio::spawn(shutdown.wrap_trigger_shutdown("stop", future::ready(()))); 218 | 219 | assert!(shutdown.wait_shutdown_triggered().await == "stop"); 220 | assert!(shutdown.wait_shutdown_complete().await == "stop"); 221 | }); 222 | } 223 | --------------------------------------------------------------------------------