OpenSource/axum
OpenSource/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2024-11-25 00:28:07 +01:00
Code Issues Projects Releases Packages Wiki Activity

better readiness handling and less boxing

Browse source
This commit is contained in:
David Pedersen 2021-05-30 01:11:18 +02:00
parent 33f2e5f661
commit 0513b56faf
2 changed files with 90 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
Cargo.toml
View file

@ -7,15 +7,16 @@ edition = "2018"
[dependencies] [dependencies]
async-trait = "0.1" async-trait = "0.1"
bytes = "1.0" bytes = "1.0"
futures-util = "0.3"
http = "0.2" http = "0.2"
http-body = "0.4" http-body = "0.4"
hyper = "0.14" hyper = "0.14"
pin-project = "1.0"
serde = "1.0" serde = "1.0"
serde_urlencoded = "0.7"
serde_json = "1.0" serde_json = "1.0"
futures-util = "0.3" serde_urlencoded = "0.7"
tower = { version = "0.4", features = ["util"] }
thiserror = "1.0" thiserror = "1.0"
tower = { version = "0.4", features = ["util"] }
[dev-dependencies] [dev-dependencies]
tokio = { version = "1.6.1", features = ["macros", "rt"] } tokio = { version = "1.6.1", features = ["macros", "rt"] }

86
src/lib.rs
View file

