7962: examples: add graceful shutdown example#73
7962: examples: add graceful shutdown example#73martin-augment wants to merge 1 commit intomasterfrom
Conversation
Add an example demonstrating graceful shutdown using `tokio::signal`. A TCP echo server listens for connections and shuts down cleanly on Ctrl+C (SIGINT), notifying active connections via a broadcast channel before exiting. This is a common production pattern but was not covered by any existing example.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a new example that illustrates the implementation of graceful shutdown in a Tokio-based application. The primary goal is to provide a clear, runnable demonstration of how to handle termination signals effectively, ensuring that server processes can shut down cleanly without abruptly terminating active client connections. This enhances the robustness and reliability of applications by preventing data loss or incomplete operations during shutdown sequences. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
WalkthroughA new example demonstrating graceful shutdown of a TCP echo server has been added. The example includes a Cargo.toml entry referencing a new graceful-shutdown.rs file. The implementation creates a tokio-based TCP server listening on 127.0.0.1:6142 that accepts client connections and echoes received data. It uses a broadcast channel to distribute shutdown signals to all active connections when the process receives an interrupt signal. Upon shutdown, the server stops accepting new connections, signals existing connections to terminate with a farewell message, and waits briefly for pending operations to complete before exiting. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
🤖 Augment PR SummarySummary: Adds a new Tokio example demonstrating graceful shutdown behavior using Changes:
🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| //! Graceful shutdown example using `tokio::signal`. | ||
| //! | ||
| //! This starts a TCP echo server that shuts down cleanly when it receives | ||
| //! Ctrl+C (SIGINT). In-flight connections are allowed to finish before |
There was a problem hiding this comment.
The docs say “In-flight connections are allowed to finish”, but on shutdown each connection breaks as soon as shutdown_rx.recv() completes, which interrupts active sessions rather than letting them complete naturally.
Severity: low
Other Locations
examples/graceful-shutdown.rs:60
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The example does not allow the clients to gracefully finish. It just waits for 100ms before dropping the Tokio runtime. The example should either collect all JoinResults and await them or use shutdown_tx.closed() to be notified when the subscribers are all doppped.
|
|
||
| let mut shutdown_rx = shutdown_tx.subscribe(); | ||
|
|
||
| tokio::spawn(async move { |
There was a problem hiding this comment.
The JoinHandle from tokio::spawn is dropped, so main can’t reliably know when all connection tasks have finished before returning; remaining tasks may be cancelled when the Tokio runtime shuts down.
Severity: medium
Other Locations
examples/graceful-shutdown.rs:84
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The example does not allow the clients to gracefully finish. It just waits for 100ms before dropping the Tokio runtime. The example should either collect all JoinResults and await them or use shutdown_tx.closed() to be notified when the subscribers are all doppped.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new example for graceful shutdown of a TCP server. However, the echo server logic is vulnerable to a Denial of Service (DoS) attack because it reads lines from clients into an unbounded buffer, which is a critical consideration for production use. Additionally, the implementation can be made more robust by replacing the fixed-duration sleep with a more reliable mechanism to wait for connections to close, and improving error handling within the connection handler.
| loop { | ||
| tokio::select! { | ||
| // Echo lines back to the client. | ||
| result = reader.read_line(&mut line) => { |
There was a problem hiding this comment.
The use of reader.read_line(&mut line) without a length limit is vulnerable to Denial of Service (DoS) attacks. A malicious client can send a very large amount of data without a newline character, causing the server to allocate an unbounded amount of memory for the line string, eventually leading to an Out-of-Memory (OOM) crash. In a production-ready server, you should use a method that enforces a maximum line length, such as using tokio_util::codec::LinesCodec or manually limiting the number of bytes read.
There was a problem hiding this comment.
value:good-to-have; category:bug; feedback: The Gemini Ai reviewer is correct! A bad client could cause denial of service attack by sending a huge line. But since this is an example of a graceful shutdown it is OK to leave it as is.
| // Drop the sender and wait briefly for connections to finish. | ||
| drop(shutdown_tx); | ||
| tokio::time::sleep(std::time::Duration::from_millis(100)).await; |
There was a problem hiding this comment.
Using a fixed-duration sleep for graceful shutdown is not robust. Connections might take longer than the specified duration to close, leading to abrupt termination. A more reliable approach is to wait for all connection handlers to actually finish. You can achieve this by waiting for all broadcast::Receiver handles to be dropped, which indicates that the connection handlers have terminated. The Sender::closed() method provides a future that completes when all receivers are gone.
// Wait for all active connections to finish processing.
// The `closed()` future will complete when all receivers have been dropped.
shutdown_tx.closed().await;There was a problem hiding this comment.
value:useful; category:bug; feedback: The Gemini AI reviewer is correct! The example does not allow the clients to gracefully finish. It just waits for 100ms before dropping the Tokio runtime. The example should either collect all JoinResults and await them or use shutdown_tx.closed() to be notified when the subscribers are all doppped.
| Ok(_) => { | ||
| let _ = writer.write_all(line.as_bytes()).await; | ||
| line.clear(); | ||
| } |
There was a problem hiding this comment.
The result of writer.write_all() is currently ignored. If an error occurs during the write operation (e.g., the client has disconnected), the loop will continue, potentially leading to more errors on subsequent operations. It's better to handle this error and break the loop to close the connection cleanly.
Ok(_) => {
if writer.write_all(line.as_bytes()).await.is_err() {
// Client disconnected.
break;
}
line.clear();
}There was a problem hiding this comment.
value:good-to-have; category:bug; feedback: The Gemini Ai reviewer is correct! A failed write indicates that the client is gone and the server should stop trying to write to it.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
|
|
||
| // Drop the sender and wait briefly for connections to finish. | ||
| drop(shutdown_tx); | ||
| tokio::time::sleep(std::time::Duration::from_millis(100)).await; |
There was a problem hiding this comment.
Spawned tasks not awaited, shutdown not actually graceful
Medium Severity
The documentation promises "in-flight connections are allowed to finish before the process exits," but the JoinHandles from tokio::spawn are discarded, and the code relies on an arbitrary 100ms sleep instead of actually awaiting task completion. When main returns, the Tokio runtime is dropped and any still-running tasks are cancelled. This means the example — whose entire purpose is to teach graceful shutdown — does not actually implement graceful shutdown. A JoinSet or collected JoinHandles awaited after signaling would be the correct approach.
Additional Locations (2)
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The example does not allow the clients to gracefully finish. It just waits for 100ms before dropping the Tokio runtime. The example should either collect all JoinResults and await them or use shutdown_tx.closed() to be notified when the subscribers are all doppped.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
examples/graceful-shutdown.rs (2)
36-38: Consider handling transient accept errors.Using
?onaccept()will terminate the server on any error, including transient ones likeEMFILE(too many open files). For a production server, you'd typically log and continue. For this example, the simplicity is appropriate.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/graceful-shutdown.rs` around lines 36 - 38, The accept branch currently uses result? which will abort on any accept error; change the handling of listener.accept() (the match arm with result = listener.accept() => { let (socket, addr) = result?; ... }) to explicitly match on result, logging transient errors and continuing the loop instead of returning early—i.e., if let Ok((socket, addr)) = result { ... } else { log::warn!("accept failed: {}", err); maybe backoff or continue; } so that transient errors (e.g., EMFILE) don't shut down the server.
42-68: Spawned tasks are not tracked.The spawned tasks'
JoinHandles are discarded, so there's no way to await their completion. The current 100ms sleep is a best-effort workaround. For a production pattern, consider storing handles in aJoinSetand awaiting them during shutdown.For an example demonstrating the broadcast-based shutdown signal pattern, this is acceptable. However, if this example is meant to show a complete production-ready pattern, tracking tasks would make it more robust.
♻️ Optional improvement using JoinSet
+use tokio::task::JoinSet; + #[tokio::main] async fn main() -> Result<(), Box<dyn Error>> { let listener = TcpListener::bind("127.0.0.1:6142").await?; println!("listening on 127.0.0.1:6142"); let (shutdown_tx, _) = broadcast::channel::<()>(1); + let mut tasks = JoinSet::new(); loop { tokio::select! { result = listener.accept() => { let (socket, addr) = result?; println!("accepted connection from {addr}"); let mut shutdown_rx = shutdown_tx.subscribe(); - tokio::spawn(async move { + tasks.spawn(async move { // ... task body unchanged ... }); } _ = tokio::signal::ctrl_c() => { println!("\nshutdown signal received, closing listener"); break; } } } let _ = shutdown_tx.send(()); drop(shutdown_tx); - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + // Wait for all connections to finish + while tasks.join_next().await.is_some() {} println!("shutdown complete"); Ok(()) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/graceful-shutdown.rs` around lines 42 - 68, The server spawns per-connection tasks with tokio::spawn and drops their JoinHandle, so shutdown only uses a 100ms sleep; change this by collecting connection task handles (or use tokio::task::JoinSet) and await them during shutdown: store each tokio::spawn() handle (from the block that uses reader.read_line, writer.write_all, shutdown_rx) into a JoinSet (or Vec<JoinHandle<()>>) when accepting connections and, after signaling shutdown, join/await the set so all spawned tasks finish before exit; ensure the shutdown path awaits the JoinSet.join_all/await loop rather than relying on sleep.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@examples/graceful-shutdown.rs`:
- Around line 36-38: The accept branch currently uses result? which will abort
on any accept error; change the handling of listener.accept() (the match arm
with result = listener.accept() => { let (socket, addr) = result?; ... }) to
explicitly match on result, logging transient errors and continuing the loop
instead of returning early—i.e., if let Ok((socket, addr)) = result { ... } else
{ log::warn!("accept failed: {}", err); maybe backoff or continue; } so that
transient errors (e.g., EMFILE) don't shut down the server.
- Around line 42-68: The server spawns per-connection tasks with tokio::spawn
and drops their JoinHandle, so shutdown only uses a 100ms sleep; change this by
collecting connection task handles (or use tokio::task::JoinSet) and await them
during shutdown: store each tokio::spawn() handle (from the block that uses
reader.read_line, writer.write_all, shutdown_rx) into a JoinSet (or
Vec<JoinHandle<()>>) when accepting connections and, after signaling shutdown,
join/await the set so all spawned tasks finish before exit; ensure the shutdown
path awaits the JoinSet.join_all/await loop rather than relying on sleep.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: c7950e29-b284-4d40-9dc9-dcb7bbe30407
📒 Files selected for processing (2)
examples/Cargo.tomlexamples/graceful-shutdown.rs
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The example does not allow the clients to gracefully finish. It just waits for 100ms before dropping the Tokio runtime. The example should either collect all JoinResults and await them or use shutdown_tx.closed() to be notified when the subscribers are all doppped. |
value:good-to-have; category:bug; feedback: The CodeRabbit Ai reviewer is correct! A failed accept() is not a reason to stop listening for new connections. But since this is an example of a graceful shutdown it is OK to return the error. |
Code Review: examples/graceful-shutdown.rsThis is a nice addition — a graceful shutdown example fills a real gap in the Tokio example collection. The overall structure and broadcast-channel approach are sound. A few issues need addressing before merging. Bug: Port conflict with
|
| Severity | Issue |
|---|---|
| Medium | Port 6142 conflicts with chat.rs |
| Medium | 100 ms sleep does not guarantee all tasks finish before exit |
| Medium | Write errors in the echo path are silently ignored |
| Minor | broadcast channel capacity deserves a comment |
| Nit | Cargo.toml ordering |
2 participants
7962: To review by AI