feat(volo-http): support extract in timeout handler (#268)

* feat(volo-http): support extract in timeout handler * chore(volo-http): re-export common types from other crates Signed-off-by: Yu Li <[email protected]>
cloudwego · Dec 1, 2023 · f0056d7 · f0056d7
1 parent e20f3c5
commit f0056d7
Show file tree

Hide file tree

Showing 8 changed files with 139 additions and 68 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
@@ -81,9 +81,6 @@ path = "src/http/simple.rs"
 anyhow.workspace = true
 async-stream.workspace = true
 bytes.workspace = true
-http = "1"
-http-body-util = "0.1"
-hyper = { version = "1", features = ["server", "http1", "http2"] }
 lazy_static.workspace = true
 metainfo.workspace = true
 motore.workspace = true

diff --git a/examples/src/http/simple.rs b/examples/src/http/simple.rs
@@ -1,16 +1,10 @@
 use std::{net::SocketAddr, time::Duration};
 
-use bytes::Bytes;
-use http::{Method, Response, StatusCode, Uri};
-use http_body_util::Full;
 use serde::Deserialize;
 use volo_http::{
     layer::TimeoutLayer,
-    param::Params,
-    request::Json,
     route::{get, post, MethodRouter, Router},
-    server::Server,
-    HttpContext,
+    Address, Bytes, Json, Method, Params, Server, StatusCode, Uri,
 };
 
 async fn hello() -> &'static str {
@@ -57,8 +51,8 @@ async fn timeout_test() {
     tokio::time::sleep(Duration::from_secs(5)).await
 }
 
-fn timeout_handler(ctx: &HttpContext) -> StatusCode {
-    tracing::info!("Timeout on `{}`, peer: {}", ctx.uri, ctx.peer);
+fn timeout_handler(uri: Uri, peer: Address) -> StatusCode {
+    tracing::info!("Timeout on `{}`, peer: {}", uri, peer);
     StatusCode::INTERNAL_SERVER_ERROR
 }
 
@@ -70,16 +64,11 @@ async fn main() {
 
     tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
 
-    let timeout_response = Response::builder()
-        .status(StatusCode::INTERNAL_SERVER_ERROR)
-        .body(Full::new(Bytes::new()))
-        .unwrap();
-
     let app = Router::new()
         .route(
             "/",
-            get(hello).layer(TimeoutLayer::new(Duration::from_secs(1), move |_| {
-                timeout_response
+            get(hello).layer(TimeoutLayer::new(Duration::from_secs(1), || {
+                StatusCode::INTERNAL_SERVER_ERROR
             })),
         )
         .route("/:echo", get(echo))

diff --git a/volo-http/src/extract.rs b/volo-http/src/extract.rs
@@ -1,5 +1,6 @@
 use futures_util::Future;
 use http::{Method, Response, Uri};
+use volo::net::Address;
 
 use crate::{response::IntoResponse, HttpContext, Params, State};
 
@@ -23,6 +24,17 @@ where
     }
 }
 
+impl<S> FromContext<S> for Address
+where
+    S: Send + Sync,
+{
+    type Rejection = Response<()>; // Infallible
+
+    async fn from_context(context: &HttpContext, _state: &S) -> Result<Address, Self::Rejection> {
+        Ok(context.peer.clone())
+    }
+}
+
 impl<S> FromContext<S> for Uri
 where
     S: Send + Sync,

diff --git a/volo-http/src/handler.rs b/volo-http/src/handler.rs
@@ -6,6 +6,7 @@ use motore::Service;
 
 use crate::{
     extract::FromContext,
+    macros::{all_the_tuples, all_the_tuples_no_last_special_case},
     request::FromRequest,
     response::{IntoResponse, RespBody},
     DynError, DynService, HttpContext,
@@ -79,31 +80,7 @@ macro_rules! impl_handler {
     };
 }
 
-impl_handler!([], T1);
-impl_handler!([T1], T2);
-impl_handler!([T1, T2], T3);
-impl_handler!([T1, T2, T3], T4);
-impl_handler!([T1, T2, T3, T4], T5);
-impl_handler!([T1, T2, T3, T4, T5], T6);
-impl_handler!([T1, T2, T3, T4, T5, T6], T7);
-impl_handler!([T1, T2, T3, T4, T5, T6, T7], T8);
-impl_handler!([T1, T2, T3, T4, T5, T6, T7, T8], T9);
-impl_handler!([T1, T2, T3, T4, T5, T6, T7, T8, T9], T10);
-impl_handler!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10], T11);
-impl_handler!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11], T12);
-impl_handler!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12], T13);
-impl_handler!(
-    [T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13],
-    T14
-);
-impl_handler!(
-    [T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14],
-    T15
-);
-impl_handler!(
-    [T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15],
-    T16
-);
+all_the_tuples!(impl_handler);
 
 // Use an extra trait with less generic types for hiding the type of handler
 pub struct DynHandler<S>(Box<dyn ErasedIntoRoute<S>>);
