Skip to main content

siderite_core/
server.rs

1//! Owned HTTP connection and HTTP/2 stream tasks.
2
3mod readiness;
4pub(crate) mod tasks;
5pub use readiness::{Readiness, ServerPhase};
6#[cfg(all(test, unix))]
7mod signal_tests;
8#[cfg(test)]
9mod tests;
10
11use crate::{App, ServerError};
12use hyper_util::rt::TokioIo;
13use hyper_util::server::conn::auto::Builder;
14use hyper_util::service::TowerToHyperService;
15use std::future::Future;
16use std::net::SocketAddr as Peer;
17use std::time::Duration;
18use tokio::net::TcpListener;
19use tokio::sync::watch;
20use tokio::time::Instant;
21use tower::Service;
22
23/// Hard transport admission bounds applied by [`App::run`].
24#[derive(Debug, Clone, Copy)]
25pub struct ServerLimits {
26    /// Maximum accepted sockets, including idle keep-alive connections.
27    pub max_connections: usize,
28    /// Concurrent HTTP/2 streams permitted on each connection.
29    pub max_http2_streams: u32,
30    /// Active or pending WebSocket upgrades across the app tree.
31    pub max_websockets: usize,
32}
33
34impl Default for ServerLimits {
35    fn default() -> Self {
36        Self {
37            max_connections: 1024,
38            max_http2_streams: 100,
39            max_websockets: 128,
40        }
41    }
42}
43
44impl App {
45    /// Configure bounded transport admission before binding the server.
46    ///
47    /// Args:
48    ///     limits: Accepted socket and per-connection HTTP/2 limits.
49    ///
50    /// Returns:
51    ///     This application; invalid bounds become configuration errors.
52    #[must_use]
53    pub fn server_limits(mut self, limits: ServerLimits) -> Self {
54        if limits.capacity().is_none() {
55            self.config_errors.push("invalid server limits".to_owned());
56        } else {
57            self.server_limits = limits;
58        }
59        self
60    }
61}
62
63impl ServerLimits {
64    fn capacity(self) -> Option<usize> {
65        let streams = usize::try_from(self.max_http2_streams).ok()?;
66        let tasks = self.max_connections.checked_mul(streams.checked_add(4)?)?;
67        (self.max_connections > 0
68            && streams > 0
69            && tasks <= 1_000_000
70            && self.max_websockets <= 1_000_000)
71            .then_some(tasks)
72    }
73}
74
75#[derive(Clone)]
76pub(crate) struct ServerOwners {
77    connections: tasks::TaskOwner,
78    streams: tasks::TaskOwner,
79    pub(crate) readiness: Readiness,
80}
81
82impl ServerOwners {
83    pub(crate) fn new(limits: ServerLimits) -> Self {
84        Self {
85            connections: tasks::TaskOwner::new(limits.max_connections),
86            streams: tasks::TaskOwner::new(limits.capacity().unwrap_or(1)),
87            readiness: Readiness::new(),
88        }
89    }
90
91    pub(crate) async fn stop_at(&self, deadline: Instant) -> bool {
92        self.readiness.draining();
93        let connections = self.connections.stop_at(deadline).await;
94        let streams = self.streams.stop_at(deadline).await;
95        connections && streams
96    }
97}
98
99struct AbortOnDrop(ServerOwners);
100
101impl Drop for AbortOnDrop {
102    fn drop(&mut self) {
103        self.0.readiness.draining();
104        self.0.connections.abort();
105        self.0.streams.abort();
106    }
107}
108
109pub(crate) async fn serve<F>(
110    listener: TcpListener,
111    router: axum::Router,
112    limits: ServerLimits,
113    owners: ServerOwners,
114    budget: Duration,
115    shutdown: F,
116) -> (Result<(), ServerError>, Instant)
117where
118    F: Future<Output = ()> + Send,
119{
120    let _guard = AbortOnDrop(owners.clone());
121    let streams = &owners.streams;
122    let (draining, receiver) = watch::channel(false);
123    let workers = &owners.connections;
124    let mut service = router.into_make_service_with_connect_info::<Peer>();
125    tokio::pin!(shutdown);
126    let error = loop {
127        tokio::select! {
128            biased;
129            _ = &mut shutdown => break None,
130            joined = async {
131                match workers.join_next().await {
132                    Some(result) => Some(result),
133                    None => std::future::pending().await,
134                }
135            } => {
136                if let Some(Err(error)) = joined {
137                    tracing::error!(%error, "HTTP connection task failed");
138                }
139            }
140            accepted = async {
141                let permit = workers.handle().acquire().await;
142                let accepted = listener.accept().await;
143                (permit, accepted)
144            } => {
145                let (permit, accepted) = accepted;
146                let (socket, peer) = match accepted {
147                    Ok(accepted) => accepted,
148                    Err(error) => break Some(ServerError::Serve(error)),
149                };
150                let Some(permit) = permit else { break None; };
151                let per_connection = match service.call(peer).await {
152                    Ok(service) => service,
153                    Err(error) => match error {},
154                };
155                let executor = streams.handle();
156                let mut stopping = receiver.clone();
157                workers.handle().spawn(permit, async move {
158                    let mut builder = Builder::new(executor);
159                    builder.http2()
160                        .max_concurrent_streams(limits.max_http2_streams);
161                    let io = TokioIo::new(socket);
162                    let service = TowerToHyperService::new(per_connection);
163                    let connection = builder
164                        .serve_connection_with_upgrades(io, service);
165                    tokio::pin!(connection);
166                    let result = tokio::select! {
167                        result = &mut connection => result,
168                        _ = stopping.changed() => {
169                            connection.as_mut().graceful_shutdown();
170                            connection.await
171                        }
172                    };
173                    if let Err(error) = result {
174                        tracing::debug!(%error, "HTTP connection ended");
175                    }
176                });
177            }
178        }
179    };
180    drop(listener);
181    owners.readiness.draining();
182    let deadline = Instant::now() + budget;
183    let _ = draining.send(true);
184    let timed_out = !owners.stop_at(deadline).await;
185    let result = error.map_or_else(
186        || {
187            if timed_out {
188                Err(ServerError::ShutdownTimeout)
189            } else {
190                Ok(())
191            }
192        },
193        Err,
194    );
195    (result, deadline)
196}