@ -31,10 +31,12 @@ use async_trait::async_trait;
use bytes::Bytes; use bytes::Bytes;
use futures_util::{future, ready}; use futures_util::{future, ready};
use http::{Method, Request, Response, StatusCode}; use http::{Method, Request, Response, StatusCode};
use pin_project::pin_project;
use serde::{de::DeserializeOwned, Deserialize, Serialize}; use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::{ use std::{
future::Future, future::Future,
marker::PhantomData, marker::PhantomData,
pin::Pin,
task::{Context, Poll}, task::{Context, Poll},
}; };
use tower::{Service, ServiceExt}; use tower::{Service, ServiceExt};
@ -155,21 +157,20 @@ pub enum Error {
DeserializeQueryString(#[from] serde_urlencoded::de::Error), DeserializeQueryString(#[from] serde_urlencoded::de::Error),
} }
// TODO(david): make this trait sealed
#[async_trait] #[async_trait]
pub trait Handler<Out> { pub trait Handler<Out> {
async fn call(self, req: Request<Body>) -> Result<Response<Body>, Error>; async fn call(self, req: Request<Body>) -> Result<Response<Body>, Error>;
} }
#[async_trait] #[async_trait]
#[allow(non_snake_case)]
impl<F, Fut> Handler<()> for F impl<F, Fut> Handler<()> for F
where where
F: Fn(Request<Body>) -> Fut + Send + Sync, F: Fn(Request<Body>) -> Fut + Send + Sync,
Fut: Future<Output = Result<Response<Body>, Error>> + Send, Fut: Future<Output = Result<Response<Body>, Error>> + Send,
{ {
async fn call(self, req: Request<Body>) -> Result<Response<Body>, Error> { async fn call(self, req: Request<Body>) -> Result<Response<Body>, Error> {
let res = self(req).await?; self(req).await
Ok(res)
} }
} }
@ -253,18 +254,35 @@ where
} }
} }
#[async_trait]
pub trait FromRequest: Sized { pub trait FromRequest: Sized {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error>; type Future: Future<Output = Result<Self, Error>> + Send;
fn from_request(req: &mut Request<Body>) -> Self::Future;
} }
#[async_trait]
impl<T> FromRequest for Option<T> impl<T> FromRequest for Option<T>
where where
T: FromRequest, T: FromRequest,
{ {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> { type Future = OptionFromRequestFuture<T::Future>;
Ok(T::from_request(req).await.ok())
fn from_request(req: &mut Request<Body>) -> Self::Future {
OptionFromRequestFuture(T::from_request(req))
}
}
#[pin_project]
pub struct OptionFromRequestFuture<F>(#[pin] F);
impl<F, T> Future for OptionFromRequestFuture<F>
where
F: Future<Output = Result<T, Error>>,
{
type Output = Result<Option<T>, Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let value = ready!(self.project().0.poll(cx));
Poll::Ready(Ok(value.ok()))
} }
} }
@ -277,15 +295,20 @@ impl<T> Query<T> {
} }
} }
#[async_trait]
impl<T> FromRequest for Query<T> impl<T> FromRequest for Query<T>
where where
T: DeserializeOwned, T: DeserializeOwned + Send,
{ {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> { type Future = future::Ready<Result<Self, Error>>;
fn from_request(req: &mut Request<Body>) -> Self::Future {
let result = (|| {
let query = req.uri().query().ok_or(Error::QueryStringMissing)?; let query = req.uri().query().ok_or(Error::QueryStringMissing)?;
let value = serde_urlencoded::from_str(query)?; let value = serde_urlencoded::from_str(query)?;
Ok(Query(value)) Ok(Query(value))
})();
future::ready(result)
} }
} }
@ -298,21 +321,24 @@ impl<T> Json<T> {
} }
} }
#[async_trait]
impl<T> FromRequest for Json<T> impl<T> FromRequest for Json<T>
where where
T: DeserializeOwned, T: DeserializeOwned,
{ {
async fn from_request(req: &mut Request<Body>) -> Result<Self, Error> { type Future = future::BoxFuture<'static, Result<Self, Error>>;
fn from_request(req: &mut Request<Body>) -> Self::Future {
// TODO(david): require the body to have `content-type: application/json` // TODO(david): require the body to have `content-type: application/json`
let body = std::mem::take(req.body_mut()); let body = std::mem::take(req.body_mut());
Box::pin(async move {
let bytes = hyper::body::to_bytes(body) let bytes = hyper::body::to_bytes(body)
.await .await
.map_err(Error::ConsumeBody)?; .map_err(Error::ConsumeBody)?;
let value = serde_json::from_slice(&bytes).map_err(Error::DeserializeRequestBody)?; let value = serde_json::from_slice(&bytes).map_err(Error::DeserializeRequestBody)?;
Ok(Json(value)) Ok(Json(value))
})
} }
} }
@ -331,11 +357,10 @@ impl<R> Service<R> for EmptyRouter {
fn call(&mut self, _req: R) -> Self::Future { fn call(&mut self, _req: R) -> Self::Future {
let mut res = Response::new(Body::empty()); let mut res = Response::new(Body::empty());
*res.status_mut() = StatusCode::NOT_FOUND; *res.status_mut() = StatusCode::NOT_FOUND;
future::ready(Ok(res)) future::ok(res)
} }
} }
#[derive(Clone)]
pub struct Route<H, F> { pub struct Route<H, F> {
handler: H, handler: H,
route_spec: RouteSpec, route_spec: RouteSpec,
@ -344,6 +369,23 @@ pub struct Route<H, F> {
fallback_ready: bool, fallback_ready: bool,
} }
impl<H, F> Clone for Route<H, F>
where
H: Clone,
F: Clone,
{
fn clone(&self) -> Self {
Self {
handler: self.handler.clone(),
fallback: self.fallback.clone(),
route_spec: self.route_spec.clone(),
// important to reset readiness when cloning
handler_ready: false,
fallback_ready: false,
}
}
}
#[derive(Clone)] #[derive(Clone)]
struct RouteSpec { struct RouteSpec {
method: Method, method: Method,
@ -367,6 +409,7 @@ where
type Future = future::Either<H::Future, F::Future>; type Future = future::Either<H::Future, F::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
loop {
if !self.handler_ready { if !self.handler_ready {
ready!(self.handler.poll_ready(cx))?; ready!(self.handler.poll_ready(cx))?;
self.handler_ready = true; self.handler_ready = true;
@ -377,14 +420,25 @@ where
self.fallback_ready = true; self.fallback_ready = true;
} }
Poll::Ready(Ok(())) if self.handler_ready && self.fallback_ready {
return Poll::Ready(Ok(()));
}
}
} }
fn call(&mut self, req: Request<Body>) -> Self::Future { fn call(&mut self, req: Request<Body>) -> Self::Future {
if self.route_spec.matches(&req) { if self.route_spec.matches(&req) {
assert!(
self.handler_ready,
"handler not ready. Did you forget to call `poll_ready`?"
);
self.handler_ready = false; self.handler_ready = false;
future::Either::Left(self.handler.call(req)) future::Either::Left(self.handler.call(req))
} else { } else {
assert!(
self.fallback_ready,
"fallback not ready. Did you forget to call `poll_ready`?"
);
self.fallback_ready = false; self.fallback_ready = false;
future::Either::Right(self.fallback.call(req)) future::Either::Right(self.fallback.call(req))
} }