@@ -272,3 +249,43 @@ where
         Ok(self.handler.clone().call(cx, req, &self.state).await)
     }
 }
+
+pub trait HandlerWithoutRequest<T>: Sized {
+    fn call(self, context: &HttpContext) -> impl Future<Output = Response<RespBody>> + Send;
+}
+
+impl<F, Res> HandlerWithoutRequest<()> for F
+where
+    F: FnOnce() -> Res + Clone + Send,
+    Res: IntoResponse,
+{
+    async fn call(self, _context: &HttpContext) -> Response<RespBody> {
+        self().into_response()
+    }
+}
+
+macro_rules! impl_handler_without_request {
+    (
+        $($ty:ident),* $(,)?
+    ) => {
+        #[allow(non_snake_case, unused_mut, unused_variables)]
+        impl<F, Res, $($ty,)*> HandlerWithoutRequest<($($ty,)*)> for F
+        where
+            F: FnOnce($($ty,)*) -> Res + Clone + Send,
+            Res: IntoResponse,
+            $( for<'r> $ty: FromContext<()> + Send + 'r, )*
+        {
+            async fn call(self, context: &HttpContext) -> Response<RespBody> {
+                $(
+                    let $ty = match $ty::from_context(context, &()).await {
+                        Ok(value) => value,
+                        Err(rejection) => return rejection.into_response(),
+                    };
+                )*
+                self($($ty,)*).into_response()
+            }
+        }
+    };
+}
+
+all_the_tuples_no_last_special_case!(impl_handler_without_request);
diff --git a/volo-http/src/layer.rs b/volo-http/src/layer.rs
@@ -1,11 +1,12 @@
-use std::time::Duration;
+use std::{marker::PhantomData, time::Duration};
 
 use http::{Method, Request, Response, StatusCode};
 use http_body_util::Full;
 use hyper::body::{Bytes, Incoming};
 use motore::{layer::Layer, service::Service};
 
 use crate::{
+    handler::HandlerWithoutRequest,
     response::{IntoResponse, RespBody},
     HttpContext,
 };
@@ -93,50 +94,57 @@ where
 }
 
 #[derive(Clone)]
