Skip to content

Commit

Permalink
fix: possible race condition during graceful shutdown (#780)
Browse files Browse the repository at this point in the history
This PR just makes the behavior more reliable, also should address #661

The possible race condition:
- `notify_waiters()` will notify only already registered waiters

Thus, if it's called after
```rust
drop(acceptor);
if alive_connections.load(Ordering::Acquire) > 0 {
```
But before `notify.notified()` - it might be never notified.
  • Loading branch information
DDtKey authored Mar 30, 2024
1 parent c4137ed commit dedb994
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions poem/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ where
let server_graceful_shutdown_token = server_graceful_shutdown_token.clone();

tokio::spawn(async move {
let serve_connection = serve_connection(socket, local_addr, remote_addr, scheme, ep, server_graceful_shutdown_token, idle_timeout);
let serve_connection = serve_connection(socket, local_addr, remote_addr, scheme, ep, server_graceful_shutdown_token.clone(), idle_timeout);

if timeout.is_some() {
tokio::select! {
Expand All @@ -178,8 +178,11 @@ where
}

if alive_connections.fetch_sub(1, Ordering::Acquire) == 1 {
// We have to notify only if there is a registered waiter on shutdown
notify.notify_waiters();
// notify only if shutdown is initiated, to prevent notification when server is active.
// It's a valid state to have 0 alive connections when server is not shutting down.
if server_graceful_shutdown_token.is_cancelled() {
notify.notify_one();
}
}
});
}
Expand Down

0 comments on commit dedb994

Please sign in to comment.