OpenSource/axum
OpenSource/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2024-10-23 17:36:39 +02:00
Code Issues Projects Releases Packages Wiki Activity

Prepare serve for potentially supporting graceful shutdown (#2357)

Browse source
This commit is contained in:
David Pedersen 2023-11-26 13:21:35 +01:00 committed by GitHub
parent da82e887cd
commit 15781fe22b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 125 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
axum-macros/tests/debug_handler/fail/argument_not_extractor.stderr
View file

@ -12,8 +12,8 @@ error[E0277]: the trait bound `bool: FromRequestParts<()>` is not satisfied
<Method as FromRequestParts<S>>
<axum::http::request::Parts as FromRequestParts<S>>
<Uri as FromRequestParts<S>>
<ConnectInfo<T> as FromRequestParts<S>>
<Version as FromRequestParts<S>>
<ConnectInfo<T> as FromRequestParts<S>>
<Extensions as FromRequestParts<S>>
and $N others
= note: required for `bool` to implement `FromRequest<(), axum_core::extract::private::ViaParts>`

2
axum-macros/tests/from_request/fail/parts_extracting_body.stderr
View file

@ -13,6 +13,6 @@ error[E0277]: the trait bound `String: FromRequestParts<S>` is not satisfied
<Method as FromRequestParts<S>>
<axum::http::request::Parts as FromRequestParts<S>>
<Uri as FromRequestParts<S>>
<ConnectInfo<T> as FromRequestParts<S>>
<Version as FromRequestParts<S>>
<ConnectInfo<T> as FromRequestParts<S>>
and $N others

4
axum/benches/benches.rs
View file

@ -5,6 +5,7 @@ use axum::{
};
use serde::{Deserialize, Serialize};
use std::{
future::IntoFuture,
io::BufRead,
process::{Command, Stdio},
};
@ -161,7 +162,8 @@ impl BenchmarkBuilder {
let addr = listener.local_addr().unwrap();
std::thread::spawn(move || {
rt.block_on(axum::serve(listener, app)).unwrap();
rt.block_on(axum::serve(listener, app).into_future())
.unwrap();
});
let mut cmd = Command::new("rewrk");

148
axum/src/serve.rs
View file

@ -2,8 +2,9 @@
use std::{
convert::Infallible,
future::Future,
future::{Future, IntoFuture},
io,
marker::PhantomData,
net::SocketAddr,
pin::Pin,
task::{Context, Poll},
@ -86,48 +87,129 @@ use tower_service::Service;
/// [`HandlerWithoutStateExt::into_make_service_with_connect_info`]: crate::handler::HandlerWithoutStateExt::into_make_service_with_connect_info
/// [`HandlerService::into_make_service_with_connect_info`]: crate::handler::HandlerService::into_make_service_with_connect_info
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
pub async fn serve<M, S>(tcp_listener: TcpListener, mut make_service: M) -> io::Result<()>
pub fn serve<M, S>(tcp_listener: TcpListener, make_service: M) -> Serve<M, S>
where
M: for<'a> Service<IncomingStream<'a>, Error = Infallible, Response = S>,
S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
S::Future: Send,
{
loop {
let (tcp_stream, remote_addr) = tcp_listener.accept().await?;
let tcp_stream = TokioIo::new(tcp_stream);
Serve {
tcp_listener,
make_service,
_marker: PhantomData,
}
}
poll_fn(|cx| make_service.poll_ready(cx))
.await
.unwrap_or_else(|err| match err {});
/// Future returned by [`serve`].
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
pub struct Serve<M, S> {
tcp_listener: TcpListener,
make_service: M,
_marker: PhantomData<S>,
}
let tower_service = make_service
.call(IncomingStream {
tcp_stream: &tcp_stream,
remote_addr,
})
.await
.unwrap_or_else(|err| match err {});
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
impl<M, S> std::fmt::Debug for Serve<M, S>
where
M: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let Self {
tcp_listener,
make_service,
_marker: _,
} = self;
let hyper_service = TowerToHyperService {
service: tower_service,
};
f.debug_struct("Serve")
.field("tcp_listener", tcp_listener)
.field("make_service", make_service)
.finish()
}
}
tokio::task::spawn(async move {
match Builder::new(TokioExecutor::new())
// upgrades needed for websockets
.serve_connection_with_upgrades(tcp_stream, hyper_service)
.await
{
Ok(()) => {}
Err(_err) => {
// This error only appears when the client doesn't send a request and
// terminate the connection.
//
// If client sends one request then terminate connection whenever, it doesn't
// appear.
}
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
impl<M, S> IntoFuture for Serve<M, S>
where
M: for<'a> Service<IncomingStream<'a>, Error = Infallible, Response = S> + Send + 'static,
for<'a> <M as Service<IncomingStream<'a>>>::Future: Send,
S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
S::Future: Send,
{
type Output = io::Result<()>;
type IntoFuture = private::ServeFuture;
fn into_future(self) -> Self::IntoFuture {
private::ServeFuture(Box::pin(async move {
let Self {
tcp_listener,
mut make_service,
_marker: _,
} = self;
loop {
let (tcp_stream, remote_addr) = tcp_listener.accept().await?;
let tcp_stream = TokioIo::new(tcp_stream);
poll_fn(|cx| make_service.poll_ready(cx))
.await
.unwrap_or_else(|err| match err {});
let tower_service = make_service
.call(IncomingStream {
tcp_stream: &tcp_stream,
remote_addr,
})
.await
.unwrap_or_else(|err| match err {});
let hyper_service = TowerToHyperService {
service: tower_service,
};
tokio::task::spawn(async move {
match Builder::new(TokioExecutor::new())
// upgrades needed for websockets
.serve_connection_with_upgrades(tcp_stream, hyper_service)
.await
{
Ok(()) => {}
Err(_err) => {
// This error only appears when the client doesn't send a request and
// terminate the connection.
//
// If client sends one request then terminate connection whenever, it doesn't
// appear.
}
}
});
}
});
}))
}
}
mod private {
use std::{
future::Future,
io,
pin::Pin,
task::{Context, Poll},
};
pub struct ServeFuture(pub(super) futures_util::future::BoxFuture<'static, io::Result<()>>);
impl Future for ServeFuture {
type Output = io::Result<()>;
#[inline]
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.0.as_mut().poll(cx)
}
}
impl std::fmt::Debug for ServeFuture {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ServeFuture").finish_non_exhaustive()
}
}
}

7
examples/testing-websockets/src/main.rs
View file

@ -92,7 +92,10 @@ where
#[cfg(test)]
mod tests {
use super::*;
use std::net::{Ipv4Addr, SocketAddr};
use std::{
future::IntoFuture,
net::{Ipv4Addr, SocketAddr},
};
use tokio_tungstenite::tungstenite;
// We can integration test one handler by running the server in a background task and
@ -103,7 +106,7 @@ mod tests {
.await
.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(axum::serve(listener, app()));
tokio::spawn(axum::serve(listener, app()).into_future());
let (mut socket, _response) =
tokio_tungstenite::connect_async(format!("ws://{addr}/integration-testable"))