Skip to main content

kerniq_runtime/
server.rs

1//! The HTTP accept loop.
2//!
3//! Kerniq drives Hyper directly (through `hyper-util`'s HTTP/1 + HTTP/2
4//! auto-detecting connection builder) instead of using `axum::serve`,
5//! because it needs three things `axum::serve` does not offer:
6//!
7//! * a **connection limit** (`max_connections`),
8//! * a **bounded drain**: in-flight connections get `shutdown_timeout` to
9//!   finish, after which they are closed,
10//! * control over the header read timeout.
11//!
12//! The per-request path is identical to `axum::serve`: each connection gets a
13//! clone of the `axum::Router` (an `Arc` bump) and requests are dispatched
14//! straight into it.
15
16use std::future::Future;
17use std::io;
18use std::net::SocketAddr;
19use std::sync::Arc;
20use std::time::Duration;
21
22use axum::Router;
23use axum::extract::ConnectInfo;
24use http::Request;
25use hyper::body::Incoming;
26use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
27use hyper_util::server::conn::auto::Builder;
28use hyper_util::server::graceful::GracefulShutdown;
29use hyper_util::service::TowerToHyperService;
30use kerniq_core::{BoxError, ServerConfig, ServerError};
31use tokio::net::{TcpListener, TcpStream};
32use tokio::sync::Semaphore;
33use tokio_util::sync::CancellationToken;
34use tower::ServiceExt;
35
36use crate::Shutdown;
37
38/// Connection-level options, usually derived from a [`ServerConfig`].
39#[derive(Debug, Clone)]
40pub struct ServeOptions {
41    /// Maximum number of open connections (`None` = unlimited).
42    pub max_connections: Option<usize>,
43    /// How long to wait for in-flight connections during shutdown.
44    pub shutdown_timeout: Duration,
45    /// Maximum time to receive request headers on HTTP/1 connections
46    /// (`Duration::ZERO` disables the timeout).
47    pub header_read_timeout: Duration,
48    /// Set `TCP_NODELAY` on accepted sockets.
49    pub tcp_nodelay: bool,
50    /// Insert `ConnectInfo<SocketAddr>` into each request.
51    pub connect_info: bool,
52}
53
54impl Default for ServeOptions {
55    fn default() -> Self {
56        Self::from(&ServerConfig::default())
57    }
58}
59
60impl From<&ServerConfig> for ServeOptions {
61    fn from(config: &ServerConfig) -> Self {
62        Self {
63            max_connections: config.max_connections,
64            shutdown_timeout: config.shutdown_timeout,
65            header_read_timeout: config.header_read_timeout,
66            tcp_nodelay: config.tcp_nodelay,
67            connect_info: config.connect_info,
68        }
69    }
70}
71
72/// Binds a TCP listener. Host names are resolved; the first address that
73/// binds successfully is used.
74pub async fn bind(address: &str) -> Result<TcpListener, ServerError> {
75    if address.trim().is_empty() {
76        return Err(ServerError::InvalidAddress {
77            address: address.to_owned(),
78            reason: "address is empty".to_owned(),
79        });
80    }
81    TcpListener::bind(address)
82        .await
83        .map_err(|source| match source.kind() {
84            io::ErrorKind::InvalidInput => ServerError::InvalidAddress {
85                address: address.to_owned(),
86                reason: source.to_string(),
87            },
88            _ => ServerError::Bind {
89                address: address.to_owned(),
90                source,
91            },
92        })
93}
94
95/// Serves `router` on `listener` until `shutdown` is triggered, then drains
96/// in-flight connections for up to `options.shutdown_timeout`.
97///
98/// Shutdown sequence:
99///
100/// 1. stop accepting (the listener is dropped, so new connections are
101///    refused by the OS),
102/// 2. ask every open connection to finish: HTTP/1 connections close after
103///    their in-flight response, HTTP/2 connections send `GOAWAY`,
104/// 3. after `shutdown_timeout`, forcibly close whatever is left.
105pub async fn serve(
106    listener: TcpListener,
107    router: Router,
108    options: ServeOptions,
109    shutdown: Shutdown,
110) -> Result<(), ServerError> {
111    let local_addr = listener.local_addr()?;
112    let limiter = options.max_connections.map(|n| Arc::new(Semaphore::new(n)));
113    let graceful = GracefulShutdown::new();
114    let force_close = CancellationToken::new();
115
116    let mut builder = Builder::new(TokioExecutor::new());
117    // The header read timeout protects against slow-header (slowloris)
118    // clients. It arms a timer for every request on a keep-alive connection,
119    // which measurably costs throughput on trivial endpoints (see
120    // docs/performance.md), so it can be disabled with a zero duration.
121    if !options.header_read_timeout.is_zero() {
122        builder
123            .http1()
124            .timer(TokioTimer::new())
125            .header_read_timeout(options.header_read_timeout);
126        builder.http2().timer(TokioTimer::new());
127    }
128
129    tracing::info!(address = %local_addr, "kerniq server listening");
130
131    loop {
132        let permit = match &limiter {
133            Some(semaphore) => tokio::select! {
134                biased;
135                () = shutdown.wait() => break,
136                permit = semaphore.clone().acquire_owned() => match permit {
137                    Ok(permit) => Some(permit),
138                    Err(_) => break, // semaphore closed: cannot happen, but stop cleanly
139                },
140            },
141            None => None,
142        };
143
144        let (stream, peer) = tokio::select! {
145            biased;
146            () = shutdown.wait() => break,
147            accepted = listener.accept() => match accepted {
148                Ok(pair) => pair,
149                Err(err) => {
150                    handle_accept_error(err).await;
151                    continue;
152                }
153            },
154        };
155
156        if options.tcp_nodelay
157            && let Err(err) = stream.set_nodelay(true)
158        {
159            tracing::debug!(error = %err, "failed to set TCP_NODELAY");
160        }
161
162        let connection = serve_connection(
163            &builder,
164            &graceful,
165            stream,
166            peer,
167            router.clone(),
168            options.connect_info,
169        );
170        let force_close = force_close.clone();
171        tokio::spawn(async move {
172            tokio::select! {
173                result = connection => {
174                    if let Err(err) = result {
175                        tracing::trace!(peer = %peer, error = %err, "connection closed with error");
176                    }
177                }
178                () = force_close.cancelled() => {
179                    tracing::debug!(peer = %peer, "connection force-closed after shutdown timeout");
180                }
181            }
182            drop(permit);
183        });
184    }
185
186    drop(listener);
187    let open = graceful.count();
188    tracing::info!(
189        open_connections = open,
190        "stopped accepting connections, draining"
191    );
192
193    match tokio::time::timeout(options.shutdown_timeout, graceful.shutdown()).await {
194        Ok(()) => tracing::info!("all connections drained"),
195        Err(_) => {
196            tracing::warn!(
197                timeout = ?options.shutdown_timeout,
198                "shutdown timeout elapsed, closing remaining connections"
199            );
200            force_close.cancel();
201        }
202    }
203    Ok(())
204}
205
206fn serve_connection(
207    builder: &Builder<TokioExecutor>,
208    graceful: &GracefulShutdown,
209    stream: TcpStream,
210    peer: SocketAddr,
211    router: Router,
212    connect_info: bool,
213) -> impl Future<Output = Result<(), BoxError>> + Send + 'static {
214    let service = router.map_request(move |mut request: Request<Incoming>| {
215        if connect_info {
216            request.extensions_mut().insert(ConnectInfo(peer));
217        }
218        request
219    });
220    let connection = builder
221        .serve_connection_with_upgrades(TokioIo::new(stream), TowerToHyperService::new(service))
222        .into_owned();
223    graceful.watch(connection)
224}
225
226/// Errors on `accept` are either per-connection (the peer went away) or
227/// resource exhaustion (EMFILE). The former are ignored; for the latter we
228/// back off briefly so a full file-descriptor table does not spin the loop.
229async fn handle_accept_error(err: io::Error) {
230    if matches!(
231        err.kind(),
232        io::ErrorKind::ConnectionRefused
233            | io::ErrorKind::ConnectionAborted
234            | io::ErrorKind::ConnectionReset
235    ) {
236        return;
237    }
238    tracing::error!(error = %err, "failed to accept connection");
239    tokio::time::sleep(Duration::from_millis(50)).await;
240}