-pub struct TimeoutLayer<F> {
+pub struct TimeoutLayer<H, T> {
     duration: Duration,
-    handler: F,
+    handler: H,
+    _marker: PhantomData<T>,
 }
 
-impl<F> TimeoutLayer<F> {
-    pub fn new<T>(duration: Duration, handler: F) -> Self
+impl<H, T> TimeoutLayer<H, T> {
+    pub fn new(duration: Duration, handler: H) -> Self
     where
-        F: FnOnce(&HttpContext) -> T + Clone + Sync,
-        T: IntoResponse,
+        H: HandlerWithoutRequest<T> + Clone + Send + Sync + 'static,
     {
-        Self { duration, handler }
+        Self {
+            duration,
+            handler,
+            _marker: PhantomData,
+        }
     }
 }
 
-impl<S, F, T> Layer<S> for TimeoutLayer<F>
+impl<S, H, T> Layer<S> for TimeoutLayer<H, T>
 where
     S: Service<HttpContext, Incoming, Response = Response<RespBody>> + Send + Sync + 'static,
-    F: FnOnce(&HttpContext) -> T + Clone + Sync,
-    T: IntoResponse,
+    H: HandlerWithoutRequest<T> + Clone + Send + Sync + 'static,
+    T: Sync,
 {
-    type Service = Timeout<S, F>;
+    type Service = Timeout<S, H, T>;
 
     fn layer(self, inner: S) -> Self::Service {
         Timeout {
             service: inner,
             duration: self.duration,
             handler: self.handler,
+            _marker: PhantomData,
         }
     }
 }
 
 #[derive(Clone)]
-pub struct Timeout<S, F> {
+pub struct Timeout<S, H, T> {
     service: S,
     duration: Duration,
-    handler: F,
+    handler: H,
+    _marker: PhantomData<T>,
 }
 
-impl<S, F, T> Service<HttpContext, Incoming> for Timeout<S, F>
+impl<S, H, T> Service<HttpContext, Incoming> for Timeout<S, H, T>
 where
     S: Service<HttpContext, Incoming, Response = Response<RespBody>> + Send + Sync + 'static,
-    F: FnOnce(&HttpContext) -> T + Clone + Sync,
-    T: IntoResponse,
+    S::Error: Send,
+    H: HandlerWithoutRequest<T> + Clone + Send + Sync + 'static,
+    T: Sync,
 {
     type Response = S::Response;
 
@@ -153,7 +161,7 @@ where
         tokio::select! {
             resp = fut_service => resp,
             _ = fut_timeout => {
-                Ok((self.handler.clone())(cx).into_response())
+                Ok(self.handler.clone().call(cx).await.into_response())
             },
         }
     }

diff --git a/volo-http/src/lib.rs b/volo-http/src/lib.rs
@@ -7,10 +7,15 @@ pub mod response;
 pub mod route;
 pub mod server;
 
-use http::{Extensions, HeaderMap, HeaderValue, Method, Uri, Version};
+mod macros;
+
+pub use bytes::Bytes;
+use http::{Extensions, HeaderMap, HeaderValue, Version};
+pub use http::{Method, StatusCode, Uri};
 use hyper::{body::Incoming, Response};
-use param::Params;
-use volo::net::Address;
+pub use volo::net::Address;
+
+pub use crate::{param::Params, request::Json, server::Server};
 
 mod private {
     #[derive(Debug, Clone, Copy)]

diff --git a/volo-http/src/macros.rs b/volo-http/src/macros.rs
@@ -0,0 +1,46 @@
+#[rustfmt::skip]
+macro_rules! all_the_tuples {
+    ($name:ident) => {
+        $name!([], T1);
+        $name!([T1], T2);
+        $name!([T1, T2], T3);
+        $name!([T1, T2, T3], T4);
+        $name!([T1, T2, T3, T4], T5);
+        $name!([T1, T2, T3, T4, T5], T6);
+        $name!([T1, T2, T3, T4, T5, T6], T7);
+        $name!([T1, T2, T3, T4, T5, T6, T7], T8);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8], T9);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9], T10);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10], T11);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11], T12);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12], T13);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13], T14);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14], T15);
+        $name!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15], T16);
+    };
+}
+
+#[rustfmt::skip]
+macro_rules! all_the_tuples_no_last_special_case {
+    ($name:ident) => {
+        $name!(T1);
+        $name!(T1, T2);
+        $name!(T1, T2, T3);
+        $name!(T1, T2, T3, T4);
+        $name!(T1, T2, T3, T4, T5);
+        $name!(T1, T2, T3, T4, T5, T6);
+        $name!(T1, T2, T3, T4, T5, T6, T7);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
+    };
+}
+
+pub(crate) use all_the_tuples;
+pub(crate) use all_the_tuples_no_last_special_case;