1use 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#[derive(Debug, Clone)]
40pub struct ServeOptions {
41 pub max_connections: Option<usize>,
43 pub shutdown_timeout: Duration,
45 pub header_read_timeout: Duration,
48 pub tcp_nodelay: bool,
50 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
72pub 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
95pub 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 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, },
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
226async 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}