layer (#224)

cloudwego · Sep 11, 2023 · 1154315 · 1154315
1 parent f9e6cc0
commit 1154315
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 22 deletions.
diff --git a/examples/src/http/http.rs b/examples/src/http/http.rs
@@ -3,12 +3,12 @@ use std::{convert::Infallible, net::SocketAddr};
 use bytes::Bytes;
 use http::{Method, Response, StatusCode, Uri};
 use hyper::body::Incoming;
-use motore::service::service_fn;
+use motore::{service::service_fn, timeout::TimeoutLayer};
 use serde::{Deserialize, Serialize};
 use volo_http::{
     handler::HandlerService,
     request::Json,
-    route::{Route, Router},
+    route::{Route, Router, Server, ServiceLayerExt},
     HttpContext,
 };
 
@@ -69,8 +69,14 @@ async fn test(
 
 #[tokio::main(flavor = "multi_thread")]
 async fn main() {
-    Router::build()
-        .route("/", Route::builder().get(service_fn(hello)).build())
+    Router::new()
+        .route(
+            "/",
+            Route::builder()
+                .get(service_fn(hello))
+                .build()
+                .layer(TimeoutLayer::new(Some(std::time::Duration::from_secs(1)))),
+        )
         .route("/:echo", Route::builder().get(service_fn(echo)).build())
         .route("/user", Route::builder().post(service_fn(json)).build())
         .route(
@@ -80,6 +86,7 @@ async fn main() {
                 .post(HandlerService::new(test))
                 .build(),
         )
+        .layer(TimeoutLayer::new(Some(std::time::Duration::from_secs(1))))
         .serve(SocketAddr::from(([127, 0, 0, 1], 3000)))
         .await
         .unwrap();

diff --git a/volo-http/src/route.rs b/volo-http/src/route.rs
@@ -3,10 +3,11 @@ use std::{future::Future, net::SocketAddr};
 use http::{Method, Response, StatusCode};
 use http_body_util::Full;
 use hyper::{
-    body::{Bytes, Incoming},
+    body::{Body, Bytes, Incoming},
     server::conn::http1,
 };
 use hyper_util::rt::TokioIo;
+use motore::layer::Layer;
 use tokio::net::TcpListener;
 
 use crate::{
@@ -16,17 +17,11 @@ use crate::{
 
 pub type DynService = motore::BoxCloneService<HttpContext, Incoming, Response<RespBody>, DynError>;
 
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct Router {
     inner: matchit::Router<DynService>,
 }
 
-impl Router {
-    pub fn build() -> RouterBuilder {
-        Default::default()
-    }
-}
-
 impl motore::Service<(), (HttpContextInner, Incoming)> for Router {
     type Response = Response<RespBody>;
 
@@ -68,12 +63,7 @@ impl motore::Service<(), (HttpContextInner, Incoming)> for Router {
     }
 }
 
-#[derive(Default)]
-pub struct RouterBuilder {
-    routes: matchit::Router<DynService>,
-}
-
-impl RouterBuilder {
+impl Router {
     pub fn new() -> Self {
         Default::default()
     }
@@ -87,18 +77,50 @@ impl RouterBuilder {
             + Clone
             + 'static,
     {
-        if let Err(e) = self.routes.insert(uri, motore::BoxCloneService::new(route)) {
+        if let Err(e) = self.inner.insert(uri, motore::BoxCloneService::new(route)) {
             panic!("routing error: {e}");
         }
         self
     }
+}
+
+pub trait ServiceLayerExt: Sized {
+    fn layer<L>(self, l: L) -> L::Service
+    where
+        L: Layer<Self>;
+}
+
+impl<S> ServiceLayerExt for S {
+    fn layer<L>(self, l: L) -> L::Service
+    where
+        L: Layer<Self>,
+    {
+        Layer::layer(l, self)
+    }
+}
 
-    pub async fn serve(self, addr: SocketAddr) -> Result<(), DynError> {
+#[async_trait::async_trait]
+pub trait Server {
+    async fn serve(self, addr: SocketAddr) -> Result<(), DynError>;
+}
+#[async_trait::async_trait]
+impl<S, OB> Server for S
+where
+    S: motore::Service<(), (HttpContextInner, Incoming), Response = Response<OB>>
+        + Clone
+        + Send
+        + Sync
+        + 'static,
+    OB: Body<Error = DynError> + Send + 'static,
+    <OB as Body>::Data: Send,
+    <S as motore::Service<(), (HttpContextInner, Incoming)>>::Error: Into<DynError>,
+{
+    async fn serve(self, addr: SocketAddr) -> Result<(), DynError> {
         let listener = TcpListener::bind(addr).await?;
-        let router = Router { inner: self.routes };
 
+        let service = self;
         loop {
-            let s = router.clone();
+            let s = service.clone();
             let (stream, peer) = listener.accept().await?;
 
             let io = TokioIo::new(stream);