chio_http_serve/
hygiene.rs

1use std::time::Duration;
2
3use axum::error_handling::HandleErrorLayer;
4use axum::extract::DefaultBodyLimit;
5use axum::http::StatusCode;
6use axum::Router;
7use tower::limit::GlobalConcurrencyLimitLayer;
8use tower::load_shed::LoadShedLayer;
9use tower::{BoxError, ServiceBuilder};
10use tower_http::timeout::TimeoutLayer;
11
12/// Wall-clock ceiling on the drain: how long to wait for in-flight requests to
13/// finish after the listener stops accepting. Operators must size the unit
14/// `TimeoutStopSec` at least this high plus a flush margin.
15pub const DEFAULT_DRAIN_TIMEOUT: Duration = Duration::from_secs(25);
16/// Per-request processing ceiling before the request is denied with 408.
17///
18/// This stays strictly below [`DEFAULT_DRAIN_TIMEOUT`] so a request admitted just
19/// before a stop signal reaches its own 408 and completes cleanly within the
20/// drain window, rather than being severed mid-flight when the forced-drain timer
21/// force-closes the connection. A serve site that lengthens the request timeout
22/// must lengthen the drain to match.
23pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(20);
24/// Maximum concurrent in-flight requests before surplus load sheds with 503.
25pub const DEFAULT_MAX_CONCURRENT_REQUESTS: usize = 1024;
26/// Maximum simultaneously accepted TCP connections.
27pub const DEFAULT_MAX_CONNECTIONS: usize = 2048;
28
29/// Bounds and hygiene knobs for one serve site. The [`Default`] is a
30/// conservative, fail-closed posture: a request that exceeds a bound is denied,
31/// never queued.
32#[derive(Debug, Clone)]
33pub struct ServeHygieneConfig {
34    /// How long the drain waits for in-flight requests after the listener stops
35    /// accepting. See [`DEFAULT_DRAIN_TIMEOUT`].
36    pub drain_timeout: Duration,
37    /// Per-request processing timeout. `None` disables it.
38    pub request_timeout: Option<Duration>,
39    /// Maximum concurrent in-flight requests. Surplus load sheds with 503
40    /// instead of queuing. `None` disables the limit.
41    pub max_concurrent_requests: Option<usize>,
42    /// Maximum simultaneously accepted TCP connections, enforced by
43    /// [`MaxConnListener`](crate::MaxConnListener). `None` disables the cap.
44    pub max_connections: Option<usize>,
45    /// Global request body-size cap. `None` preserves each route's own limit,
46    /// which matters for sites that set a large upload limit on one route that a
47    /// global cap would clobber. A site with no route-local limit sets this
48    /// explicitly.
49    pub max_body_bytes: Option<usize>,
50}
51
52impl Default for ServeHygieneConfig {
53    fn default() -> Self {
54        Self {
55            drain_timeout: DEFAULT_DRAIN_TIMEOUT,
56            request_timeout: Some(DEFAULT_REQUEST_TIMEOUT),
57            max_concurrent_requests: Some(DEFAULT_MAX_CONCURRENT_REQUESTS),
58            max_connections: Some(DEFAULT_MAX_CONNECTIONS),
59            max_body_bytes: None,
60        }
61    }
62}
63
64/// Map the error surfaced by the load-shed / concurrency-limit stack to a
65/// status. A shed request is a deliberate overload denial (503); anything else
66/// reaching this handler is an internal fault (500).
67async fn shed_to_status(error: BoxError) -> StatusCode {
68    if error.is::<tower::load_shed::error::Overloaded>() {
69        StatusCode::SERVICE_UNAVAILABLE
70    } else {
71        StatusCode::INTERNAL_SERVER_ERROR
72    }
73}
74
75/// Wrap `router` with the configured request timeout, concurrency limit fronted
76/// by load shedding, and optional body-size cap.
77///
78/// The stack returns a plain [`Router`], so a call site never touches tower
79/// types beyond this call. Load shedding sits outside the concurrency limit so a
80/// request over the limit fails fast with 503 rather than parking: in a
81/// [`ServiceBuilder`] the first layer listed is outermost, giving
82/// `HandleError(LoadShed(ConcurrencyLimit(router)))`. Both `LoadShed` and
83/// `GlobalConcurrencyLimit` are fallible services, so the pair is wrapped in a
84/// [`HandleErrorLayer`] that turns the shed error into a response, which is what
85/// lets the result stay an infallible `Router` layer.
86pub fn apply_server_hygiene(mut router: Router, config: &ServeHygieneConfig) -> Router {
87    if let Some(limit) = config.max_body_bytes {
88        router = router.layer(DefaultBodyLimit::max(limit));
89    }
90    if let Some(timeout) = config.request_timeout {
91        router = router.layer(TimeoutLayer::with_status_code(
92            StatusCode::REQUEST_TIMEOUT,
93            timeout,
94        ));
95    }
96    if let Some(max) = config.max_concurrent_requests {
97        router = router.layer(
98            ServiceBuilder::new()
99                .layer(HandleErrorLayer::new(shed_to_status))
100                .layer(LoadShedLayer::new())
101                .layer(GlobalConcurrencyLimitLayer::new(max)),
102        );
103    }
104    router
105}
chio_http_serve/hygiene.rs

chio_http_serve/
hygiene.rs
