diff --git a/Cargo.toml b/Cargo.toml index 71ec6d4..8bc1a07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,7 +50,13 @@ http-body = "1" pin-project-lite = "0.2" # `serve_with`'s connections: HTTP/1.1 and HTTP/2 on one port, over TCP or # TLS. Already in the tree through axum and tonic. -hyper-util = { version = "0.1", features = ["server-auto", "service", "tokio"] } +hyper-util = { version = "0.1", features = ["server-auto", "server-graceful", "service", "tokio"] } +# The executor HTTP/2 connections spawn their stream tasks on: a +# `TaskTracker` a shutdown waits on, and a `CancellationToken` that ends +# them (0.7.16: `run_until_cancelled` checks the token first). Both already +# in the tree through hyper-util and tonic. +hyper = "1" +tokio-util = { version = "0.7.16", features = ["rt"] } # gRPC client (to upstream service). `tls-connect-info` gives a rustls server # stream the connection record (client certificates) a tonic handler reads, so @@ -180,6 +186,9 @@ tower = { version = "0.5", features = ["util", "limit"] } http-body-util = "0.1" # `hyper::upgrade::on`: a fallback that upgrades its connection (src/serve/tests.rs). hyper = "1" +# A raw HTTP/2 client that sees the GOAWAY of a graceful shutdown +# (tests/shutdown.rs); already in the tree through hyper. +h2 = "0.4" # Trailer frames for the hand-written upstream response in # tests/upstream_controls.rs (tonic's server API cannot set success trailers). http-body = "1" diff --git a/README.md b/README.md index fdcb346..3f2a9ef 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,8 @@ tonic services. through, or translated to gRPC for an upstream that speaks only gRPC - Built-in TLS and mTLS, and a cap on open connections; or your own TLS, with the client's address and certificate reaching the upstream +- Graceful shutdown: calls and streams in flight finish within a bounded + drain, HTTP/2 clients get a GOAWAY ([Shutting down](#shutting-down)) - Guards you scope to the traffic they cover (transcoded calls, the proxy's own endpoints, native gRPC, the fallback), rejecting in the protocol of the request ([Guards and scopes](#guards-and-scopes)): @@ -90,6 +92,9 @@ binary uses. Unset, `TOKIO_WORKER_THREADS` decides, else the number of CPUs available to the process. The startup log (`RUST_LOG=info`) shows the count and where it came from. +SIGTERM or Ctrl-C stops the binary gracefully: it stops accepting, lets the +calls in flight finish for up to `listen.drain_timeout_secs`, and exits 0. + ## Configuration ```yaml @@ -105,6 +110,10 @@ listen: idle_timeout_secs: 60 # Seconds an HTTP/1.1 client has to send a request's headers. header_read_timeout_secs: 30 + # Seconds a shutdown (SIGTERM, Ctrl-C) waits for requests and streams in + # flight to finish; the connections still open after it are closed. + # 0 waits for all of them. + drain_timeout_secs: 25 # Optional: TLS on the listener (REST and gRPC share the port; ALPN offers # h2 and http/1.1). With client_ca_file, client certificates are verified # (mTLS) and reach an in-process tonic upstream as Request::peer_certs. @@ -727,7 +736,7 @@ returns `application/jwk-set+json` (RFC 7517 §8.5). command-line dependencies come with the `cli` feature. The library runs on your tokio runtime and logs through `tracing` to the subscriber you set up. -```rust +```rust,no_run use std::path::Path; use structured_proxy::ProxyServer; @@ -738,8 +747,13 @@ async fn main() -> anyhow::Result<()> { // `ProxyConfig`. let server = ProxyServer::from_file(Path::new("my-service.yaml"))?; - // Run the proxy on the configured listen address. - server.serve().await?; + // Run the proxy on the configured listen address until Ctrl-C, then + // drain (see "Shutting down"). + server + .serve_with_shutdown(async { + tokio::signal::ctrl_c().await.ok(); + }) + .await?; Ok(()) } ``` @@ -858,6 +872,48 @@ listener verified reaches a tonic handler in process as `Request::peer_certs`. TLS needs a rustls crypto provider: the one a crypto backend feature brings, or the one your process installed (see [TLS crypto](#tls-crypto)). +### Shutting down + +`serve` and `serve_with` run until their future is dropped, which closes every +connection at once. For a graceful stop, hand `serve_with_shutdown` (or +`ProxyServer::serve_with_shutdown`) a future that completes when the process +should stop: + +```rust +use std::time::Duration; +use structured_proxy::{ProxyServer, ServeOptions}; + +# async fn run(grpc: tonic::service::Routes) -> anyhow::Result<()> { +let proxy = ProxyServer::from_file(std::path::Path::new("my-service.yaml"))?.service(grpc)?; +let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?; +let options = ServeOptions::new().drain_timeout(Some(Duration::from_secs(20))); +let stop = async { + tokio::signal::ctrl_c().await.ok(); +}; +structured_proxy::serve_with_shutdown(listener, proxy, options, stop).await?; +# Ok(()) +# } +``` + +Once the future completes: + +- the listening socket closes, so new connections are refused, and a + connection still in its TLS handshake or waiting for a `max_connections` + slot is dropped; +- HTTP/2 connections get a GOAWAY, so their clients open no new calls, and + HTTP/1.1 connections close after the response in progress; +- calls and streams in flight run to their end, and `serve_with_shutdown` + returns once every connection has closed; +- after `drain_timeout` (`listen.drain_timeout_secs`; `None` or 0 waits + without a bound) the connections still open are closed. + +Past its grace period an orchestrator kills the process along with its calls, +so keep the drain below it. The default of 25 s fits the 30 s Kubernetes gives +a pod after SIGTERM; with a longer `terminationGracePeriodSeconds` the drain +can grow with it. A connection a fallback upgraded (a +WebSocket) belongs to the fallback's task and closes when that task lets it +go. + ### Behind your own TLS For a server of your own (another TLS stack, a Unix socket), run the service @@ -941,7 +997,7 @@ a long server stream short. `ProxyServer::router` returns the proxy's HTTP routes in front of the configured upstream address, to serve or to merge into your own axum `Router`: -```rust +```rust,no_run use std::path::Path; use structured_proxy::{config::ProxyConfig, ProxyServer}; @@ -1033,6 +1089,7 @@ single-backend while `jsonwebtoken` sees two. Settle it once at the top of ```rust # fn main() { +# #[cfg(feature = "builtin_jwt")] structured_proxy::install_default_crypto_provider(); # } ``` @@ -1151,7 +1208,7 @@ verifies certificates with its own patched `rustls-webpki`. At startup the proxy reads your proto descriptors and turns every `google.api.http` rule into a REST route. Each request is then sorted once: -``` +```text REST, gRPC and gRPC-Web clients (HTTP/1.1, HTTP/2, optional TLS) │ ┌──────────────▼──────────────┐ diff --git a/packaging/config.yaml b/packaging/config.yaml index e8ca3c5..f739d1e 100644 --- a/packaging/config.yaml +++ b/packaging/config.yaml @@ -21,6 +21,10 @@ # HTTP listen address for the transcoded REST surface. listen: http: "0.0.0.0:8080" + # Seconds a stop (systemctl stop, SIGTERM) waits for calls in flight to + # finish before closing their connections. Keep it below TimeoutStopSec of + # the unit (30). + # drain_timeout_secs: 25 # Upstream gRPC service this proxy transcodes to. REQUIRED — replace with your # service address. diff --git a/packaging/structured-proxy.service b/packaging/structured-proxy.service index 63ae0bc..bcd29b8 100644 --- a/packaging/structured-proxy.service +++ b/packaging/structured-proxy.service @@ -19,6 +19,9 @@ Environment=RUST_LOG=info Restart=on-failure RestartSec=5 LimitNOFILE=65536 +# On SIGTERM the proxy drains for up to listen.drain_timeout_secs (25 by +# default), then exits; keep this above it, so systemd does not kill calls the +# drain would still finish. TimeoutStopSec=30 StandardOutput=journal StandardError=journal diff --git a/src/auth/jwks/tests.rs b/src/auth/jwks/tests.rs index 81e25a3..e16b5fd 100644 --- a/src/auth/jwks/tests.rs +++ b/src/auth/jwks/tests.rs @@ -207,7 +207,10 @@ async fn a_lookup_during_a_refresh_waits_for_its_keys() { // refresh is in flight when a second lookup arrives: the second must not // answer from the aged set in the meantime. let (endpoint, uri) = endpoint().await; - let interval = Duration::from_millis(50); + // The second lookup checks the throttle only once the held refresh ends, + // against when that refresh started. The interval must outlast that on a + // loaded machine, or the second lookup may rightly refresh a third time. + let interval = Duration::from_millis(500); let cache = Arc::new( JwksCache::new(uri) .unwrap() diff --git a/src/config.rs b/src/config.rs index f0aa15c..65dec8e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -479,6 +479,15 @@ pub struct ListenConfig { /// 1. Default: 30. #[serde(default = "default_header_read_timeout_secs")] pub header_read_timeout_secs: u64, + /// Seconds a graceful shutdown waits for open connections to finish what + /// they serve; the ones still open after it are closed. 0 waits for all. + /// Default: 25, below the 30 s grace period Kubernetes gives by default. + #[serde(default = "default_drain_timeout_secs")] + pub drain_timeout_secs: u64, +} + +fn default_drain_timeout_secs() -> u64 { + 25 } fn default_idle_timeout_secs() -> u64 { @@ -547,6 +556,7 @@ impl Default for ListenConfig { tls: None, idle_timeout_secs: default_idle_timeout_secs(), header_read_timeout_secs: default_header_read_timeout_secs(), + drain_timeout_secs: default_drain_timeout_secs(), } } } diff --git a/src/lib.rs b/src/lib.rs index 1d9757b..9607bff 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -53,6 +53,11 @@ compile_error!( (or neither, and inject a verifier with `ProxyServer::with_token_verifier`)" ); +/// The README's Rust examples, compiled as doc tests. +#[cfg(doctest)] +#[doc = include_str!("../README.md")] +struct ReadmeDoctests; + pub mod auth; pub mod config; mod cors; @@ -73,7 +78,7 @@ pub mod upstream; /// [`install_default_crypto_provider`] for when a call is needed. #[cfg(feature = "builtin_jwt")] pub use auth::crypto::install_default_crypto_provider; -pub use serve::{serve, serve_with, ServeOptions}; +pub use serve::{serve, serve_with, serve_with_shutdown, ServeOptions}; pub use service::{ConnectionInfo, ProxyService}; use axum::extract::State; @@ -1152,7 +1157,11 @@ impl ProxyServer { (listen.idle_timeout_secs > 0) .then(|| Duration::from_secs(listen.idle_timeout_secs)), ) - .header_read_timeout(Duration::from_secs(listen.header_read_timeout_secs)); + .header_read_timeout(Duration::from_secs(listen.header_read_timeout_secs)) + .drain_timeout( + (listen.drain_timeout_secs > 0) + .then(|| Duration::from_secs(listen.drain_timeout_secs)), + ); if let Some(max) = listen.max_connections { anyhow::ensure!(max > 0, "listen.max_connections must be at least 1"); options = options.max_connections(max); @@ -1183,6 +1192,31 @@ impl ProxyServer { /// [`serve_options`](Self::serve_options) reject, an invalid listen /// address, or a listener that fails. pub async fn serve(&self) -> anyhow::Result<()> { + self.serve_with_shutdown(std::future::pending()).await + } + + /// [`serve`](Self::serve) until `signal` completes, then shut down with + /// the drain of `listen.drain_timeout_secs` (see [`serve_with_shutdown`]). + /// + /// # Errors + /// + /// What [`serve`](Self::serve) returns. + /// + /// # Examples + /// + /// ```no_run + /// # async fn run(server: structured_proxy::ProxyServer) -> anyhow::Result<()> { + /// server + /// .serve_with_shutdown(async { + /// tokio::signal::ctrl_c().await.ok(); + /// }) + /// .await + /// # } + /// ``` + pub async fn serve_with_shutdown( + &self, + signal: impl std::future::Future, + ) -> anyhow::Result<()> { let service = self.service(self.upstream()?)?; let options = self.serve_options()?; let addr: SocketAddr = self.config.listen.http.parse()?; @@ -1194,7 +1228,8 @@ impl ProxyServer { self.config.service.name, addr ); - serve_with(listener, service, options).await?; + serve_with_shutdown(listener, service, options, signal).await?; + tracing::info!("{} stopped", self.config.service.name); Ok(()) } } diff --git a/src/main.rs b/src/main.rs index c7ea721..e0258f1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -52,5 +52,38 @@ fn main() -> anyhow::Result<()> { "Starting structured-proxy" ); - rt.block_on(server.serve()) + rt.block_on(server.serve_with_shutdown(stop_requested())) +} + +/// Completes on Ctrl-C, and on SIGTERM on Unix (what container runtimes and +/// service managers send to stop a process). +async fn stop_requested() { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + match signal(SignalKind::terminate()) { + Ok(mut term) => { + tokio::select! { + () = ctrl_c() => {} + _ = term.recv() => {} + } + } + Err(e) => { + tracing::warn!(error = %e, "cannot listen for SIGTERM; stopping on Ctrl-C only"); + ctrl_c().await; + } + } + } + #[cfg(not(unix))] + ctrl_c().await; + tracing::info!("stop requested; draining connections"); +} + +/// Completes on Ctrl-C. A handler that cannot be installed never completes: +/// failing to listen for the signal is no reason to stop serving. +async fn ctrl_c() { + if let Err(e) = tokio::signal::ctrl_c().await { + tracing::warn!(error = %e, "cannot listen for Ctrl-C"); + std::future::pending::<()>().await; + } } diff --git a/src/serve.rs b/src/serve.rs index 728658b..b9de352 100644 --- a/src/serve.rs +++ b/src/serve.rs @@ -1,18 +1,22 @@ //! Running a [`ProxyService`] on a TCP listener: HTTP/1.1 and HTTP/2 on one -//! port, optionally behind TLS, with an optional cap on open connections. +//! port, optionally behind TLS, with an optional cap on open connections and +//! a graceful shutdown. +use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Duration; -use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer}; +use hyper_util::rt::{TokioIo, TokioTimer}; use hyper_util::server::conn::auto::Builder; +use hyper_util::server::graceful::GracefulConnection; use hyper_util::service::TowerToHyperService; use pin_project_lite::pin_project; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::net::TcpListener; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore}; +use tokio::task::JoinSet; use tonic::transport::server::Connected; use crate::service::{ConnectionInfo, ProxyService}; @@ -40,6 +44,7 @@ pub struct ServeOptions { idle_timeout: Option, header_read_timeout: Duration, tls_handshake_timeout: Duration, + drain_timeout: Option, } impl Default for ServeOptions { @@ -50,6 +55,9 @@ impl Default for ServeOptions { idle_timeout: Some(Duration::from_secs(60)), header_read_timeout: Duration::from_secs(30), tls_handshake_timeout: Duration::from_secs(10), + // Below the 30 s a Kubernetes pod gets after SIGTERM by default, + // so the drain ends before the kill. + drain_timeout: Some(Duration::from_secs(25)), } } } @@ -57,11 +65,21 @@ impl Default for ServeOptions { impl ServeOptions { /// Cleartext, with no limit on connections; a connection idle for 60 s is /// closed, a client gets 30 s to send the headers of an HTTP/1.1 request - /// and 10 s to finish a TLS handshake. + /// and 10 s to finish a TLS handshake, and a shutdown waits at most 25 s + /// for open connections. pub fn new() -> Self { Self::default() } + /// How long a shutdown ([`serve_with_shutdown`]) waits for the + /// connections still open to finish what they serve; the ones open after + /// it are closed. `None` waits for all of them. + #[must_use] + pub fn drain_timeout(mut self, timeout: Option) -> Self { + self.drain_timeout = timeout; + self + } + /// Close a connection that has had no request in flight for `timeout` /// (gracefully: HTTP/2 gets a GOAWAY); `None` keeps idle connections /// open. A request counts until its response body ends, so a stream keeps @@ -120,9 +138,10 @@ impl ServeOptions { } } -/// Serve `service` on `listener` until the listener fails: cleartext HTTP/1.1 -/// and HTTP/2 on the same port, so REST clients and native gRPC clients share -/// it. [`serve_with`] adds TLS and a connection limit. +/// Serve `service` on `listener`: cleartext HTTP/1.1 and HTTP/2 on the same +/// port, so REST clients and native gRPC clients share it. It runs until the +/// future is dropped, which closes every connection; [`serve_with_shutdown`] +/// stops gracefully, and [`serve_with`] adds TLS and a connection limit. /// /// # Errors /// @@ -178,6 +197,52 @@ pub async fn serve_with( service: ProxyService, options: ServeOptions, ) -> std::io::Result<()> { + serve_with_shutdown(listener, service, options, std::future::pending()).await +} + +/// [`serve_with`], until `signal` resolves; then a graceful shutdown. The +/// listening socket closes, so new connections are refused; a connection +/// still in its TLS handshake or waiting for its `max_connections` slot is +/// dropped; every open connection is asked to wind down (HTTP/2 gets a GOAWAY, +/// so its client opens no new streams; HTTP/1.1 closes after the response in +/// progress), and requests and streams in flight finish. The future resolves +/// once every connection has closed, or after +/// [`drain_timeout`](ServeOptions::drain_timeout), closing the ones still +/// open. Dropping it closes every connection at once. +/// +/// A connection a fallback upgraded (a WebSocket) belongs to the fallback's +/// own task once upgraded, and closes when that task lets it go. +/// +/// # Errors +/// +/// See [`serve`]. +/// +/// # Examples +/// +/// ```no_run +/// use structured_proxy::{ProxyServer, ServeOptions}; +/// +/// # async fn run() -> anyhow::Result<()> { +/// let service = ProxyServer::from_yaml_str("service:\n name: demo\n")? +/// .service(tonic::service::Routes::default())?; +/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?; +/// let shutdown = async { +/// tokio::signal::ctrl_c().await.expect("the signal handler installs"); +/// }; +/// structured_proxy::serve_with_shutdown(listener, service, ServeOptions::new(), shutdown).await?; +/// # Ok(()) +/// # } +/// ``` +pub async fn serve_with_shutdown( + listener: TcpListener, + service: ProxyService, + options: ServeOptions, + signal: F, +) -> std::io::Result<()> +where + U: Upstream, + F: Future, +{ let slots = options .max_connections .map(|max| Arc::new(Semaphore::new(max))); @@ -187,30 +252,47 @@ pub async fn serve_with( idle_timeout: options.idle_timeout, header_read_timeout: options.header_read_timeout, }; - loop { + // Tells every connection to wind down. The connections are this future's + // own tasks, so dropping it ends them too. + let (stop, stopping) = watch::channel(false); + let mut connections = JoinSet::new(); + let streams = StreamExecutor::default(); + let _streams_end_with_this_future = streams.cancel_on_drop(); + tokio::pin!(signal); + // Every wait below also reaps finished connections, so the set holds only + // open ones even while no client arrives. + 'serve: loop { // The slot is taken before the accept, so a full server leaves new // connections in the kernel's backlog instead of accepting and // dropping them. - let slot = match &slots { - Some(slots) => Some( - slots - .clone() - .acquire_owned() - .await - .expect("the connection semaphore is never closed"), - ), - None => None, + let slot = loop { + tokio::select! { + biased; + () = &mut signal => break 'serve, + Some(_) = connections.join_next() => {} + slot = take_slot(slots.as_ref()) => break slot, + } }; - let tcp = match listener.accept().await { - Ok((tcp, _)) => tcp, - Err(error) => { - accept_failed(error).await; - continue; + let tcp = loop { + tokio::select! { + biased; + () = &mut signal => break 'serve, + Some(_) = connections.join_next() => {} + accepted = listener.accept() => match accepted { + Ok((tcp, _)) => break tcp, + Err(error) => tokio::select! { + biased; + () = &mut signal => break 'serve, + () = accept_failed(error) => {} + }, + }, } }; let service = service.clone(); let acceptor = acceptor.clone(); - tokio::spawn(async move { + let stopping = stopping.clone(); + let streams = streams.clone(); + connections.spawn(async move { // Small gRPC frames and REST answers are latency-bound. if let Err(error) = tcp.set_nodelay(true) { tracing::debug!(%error, "cannot set TCP_NODELAY"); @@ -218,13 +300,16 @@ pub async fn serve_with( match acceptor { None => { let service = service.for_connection(tcp.connect_info()); - serve_connection(SlotIo { io: tcp, slot }, service, limits).await; + let io = SlotIo { io: tcp, slot }; + serve_connection(io, service, limits, streams, stopping).await; } Some(acceptor) => { // The slot is held through the handshake, then by the - // stream. - let stream = - match tokio::time::timeout(handshake_timeout, acceptor.accept(tcp)).await { + // stream. A handshake still running at shutdown is not + // finished: the connection would only be closed again. + let handshake = tokio::time::timeout(handshake_timeout, acceptor.accept(tcp)); + let stream = tokio::select! { + handshake = handshake => match handshake { Ok(Ok(stream)) => stream, Ok(Err(error)) => { tracing::debug!(%error, "TLS handshake failed"); @@ -234,14 +319,107 @@ pub async fn serve_with( tracing::debug!("TLS handshake timed out"); return; } - }; + }, + () = stopped(stopping.clone()) => return, + }; let service = service.for_connection(ConnectionInfo::tls(stream.connect_info())); - serve_connection(SlotIo { io: stream, slot }, service, limits).await; + let io = SlotIo { io: stream, slot }; + serve_connection(io, service, limits, streams, stopping).await; } } }); } + // No new connections from here: the listening socket closes, then the + // open connections wind down. + drop(listener); + stop.send_replace(true); + streams.close(); + let drained = async { + while connections.join_next().await.is_some() {} + streams.finished().await; + }; + match options.drain_timeout { + None => drained.await, + Some(timeout) => { + if tokio::time::timeout(timeout, drained).await.is_err() { + tracing::warn!( + open = connections.len(), + "shutdown drain timed out; closing the connections still open" + ); + connections.shutdown().await; + streams.stop().await; + } + } + } + Ok(()) +} + +/// Spawns the stream tasks of HTTP/2 connections (hyper runs each stream's +/// service in a task of its own) under one tracker, so a shutdown can wait +/// for them as it does for the connections. +/// +/// Each task also ends when [`stop`](Self::stop) cancels them: a closed +/// connection resets its streams, but a stream ending on that reset is +/// hyper's behaviour, not something a bounded shutdown should rest on. +#[derive(Clone, Default)] +struct StreamExecutor { + tasks: tokio_util::task::TaskTracker, + cancel: tokio_util::sync::CancellationToken, +} + +impl StreamExecutor { + /// No new stream after this; the ones running may finish. + fn close(&self) { + self.tasks.close(); + } + + /// Resolves once every stream task has finished on its own. + async fn finished(&self) { + self.tasks.wait().await; + } + + /// Ends every stream task still running and waits for them. + async fn stop(&self) { + self.cancel.cancel(); + self.tasks.wait().await; + } + + /// Cancels every stream task when dropped, so they end with the future + /// that owns the server even if it is dropped mid-shutdown. + fn cancel_on_drop(&self) -> tokio_util::sync::DropGuard { + self.cancel.clone().drop_guard() + } +} + +impl hyper::rt::Executor for StreamExecutor +where + F: Future + Send + 'static, + F::Output: Send + 'static, +{ + fn execute(&self, future: F) { + self.tasks + .spawn(self.cancel.clone().run_until_cancelled_owned(future)); + } +} + +/// A `max_connections` slot, or none without a limit. +async fn take_slot(slots: Option<&Arc>) -> Option { + match slots { + Some(slots) => Some( + Arc::clone(slots) + .acquire_owned() + .await + .expect("the connection semaphore is never closed"), + ), + None => None, + } +} + +/// Resolves once the server shuts down. +async fn stopped(mut stopping: watch::Receiver) { + // A closed channel means the serve future is gone: stop as well. + stopping.wait_for(|stop| *stop).await.ok(); } pin_project! { @@ -309,13 +487,19 @@ struct ConnectionLimits { header_read_timeout: Duration, } -/// HTTP/1.1 or HTTP/2, whichever the client speaks, on one connection. -async fn serve_connection(io: I, service: ProxyService, limits: ConnectionLimits) -where +/// HTTP/1.1 or HTTP/2, whichever the client speaks, on one connection, until +/// it closes, idles out or the server shuts down. +async fn serve_connection( + io: I, + service: ProxyService, + limits: ConnectionLimits, + streams: StreamExecutor, + stopping: watch::Receiver, +) where U: Upstream, I: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - let mut builder = Builder::new(TokioExecutor::new()); + let mut builder = Builder::new(streams); // hyper times nothing without a timer: its own default header read // timeout is dropped with a warning. builder @@ -325,9 +509,11 @@ where builder.http2().timer(TokioTimer::new()); let served = match limits.idle_timeout { None => { - builder - .serve_connection_with_upgrades(TokioIo::new(io), TowerToHyperService::new(service)) - .await + let connection = builder.serve_connection_with_upgrades( + TokioIo::new(io), + TowerToHyperService::new(service), + ); + drive(connection, std::future::pending(), stopping).await } Some(timeout) => { let activity = Arc::new(idle::Activity::default()); @@ -339,16 +525,7 @@ where TokioIo::new(io), TowerToHyperService::new(service), ); - tokio::pin!(connection); - tokio::select! { - served = connection.as_mut() => served, - () = activity.idle_for(timeout) => { - // HTTP/2 gets a GOAWAY, HTTP/1.1 closes after the request - // it is reading, if any. - connection.as_mut().graceful_shutdown(); - connection.await - } - } + drive(connection, activity.idle_for(timeout), stopping).await } }; if let Err(error) = served { @@ -356,6 +533,28 @@ where } } +/// Serve `connection` until it ends; when `idle` resolves or the server +/// shuts down first, it is asked to wind down (HTTP/2 gets a GOAWAY, HTTP/1.1 +/// closes after the request in progress) and served until it has. +async fn drive( + connection: C, + idle: impl Future, + stopping: watch::Receiver, +) -> Result<(), C::Error> { + tokio::pin!(connection); + tokio::select! { + served = connection.as_mut() => served, + () = idle => { + connection.as_mut().graceful_shutdown(); + connection.await + } + () = stopped(stopping) => { + connection.as_mut().graceful_shutdown(); + connection.await + } + } +} + /// A failed accept: a connection the client gave up on is skipped; anything /// else (a full file descriptor table) is logged and retried a second later, /// the pause that lets open connections close. diff --git a/src/serve/tests.rs b/src/serve/tests.rs index 4c5587b..7ad44cb 100644 --- a/src/serve/tests.rs +++ b/src/serve/tests.rs @@ -185,6 +185,24 @@ async fn without_a_limit_connections_are_served_together() { assert_eq!(status_line(&mut first).await, "HTTP/1.1 200 OK"); } +#[tokio::test] +async fn stopping_the_streams_ends_one_that_never_finishes() { + // A stream task that does not end when its connection closes (a service + // or body that never wakes) must still not hold the shutdown open. + let streams = StreamExecutor::default(); + let (dropped, gone) = tokio::sync::oneshot::channel::<()>(); + hyper::rt::Executor::execute(&streams, async move { + let _dropped = dropped; + std::future::pending::<()>().await; + }); + streams.close(); + tokio::time::timeout(Duration::from_secs(5), streams.stop()) + .await + .expect("stopping ends the stream task"); + // The sender went with the task's future: the future was dropped. + assert!(gone.await.is_err()); +} + #[test] #[should_panic(expected = "max_connections must be at least 1")] fn a_limit_of_zero_is_refused() { @@ -194,16 +212,23 @@ fn a_limit_of_zero_is_refused() { #[test] fn connection_timeouts_come_from_the_config() { let server = crate::ProxyServer::from_yaml_str( - "listen:\n idle_timeout_secs: 0\n header_read_timeout_secs: 5\n", + "listen:\n idle_timeout_secs: 0\n header_read_timeout_secs: 5\n drain_timeout_secs: 0\n", ) .unwrap(); let options = server.serve_options().unwrap(); - // 0 keeps idle connections open. + // 0 keeps idle connections open, and a shutdown waits for every one. assert_eq!(options.idle_timeout, None); assert_eq!(options.header_read_timeout, Duration::from_secs(5)); + assert_eq!(options.drain_timeout, None); let defaults = crate::ProxyServer::new().serve_options().unwrap(); assert_eq!(defaults.idle_timeout, Some(Duration::from_secs(60))); assert_eq!(defaults.header_read_timeout, Duration::from_secs(30)); + assert_eq!(defaults.drain_timeout, Some(Duration::from_secs(25))); + let set = crate::ProxyServer::from_yaml_str("listen:\n drain_timeout_secs: 7\n") + .unwrap() + .serve_options() + .unwrap(); + assert_eq!(set.drain_timeout, Some(Duration::from_secs(7))); } #[test] diff --git a/tests/cli.rs b/tests/cli.rs index 4b1268e..2e7513f 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -180,6 +180,48 @@ fn the_binary_keeps_the_shared_rate_limit_store() { assert!(!output.contains("not compiled in"), "{output}"); } +#[cfg(unix)] +#[test] +fn sigterm_stops_the_proxy_cleanly() { + // What a container runtime or service manager sends to stop it. + let port = free_port(); + let config = ConfigFile::new( + "sigterm", + &format!( + "listen:\n http: \"127.0.0.1:{port}\"\n\ + upstream:\n default: \"http://127.0.0.1:9\"\n\ + descriptors: []\n" + ), + ); + let mut proxy = Running( + Command::new(BIN) + .arg("--config") + .arg(&config.0) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(), + ); + assert!(live(port).starts_with("HTTP/1.1 200")); + + let sent = Command::new("kill") + .arg("-TERM") + .arg(proxy.0.id().to_string()) + .status() + .unwrap(); + assert!(sent.success()); + let deadline = Instant::now() + Duration::from_secs(20); + let status = loop { + if let Some(status) = proxy.0.try_wait().unwrap() { + break status; + } + assert!(Instant::now() < deadline, "the proxy kept running"); + std::thread::sleep(Duration::from_millis(50)); + }; + assert!(status.success(), "{status}"); + assert!(TcpStream::connect(("127.0.0.1", port)).is_err()); +} + /// A config file for this test process holding `yaml`, removed on drop. struct ConfigFile(std::path::PathBuf); diff --git a/tests/connection_memory.rs b/tests/connection_memory.rs new file mode 100644 index 0000000..f61976e --- /dev/null +++ b/tests/connection_memory.rs @@ -0,0 +1,98 @@ +//! What the listener keeps of a connection after it closes. A test binary of +//! its own: it counts every allocation of the process. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicIsize, Ordering}; +use std::time::Duration; + +use structured_proxy::{ProxyServer, ServeOptions}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; + +/// The system allocator, keeping count of the bytes live. +struct Counting; + +static LIVE: AtomicIsize = AtomicIsize::new(0); + +// SAFETY: every call goes to `System` unchanged; the counter only observes. +unsafe impl GlobalAlloc for Counting { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + LIVE.fetch_add(layout.size() as isize, Ordering::Relaxed); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + LIVE.fetch_sub(layout.size() as isize, Ordering::Relaxed); + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new = unsafe { System.realloc(ptr, layout, new_size) }; + if !new.is_null() { + LIVE.fetch_add( + new_size as isize - layout.size() as isize, + Ordering::Relaxed, + ); + } + new + } +} + +#[global_allocator] +static ALLOCATOR: Counting = Counting; + +/// A keep-alive connection that has had one `GET /health/live` answered. +async fn served(addr: std::net::SocketAddr) -> TcpStream { + let mut stream = TcpStream::connect(addr).await.unwrap(); + stream + .write_all(b"GET /health/live HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + // The status line is all the case needs; the rest stays unread. + let mut status = [0; 12]; + stream.read_exact(&mut status).await.unwrap(); + assert_eq!(&status, b"HTTP/1.1 200"); + stream +} + +#[tokio::test] +async fn closed_connections_are_released_while_the_listener_waits() { + // After a burst the server sits in accept with no new client; what the + // closed connections held must be freed then, not at the next accept. + let service = ProxyServer::from_yaml_str("service:\n name: demo\n") + .unwrap() + .service(tonic::service::Routes::default()) + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(structured_proxy::serve_with( + listener, + service, + ServeOptions::new(), + )); + // Lazily built state (the router, runtime buffers) is in place after one. + drop(served(addr).await); + tokio::time::sleep(Duration::from_millis(100)).await; + let before = LIVE.load(Ordering::Relaxed); + + // Held open together, then closed together: they end while the server + // waits for a client that does not come. Few enough for the default + // file descriptor limit of macOS (256). + const BURST: isize = 100; + let mut open = Vec::new(); + for _ in 0..BURST { + open.push(served(addr).await); + } + drop(open); + tokio::time::sleep(Duration::from_millis(300)).await; + let kept = LIVE.load(Ordering::Relaxed) - before; + // A closed connection's task keeps about 190 bytes until it is reaped; + // the bound leaves room for allocator noise, not for the tasks. + assert!( + kept < 64 * BURST, + "{kept} bytes kept for {BURST} closed connections" + ); +} diff --git a/tests/embedded.rs b/tests/embedded.rs index c1d4ba8..0936b95 100644 --- a/tests/embedded.rs +++ b/tests/embedded.rs @@ -28,6 +28,7 @@ fn embedded_config_is_constructible() { tls: None, idle_timeout_secs: 60, header_read_timeout_secs: 30, + drain_timeout_secs: 25, }, service: ServiceConfig { name: "embedded-test".into(), diff --git a/tests/shutdown.rs b/tests/shutdown.rs new file mode 100644 index 0000000..4a4a9f2 --- /dev/null +++ b/tests/shutdown.rs @@ -0,0 +1,516 @@ +//! Graceful shutdown of the built-in listener (`serve_with_shutdown`): after +//! the signal new connections are refused, requests and streams in flight +//! finish, idle connections are closed, and dropping the serve future closes +//! everything at once. Each case runs in cleartext, behind TLS and under a +//! connection limit. + +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use http_body_util::BodyExt as _; +use hyper_util::rt::TokioIo; +use rustls::pki_types::pem::PemObject; +use rustls::pki_types::{CertificateDer, ServerName}; +use structured_proxy::{ProxyServer, ServeOptions}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio::sync::{mpsc, oneshot, watch}; +use tokio::task::JoinHandle; +use tower::ServiceExt as _; + +const TESTDATA: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/src/tls/testdata"); +const CA: &str = include_str!("../src/tls/testdata/ca.pem"); +/// How long a test waits for what should happen right away; far below the +/// drain timeout, so no case passes by running into it. +const PROMPT: Duration = Duration::from_secs(5); + +#[derive(Clone, Copy, Debug)] +struct Setup { + tls: bool, + max_connections: Option, +} + +const SETUPS: [Setup; 3] = [ + Setup { + tls: false, + max_connections: None, + }, + Setup { + tls: true, + max_connections: None, + }, + // One slot: the one connection a case opens holds it, so the accept loop + // sits waiting for a slot when the signal comes. + Setup { + tls: false, + max_connections: Some(1), + }, +]; + +trait Io: AsyncRead + AsyncWrite + Unpin + Send + 'static {} +impl Io for T {} + +/// Tells the test a handler future was dropped. +struct OnDrop(mpsc::UnboundedSender<()>); + +impl Drop for OnDrop { + fn drop(&mut self) { + // The test may have stopped listening; nothing to report then. + self.0.send(()).ok(); + } +} + +/// Routes whose answers the test holds: `/slow` answers once released, +/// `/stream` sends a first chunk and the last one once released, `/hang` +/// never answers and reports when its handler is dropped. Each reports on +/// `entered` once its handler runs. +fn held_routes( + entered: &mpsc::UnboundedSender<()>, + release: &watch::Receiver, + dropped: &mpsc::UnboundedSender<()>, +) -> axum::Router { + use axum::routing::get; + let slow = { + let (entered, release) = (entered.clone(), release.clone()); + move || { + let (entered, mut release) = (entered.clone(), release.clone()); + async move { + entered.send(()).unwrap(); + release.wait_for(|go| *go).await.unwrap(); + "done" + } + } + }; + let stream = { + let (entered, release) = (entered.clone(), release.clone()); + move || { + let (entered, mut release) = (entered.clone(), release.clone()); + async move { + entered.send(()).unwrap(); + let first = futures::stream::once(async { Ok::<_, std::io::Error>("first\n") }); + let last = futures::stream::once(async move { + release.wait_for(|go| *go).await.unwrap(); + Ok("last\n") + }); + axum::body::Body::from_stream(futures::StreamExt::chain(first, last)) + } + } + }; + let hang = { + let (entered, dropped) = (entered.clone(), dropped.clone()); + move || { + let (entered, dropped) = (entered.clone(), dropped.clone()); + async move { + let _dropped = OnDrop(dropped); + entered.send(()).unwrap(); + std::future::pending::<&'static str>().await + } + } + }; + axum::Router::new() + .route("/slow", get(slow)) + .route("/stream", get(stream)) + .route("/hang", get(hang)) +} + +/// The proxy serving [`held_routes`] until the test stops it. +struct Proxy { + addr: SocketAddr, + setup: Setup, + entered: mpsc::UnboundedReceiver<()>, + release: watch::Sender, + dropped: mpsc::UnboundedReceiver<()>, + stop: Option>, + served: JoinHandle>, +} + +impl Proxy { + /// Serve with the options `listen:` gives for `setup`, changed by + /// `configure`. Idle connections stay open, so what closes one in a case + /// is the shutdown. + async fn start(setup: Setup, configure: impl FnOnce(ServeOptions) -> ServeOptions) -> Self { + let mut listen = String::from("listen:\n idle_timeout_secs: 0\n"); + if let Some(max) = setup.max_connections { + listen.push_str(&format!(" max_connections: {max}\n")); + } + if setup.tls { + install_provider(); + listen.push_str(&format!( + " tls:\n cert_file: {TESTDATA}/ecdsa.pem\n key_file: {TESTDATA}/ecdsa.key.pem\n" + )); + } + let server = ProxyServer::from_yaml_str(&listen).unwrap(); + let options = configure(server.serve_options().unwrap()); + let (entered_tx, entered) = mpsc::unbounded_channel(); + let (release, release_rx) = watch::channel(false); + let (dropped_tx, dropped) = mpsc::unbounded_channel(); + let service = server + .service(tonic::service::Routes::default()) + .unwrap() + .with_fallback(held_routes(&entered_tx, &release_rx, &dropped_tx)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (stop, signal) = oneshot::channel::<()>(); + let served = tokio::spawn(structured_proxy::serve_with_shutdown( + listener, + service, + options, + async { + signal.await.unwrap(); + }, + )); + Self { + addr, + setup, + entered, + release, + dropped, + stop: Some(stop), + served, + } + } + + /// A connection negotiating `alpn` behind TLS. + async fn connect(&self, alpn: &[u8]) -> Box { + connect(self.addr, self.setup.tls, alpn).await + } + + /// A tonic channel: an HTTP/2 client with a connection of its own. + async fn channel(&self) -> tonic::transport::Channel { + let (addr, tls) = (self.addr, self.setup.tls); + let connector = tower::service_fn(move |_: http::Uri| async move { + Ok::<_, std::io::Error>(TokioIo::new(connect(addr, tls, b"h2").await)) + }); + tonic::transport::Endpoint::from_static("http://localhost") + .connect_with_connector(connector) + .await + .unwrap() + } + + /// Wait until `count` handlers have started. + async fn entered(&mut self, count: usize) { + for _ in 0..count { + tokio::time::timeout(PROMPT, self.entered.recv()) + .await + .expect("the request reaches its handler") + .unwrap(); + } + } + + /// Wait until `count` handler futures have been dropped. + async fn dropped(&mut self, count: usize) { + for _ in 0..count { + tokio::time::timeout(PROMPT, self.dropped.recv()) + .await + .expect("the handler is dropped") + .unwrap(); + } + } + + fn signal(&mut self) { + self.stop.take().unwrap().send(()).unwrap(); + } + + fn release(&self) { + self.release.send_replace(true); + } + + /// Wait until the listening socket is closed. + async fn refuses_connections(&self) { + tokio::time::timeout(PROMPT, async { + while TcpStream::connect(self.addr).await.is_ok() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("new connections are refused after the signal"); + } + + /// Wait for the serve future to resolve. + async fn served(&mut self) { + tokio::time::timeout(PROMPT, &mut self.served) + .await + .expect("the serve future resolves") + .unwrap() + .unwrap(); + } +} + +fn install_provider() { + if rustls::crypto::CryptoProvider::get_default().is_none() { + rustls::crypto::CryptoProvider::install_default(rustls_rustcrypto::provider()) + .expect("each test runs in its own process"); + } +} + +async fn connect(addr: SocketAddr, tls: bool, alpn: &[u8]) -> Box { + let tcp = TcpStream::connect(addr).await.unwrap(); + if !tls { + return Box::new(tcp); + } + let mut roots = rustls::RootCertStore::empty(); + for cert in CertificateDer::pem_slice_iter(CA.as_bytes()) { + roots.add(cert.unwrap()).unwrap(); + } + let mut config = + rustls::ClientConfig::builder_with_provider(Arc::new(rustls_rustcrypto::provider())) + .with_safe_default_protocol_versions() + .unwrap() + .with_root_certificates(roots) + .with_no_client_auth(); + config.alpn_protocols = vec![alpn.to_vec()]; + let stream = tokio_rustls::TlsConnector::from(Arc::new(config)) + .connect(ServerName::try_from("localhost").unwrap(), tcp) + .await + .unwrap(); + Box::new(stream) +} + +/// Read until the server closes the connection; returns what it sent. +async fn read_until_closed(io: &mut (impl AsyncRead + Unpin)) -> Vec { + let mut received = Vec::new(); + let read = tokio::time::timeout(PROMPT, io.read_to_end(&mut received)) + .await + .expect("the server closes the connection"); + // An end of stream, a reset, or TLS closed without close_notify: the + // server let go either way. + if let Err(error) = read { + use std::io::ErrorKind; + assert!( + matches!( + error.kind(), + ErrorKind::ConnectionReset | ErrorKind::UnexpectedEof + ), + "{error}" + ); + } + received +} + +async fn get(channel: &tonic::transport::Channel, path: &str) -> http::Response { + send(channel, path).await.unwrap() +} + +/// A call the server is expected to cut off: its outcome is left to the case. +async fn send( + channel: &tonic::transport::Channel, + path: &str, +) -> Result, tonic::transport::Error> { + let request = http::Request::get(path) + .body(tonic::body::Body::empty()) + .unwrap(); + channel.clone().oneshot(request).await +} + +#[tokio::test] +async fn an_http1_request_in_flight_completes_and_new_connections_are_refused() { + for setup in SETUPS { + let mut proxy = Proxy::start(setup, |options| options).await; + let mut client = proxy.connect(b"http/1.1").await; + client + .write_all(b"GET /slow HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + proxy.entered(1).await; + + proxy.signal(); + proxy.refuses_connections().await; + assert!( + !proxy.served.is_finished(), + "{setup:?}: a request is in flight" + ); + + proxy.release(); + let response = String::from_utf8(read_until_closed(&mut client).await).unwrap(); + assert!( + response.starts_with("HTTP/1.1 200"), + "{setup:?}: {response}" + ); + assert!(response.ends_with("done"), "{setup:?}: {response}"); + proxy.served().await; + } +} + +#[tokio::test] +async fn http2_calls_and_streams_in_flight_complete() { + for setup in SETUPS { + let mut proxy = Proxy::start(setup, |options| options).await; + let channel = proxy.channel().await; + let unary = tokio::spawn({ + let channel = channel.clone(); + async move { + let body = get(&channel, "/slow").await.into_body(); + body.collect().await.unwrap().to_bytes() + } + }); + let mut stream = get(&channel, "/stream").await.into_body(); + proxy.entered(2).await; + let first = stream.frame().await.unwrap().unwrap().into_data().unwrap(); + assert_eq!(first, "first\n", "{setup:?}"); + + proxy.signal(); + proxy.refuses_connections().await; + assert!(!proxy.served.is_finished(), "{setup:?}: streams are open"); + + proxy.release(); + let rest = stream.collect().await.unwrap().to_bytes(); + assert_eq!(rest, "last\n", "{setup:?}"); + assert_eq!(unary.await.unwrap(), "done", "{setup:?}"); + proxy.served().await; + } +} + +#[tokio::test] +async fn an_idle_http2_client_gets_a_goaway() { + for setup in SETUPS { + let mut proxy = Proxy::start(setup, |options| options).await; + let io = proxy.connect(b"h2").await; + let (sender, connection) = h2::client::handshake(io).await.unwrap(); + let connection = tokio::spawn(connection); + let mut sender = sender.ready().await.unwrap(); + let (response, _) = sender + .send_request( + http::Request::get("http://localhost/health/live") + .body(()) + .unwrap(), + true, + ) + .unwrap(); + assert_eq!(response.await.unwrap().status(), 200, "{setup:?}"); + + // The client keeps its connection; the server closes it. + proxy.signal(); + proxy.served().await; + let ended = tokio::time::timeout(PROMPT, connection) + .await + .expect("the client's connection ends") + .unwrap(); + assert!(ended.is_ok(), "{setup:?}: {ended:?}"); + let refused = sender + .send_request( + http::Request::get("http://localhost/health/live") + .body(()) + .unwrap(), + true, + ) + .unwrap_err(); + assert!(refused.is_go_away(), "{setup:?}: {refused}"); + assert_eq!(refused.reason(), Some(h2::Reason::NO_ERROR), "{setup:?}"); + } +} + +#[tokio::test] +async fn a_connection_that_sent_nothing_is_closed() { + // Behind TLS it is still in its handshake; in cleartext hyper has not + // seen which HTTP version it speaks. + for setup in SETUPS { + let mut proxy = Proxy::start(setup, |options| options).await; + let mut silent = TcpStream::connect(proxy.addr).await.unwrap(); + // Let the server accept it before the signal. + tokio::time::sleep(Duration::from_millis(100)).await; + + proxy.signal(); + proxy.served().await; + assert_eq!(read_until_closed(&mut silent).await, b"", "{setup:?}"); + } +} + +#[tokio::test] +async fn a_connection_waiting_for_a_slot_is_closed() { + let setup = Setup { + tls: false, + max_connections: Some(1), + }; + let mut proxy = Proxy::start(setup, |options| options).await; + let mut first = proxy.connect(b"http/1.1").await; + first + .write_all(b"GET /health/live HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + let mut head = [0; 12]; + first.read_exact(&mut head).await.unwrap(); + assert_eq!(&head, b"HTTP/1.1 200"); + // The first connection holds the one slot: the kernel completes this one, + // but the proxy never accepts it. + let mut waiting = TcpStream::connect(proxy.addr).await.unwrap(); + waiting + .write_all(b"GET /health/live HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + + proxy.signal(); + proxy.served().await; + assert_eq!(read_until_closed(&mut waiting).await, b""); + read_until_closed(&mut first).await; +} + +#[tokio::test] +async fn the_drain_timeout_closes_connections_still_busy() { + for setup in SETUPS { + let mut proxy = Proxy::start(setup, |options| { + options.drain_timeout(Some(Duration::from_millis(300))) + }) + .await; + let mut client = proxy.connect(b"http/1.1").await; + client + .write_all(b"GET /hang HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + proxy.entered(1).await; + + proxy.signal(); + proxy.served().await; + proxy.dropped(1).await; + assert_eq!(read_until_closed(&mut client).await, b"", "{setup:?}"); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_drain_timeout_ends_http2_handlers_before_returning() { + // HTTP/2 serves each stream in a task of its own: when the serve future + // resolves, those have ended too, not only their connections. + for setup in SETUPS { + let mut proxy = Proxy::start(setup, |options| { + options.drain_timeout(Some(Duration::from_millis(300))) + }) + .await; + let channel = proxy.channel().await; + let call = tokio::spawn(async move { send(&channel, "/hang").await }); + proxy.entered(1).await; + + proxy.signal(); + proxy.served().await; + assert!( + proxy.dropped.try_recv().is_ok(), + "{setup:?}: a handler outlived serve_with_shutdown" + ); + call.abort(); + } +} + +#[tokio::test] +async fn dropping_the_serve_future_ends_every_connection_and_handler() { + // Unlimited: this case opens two connections. + for setup in &SETUPS[..2] { + let mut proxy = Proxy::start(*setup, |options| options).await; + let mut http1 = proxy.connect(b"http/1.1").await; + http1 + .write_all(b"GET /hang HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + let channel = proxy.channel().await; + let http2 = tokio::spawn(async move { send(&channel, "/hang").await }); + proxy.entered(2).await; + + proxy.served.abort(); + // HTTP/2 serves each stream in a task of its own; that one ends too. + proxy.dropped(2).await; + assert_eq!(read_until_closed(&mut http1).await, b"", "{setup:?}"); + let refused = tokio::time::timeout(PROMPT, TcpStream::connect(proxy.addr)) + .await + .expect("connecting completes"); + assert!(refused.is_err(), "{setup:?}: the listener is closed"); + http2.abort(); + } +}