Improve graceful shotdown and fix http dectect issue (#354)

salvo-rs · Aug 8, 2023 · 73090e8 · 73090e8
1 parent 67b24b8
commit 73090e8
Show file tree

Hide file tree

Showing 13 changed files with 317 additions and 66 deletions.
diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml
@@ -78,6 +78,7 @@ tokio-native-tls = { workspace = true, optional = true }
 tokio-rustls = { workspace = true, optional = true }
 tokio-openssl = { workspace = true, optional = true }
 tokio-stream = { workspace = true }
+tokio-util = { workspace = true }
 tracing = { workspace = true }
 url = { workspace = true, optional = true }
 x509-parser = { workspace = true, optional = true }

diff --git a/crates/core/src/conn/joined.rs b/crates/core/src/conn/joined.rs
@@ -3,9 +3,11 @@ use std::io::{self, Result as IoResult};
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
+use std::time::Duration;
 
 use pin_project::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio_util::sync::CancellationToken;
 
 use crate::async_trait;
 use crate::conn::Holding;
@@ -117,10 +119,22 @@ where
     A: HttpConnection + Send,
     B: HttpConnection + Send,
 {
-    async fn serve(self, handler: HyperHandler, builder: Arc<HttpBuilder>) -> IoResult<()> {
+    async fn serve(
+        self,
+        handler: HyperHandler,
+        builder: Arc<HttpBuilder>,
+        server_shutdown_token: CancellationToken,
+        idle_connection_timeout: Option<Duration>,
+    ) -> IoResult<()> {
         match self {
-            JoinedStream::A(a) => a.serve(handler, builder).await,
-            JoinedStream::B(b) => b.serve(handler, builder).await,
+            JoinedStream::A(a) => {
+                a.serve(handler, builder, server_shutdown_token, idle_connection_timeout)
+                    .await
+            }
+            JoinedStream::B(b) => {
+                b.serve(handler, builder, server_shutdown_token, idle_connection_timeout)
+                    .await
+            }
         }
     }
 }

diff --git a/crates/core/src/conn/mod.rs b/crates/core/src/conn/mod.rs
@@ -67,9 +67,11 @@ cfg_feature! {
     mod sealed {
         use std::io::{Error as IoError, ErrorKind, Result as IoResult};
         use std::sync::Arc;
+        use std::time::Duration;
 
         use tokio_rustls::server::TlsStream;
         use tokio::io::{AsyncRead, AsyncWrite};
+        use tokio_util::sync::CancellationToken;
 
         use crate::async_trait;
         use crate::service::HyperHandler;
@@ -82,9 +84,11 @@ cfg_feature! {
         where
             S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
         {
-            async fn serve(self, handler: HyperHandler, builder: Arc<HttpBuilder>) -> IoResult<()> {
+            async fn serve(self, handler: HyperHandler, builder: Arc<HttpBuilder>,
+                server_shutdown_token: CancellationToken,
+                idle_connection_timeout: Option<Duration>) -> IoResult<()> {
                 builder
-                    .serve_connection(self, handler)
+                    .serve_connection(self, handler, server_shutdown_token, idle_connection_timeout)
                     .await
                     .map_err(|e| IoError::new(ErrorKind::Other, e.to_string()))
             }

diff --git a/crates/core/src/conn/native_tls/listener.rs b/crates/core/src/conn/native_tls/listener.rs
@@ -2,13 +2,15 @@
 use std::io::{Error as IoError, ErrorKind, Result as IoResult};
 use std::sync::Arc;
 use std::task::{Context, Poll};
+use std::time::Duration;
 
 use futures_util::stream::BoxStream;
 use futures_util::task::noop_waker_ref;
 use futures_util::{Stream, StreamExt};
 use http::uri::Scheme;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_native_tls::TlsStream;
+use tokio_util::sync::CancellationToken;
 
 use crate::async_trait;
 use crate::conn::{Accepted, Acceptor, Holding, HttpBuilder, IntoConfigStream, Listener};
@@ -60,9 +62,15 @@ impl<S> HttpConnection for TlsStream<S>
 where
     S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 {
-    async fn serve(self, handler: HyperHandler, builder: Arc<HttpBuilder>) -> IoResult<()> {
+    async fn serve(
+        self,
+        handler: HyperHandler,
+        builder: Arc<HttpBuilder>,
+        server_shutdown_token: CancellationToken,
+        idle_connection_timeout: Option<Duration>,
+    ) -> IoResult<()> {
         builder
-            .serve_connection(self, handler)
+            .serve_connection(self, handler, server_shutdown_token, idle_connection_timeout)
             .await
             .map_err(|e| IoError::new(ErrorKind::Other, e.to_string()))
     }
@@ -163,16 +171,16 @@ where
             http_version,
             http_scheme,
         } = self.inner.accept().await?;
-            let conn = tls_acceptor
-                .accept(conn)
-                .await
-                .map_err(|e| IoError::new(ErrorKind::Other, e.to_string()))?;
-            Ok(Accepted {
-                conn,
-                local_addr,
-                remote_addr,
-                http_version,
-                http_scheme,
-            })
+        let conn = tls_acceptor
+            .accept(conn)
+            .await
+            .map_err(|e| IoError::new(ErrorKind::Other, e.to_string()))?;
+        Ok(Accepted {
+            conn,
+            local_addr,
+            remote_addr,
+            http_version,
+            http_scheme,
+        })
     }
 }
diff --git a/crates/core/src/conn/openssl/listener.rs b/crates/core/src/conn/openssl/listener.rs
@@ -1,7 +1,7 @@
 //! openssl module
 use std::io::{Error as IoError, Result as IoResult};
 use std::sync::Arc;
-use std::task::{Context, Poll};
+use std::task::{Context, Poll};use std::time::Duration;
 
 use futures_util::stream::BoxStream;
 use futures_util::task::noop_waker_ref;
@@ -10,7 +10,7 @@ use http::uri::Scheme;
 use openssl::ssl::{Ssl, SslAcceptor};
 use tokio::io::ErrorKind;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tokio_openssl::SslStream;
+use tokio_openssl::SslStream;use tokio_util::sync::CancellationToken;
 
 use super::OpensslConfig;
 
@@ -105,9 +105,15 @@ impl<S> HttpConnection for SslStream<S>
 where
     S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
 {
-    async fn serve(self, handler: HyperHandler, builder: Arc<HttpBuilder>) -> IoResult<()> {
+    async fn serve(
+        self,
+        handler: HyperHandler,
+        builder: Arc<HttpBuilder>,
+        server_shutdown_token: CancellationToken,
+        idle_connection_timeout: Option<Duration>,
+    ) -> IoResult<()> {
         builder
-            .serve_connection(self, handler)
+            .serve_connection(self, handler, server_shutdown_token, idle_connection_timeout)
             .await
             .map_err(|e| IoError::new(ErrorKind::Other, e.to_string()))
     }