-
Notifications
You must be signed in to change notification settings - Fork 74
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Make server shutdown wait on any Detached handler futures #702
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -12,6 +12,7 @@ use super::router::HttpRouter; | |
use super::ProbeRegistration; | ||
|
||
use async_stream::stream; | ||
use debug_ignore::DebugIgnore; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. TIL, this crate seems useful! |
||
use futures::future::{ | ||
BoxFuture, FusedFuture, FutureExt, Shared, TryFutureExt, | ||
}; | ||
|
@@ -28,6 +29,7 @@ use hyper::Response; | |
use rustls; | ||
use std::convert::TryFrom; | ||
use std::future::Future; | ||
use std::mem; | ||
use std::net::SocketAddr; | ||
use std::num::NonZeroU32; | ||
use std::panic; | ||
|
@@ -39,6 +41,7 @@ use tokio::net::{TcpListener, TcpStream}; | |
use tokio::sync::oneshot; | ||
use tokio_rustls::{server::TlsStream, TlsAcceptor}; | ||
use uuid::Uuid; | ||
use waitgroup::WaitGroup; | ||
|
||
use crate::config::HandlerTaskMode; | ||
use crate::RequestInfo; | ||
|
@@ -69,6 +72,9 @@ pub struct DropshotState<C: ServerContext> { | |
pub local_addr: SocketAddr, | ||
/// Identifies how to accept TLS connections | ||
pub(crate) tls_acceptor: Option<Arc<Mutex<TlsAcceptor>>>, | ||
/// Worker for the handler_waitgroup associated with this server, allowing | ||
/// graceful shutdown to wait for all handlers to complete. | ||
pub(crate) handler_waitgroup_worker: DebugIgnore<waitgroup::Worker>, | ||
} | ||
|
||
impl<C: ServerContext> DropshotState<C> { | ||
|
@@ -96,6 +102,7 @@ pub struct HttpServerStarter<C: ServerContext> { | |
app_state: Arc<DropshotState<C>>, | ||
local_addr: SocketAddr, | ||
wrapped: WrappedHttpServerStarter<C>, | ||
handler_waitgroup: WaitGroup, | ||
} | ||
|
||
impl<C: ServerContext> HttpServerStarter<C> { | ||
|
@@ -123,6 +130,7 @@ impl<C: ServerContext> HttpServerStarter<C> { | |
default_handler_task_mode: config.default_handler_task_mode, | ||
}; | ||
|
||
let handler_waitgroup = WaitGroup::new(); | ||
let starter = match &tls { | ||
Some(tls) => { | ||
let (starter, app_state, local_addr) = | ||
|
@@ -133,11 +141,13 @@ impl<C: ServerContext> HttpServerStarter<C> { | |
private, | ||
log, | ||
tls, | ||
handler_waitgroup.worker(), | ||
)?; | ||
HttpServerStarter { | ||
app_state, | ||
local_addr, | ||
wrapped: WrappedHttpServerStarter::Https(starter), | ||
handler_waitgroup, | ||
} | ||
} | ||
None => { | ||
|
@@ -148,11 +158,13 @@ impl<C: ServerContext> HttpServerStarter<C> { | |
api, | ||
private, | ||
log, | ||
handler_waitgroup.worker(), | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is probably not something to worry about right now, but: I wonder how debuggable this is, either in situ or post mortem. Like if you walk up to a server that's shutting down and stuck waiting on one of these, would you have any way to know which request it was waiting on? I imagine eventually we'll want to elevate this to an API but that's probably way down the road. Just to show what I mean, in the past I built something like this: There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Probably not! Internally the waitgroup is just a wrapper around an AtomicWaker, which itself is just a glorified |
||
)?; | ||
HttpServerStarter { | ||
app_state, | ||
local_addr, | ||
wrapped: WrappedHttpServerStarter::Http(starter), | ||
handler_waitgroup, | ||
} | ||
} | ||
}; | ||
|
@@ -182,6 +194,15 @@ impl<C: ServerContext> HttpServerStarter<C> { | |
}); | ||
info!(self.app_state.log, "listening"); | ||
|
||
let handler_waitgroup = self.handler_waitgroup; | ||
let join_handle = async move { | ||
// After the server shuts down, we also want to wait for any | ||
// detached handler futures to complete. | ||
() = join_handle.await?; | ||
() = handler_waitgroup.wait().await; | ||
Ok(()) | ||
}; | ||
|
||
#[cfg(feature = "usdt-probes")] | ||
let probe_registration = match usdt::register_probes() { | ||
Ok(_) => { | ||
|
@@ -258,6 +279,7 @@ impl<C: ServerContext> InnerHttpServerStarter<C> { | |
api: ApiDescription<C>, | ||
private: C, | ||
log: &Logger, | ||
handler_waitgroup_worker: waitgroup::Worker, | ||
) -> Result<InnerHttpServerStarterNewReturn<C>, hyper::Error> { | ||
let incoming = AddrIncoming::bind(&config.bind_address)?; | ||
let local_addr = incoming.local_addr(); | ||
|
@@ -269,6 +291,7 @@ impl<C: ServerContext> InnerHttpServerStarter<C> { | |
log: log.new(o!("local_addr" => local_addr)), | ||
local_addr, | ||
tls_acceptor: None, | ||
handler_waitgroup_worker: DebugIgnore(handler_waitgroup_worker), | ||
}); | ||
|
||
let make_service = ServerConnectionHandler::new(app_state.clone()); | ||
|
@@ -546,6 +569,7 @@ impl<C: ServerContext> InnerHttpsServerStarter<C> { | |
private: C, | ||
log: &Logger, | ||
tls: &ConfigTls, | ||
handler_waitgroup_worker: waitgroup::Worker, | ||
) -> Result<InnerHttpsServerStarterNewReturn<C>, GenericError> { | ||
let acceptor = Arc::new(Mutex::new(TlsAcceptor::from(Arc::new( | ||
rustls::ServerConfig::try_from(tls)?, | ||
|
@@ -572,6 +596,7 @@ impl<C: ServerContext> InnerHttpsServerStarter<C> { | |
log: logger, | ||
local_addr, | ||
tls_acceptor: Some(acceptor), | ||
handler_waitgroup_worker: DebugIgnore(handler_waitgroup_worker), | ||
}); | ||
|
||
let make_service = ServerConnectionHandler::new(Arc::clone(&app_state)); | ||
|
@@ -689,6 +714,14 @@ impl<C: ServerContext> HttpServer<C> { | |
.expect("cannot close twice") | ||
.send(()) | ||
.expect("failed to send close signal"); | ||
|
||
// We _must_ explicitly drop our app state before awaiting join_future. | ||
// If we are running handlers in `Detached` mode, our `app_state` has a | ||
// `waitgroup::Worker` that they all clone, and `join_future` will await | ||
// all of them being dropped. That means we must drop our "primary" | ||
// clone of it, too! | ||
mem::drop(self.app_state); | ||
|
||
self.join_future.await | ||
} | ||
} | ||
|
@@ -875,6 +908,7 @@ async fn http_request_handle<C: ServerContext>( | |
// to completion. | ||
let (tx, rx) = oneshot::channel(); | ||
let request_log = rqctx.log.clone(); | ||
let worker = server.handler_waitgroup_worker.clone(); | ||
let handler_task = tokio::spawn(async move { | ||
let request_log = rqctx.log.clone(); | ||
let result = handler.handle_request(rqctx, request).await; | ||
|
@@ -887,6 +921,10 @@ async fn http_request_handle<C: ServerContext>( | |
"client disconnected before response returned" | ||
); | ||
} | ||
|
||
// Drop our waitgroup worker, allowing graceful shutdown to | ||
// complete (if it's waiting on us). | ||
mem::drop(worker); | ||
}); | ||
|
||
// The only way we can fail to receive on `rx` is if `tx` is | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is fine. I see this pulls in two deps. It's a shame if there's nothing in tokio or futures that can already do this. 🤷