1mod 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#[derive(Debug, Clone, Copy)]
25pub struct ServerLimits {
26 pub max_connections: usize,
28 pub max_http2_streams: u32,
30 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 #[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}