Http gracefull shutdown support

ntex-io/CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [2.2.0] - 2024-08-12
+
+* Allow to notify dispatcher from IoRef
+
 ## [2.1.0] - 2024-07-30
 
 * Optimize `Io` layout

ntex-io/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex-io"
-version = "2.1.0"
+version = "2.2.0"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Utilities for encoding and decoding frames"
 keywords = ["network", "framework", "async", "futures"]

ntex-io/src/ioref.rs

@@ -217,17 +217,24 @@ impl IoRef {
     }
 
     #[inline]
-    /// current timer handle
-    pub fn timer_handle(&self) -> timer::TimerHandle {
-        self.0.timeout.get()
+    /// Wakeup dispatcher
+    pub fn notify_dispatcher(&self) {
+        self.0.dispatch_task.wake();
+        log::trace!("{}: Timer, notify dispatcher", self.tag());
     }
 
     #[inline]
-    /// wakeup dispatcher and send keep-alive error
+    /// Wakeup dispatcher and send keep-alive error
     pub fn notify_timeout(&self) {
         self.0.notify_timeout()
     }
 
+    #[inline]
+    /// current timer handle
+    pub fn timer_handle(&self) -> timer::TimerHandle {
+        self.0.timeout.get()
+    }
+
     #[inline]
     /// Start timer
     pub fn start_timer(&self, timeout: Seconds) -> timer::TimerHandle {

ntex/CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [2.2.0] - 2024-08-12
+
+* Http gracefull shutdown support
+
 ## [2.1.0] - 2024-07-30
 
 * Better handling for connection upgrade #385

ntex/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex"
-version = "2.1.0"
+version = "2.2.0"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Framework for composable network services"
 readme = "README.md"
@@ -65,10 +65,10 @@ ntex-service = "3.0"
 ntex-macros = "0.1.3"
 ntex-util = "2"
 ntex-bytes = "0.1.27"
-ntex-server = "2.1"
+ntex-server = "2.3"
 ntex-h2 = "1.0"
 ntex-rt = "0.4.13"
-ntex-io = "2.1"
+ntex-io = "2.2"
 ntex-net = "2.0"
 ntex-tls = "2.0"
 

ntex/src/http/config.rs

@@ -234,13 +234,23 @@ impl ServiceConfig {
     }
 }
 
+bitflags::bitflags! {
+    #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
+    struct Flags: u8 {
+        /// Keep-alive enabled
+        const KA_ENABLED = 0b0000_0001;
+        /// Shutdown service
+        const SHUTDOWN   = 0b0000_0010;
+    }
+}
+
 pub(super) struct DispatcherConfig<S, C> {
+    flags: Cell<Flags>,
     pub(super) service: Pipeline<S>,
     pub(super) control: Pipeline<C>,
     pub(super) keep_alive: Seconds,
     pub(super) client_disconnect: Seconds,
     pub(super) h2config: h2::Config,
-    pub(super) ka_enabled: bool,
     pub(super) headers_read_rate: Option<ReadRate>,
     pub(super) payload_read_rate: Option<ReadRate>,
     pub(super) timer: DateService,
@@ -253,22 +263,37 @@ impl<S, C> DispatcherConfig<S, C> {
             control: control.into(),
             keep_alive: cfg.keep_alive,
             client_disconnect: cfg.client_disconnect,
-            ka_enabled: cfg.ka_enabled,
             headers_read_rate: cfg.headers_read_rate,
             payload_read_rate: cfg.payload_read_rate,
             h2config: cfg.h2config.clone(),
             timer: cfg.timer.clone(),
+            flags: Cell::new(if cfg.ka_enabled {
+                Flags::KA_ENABLED
+            } else {
+                Flags::empty()
+            }),
         }
     }
 
     /// Return state of connection keep-alive functionality
     pub(super) fn keep_alive_enabled(&self) -> bool {
-        self.ka_enabled
+        self.flags.get().contains(Flags::KA_ENABLED)
     }
 
     pub(super) fn headers_read_rate(&self) -> Option<&ReadRate> {
         self.headers_read_rate.as_ref()
     }
+
+    /// Service is shuting down
+    pub(super) fn is_shutdown(&self) -> bool {
+        self.flags.get().contains(Flags::SHUTDOWN)
+    }
+
+    pub(super) fn shutdown(&self) {
+        let mut flags = self.flags.get();
+        flags.insert(Flags::SHUTDOWN);
+        self.flags.set(flags);
+    }
 }
 
 const DATE_VALUE_LENGTH_HDR: usize = 39;

ntex/src/http/h1/dispatcher.rs

@@ -223,6 +223,12 @@ where
     B: MessageBody,
 {
     fn poll_read_request(&mut self, cx: &mut Context<'_>) -> Poll<State<F, C, S, B>> {
+        // stop dispatcher
+        if self.config.is_shutdown() {
+            log::trace!("{}: Service is shutting down", self.io.tag());
+            return Poll::Ready(self.stop());
+        }
+
         log::trace!("{}: Trying to read http message", self.io.tag());
 
         let result = match self.io.poll_recv_decode(&self.codec, cx) {

ntex/src/http/h1/service.rs

@@ -1,12 +1,12 @@
-use std::{error::Error, fmt, marker, rc::Rc};
+use std::{cell::Cell, cell::RefCell, error::Error, fmt, marker, rc::Rc};
 
 use crate::http::body::MessageBody;
 use crate::http::config::{DispatcherConfig, ServiceConfig};
 use crate::http::error::{DispatchError, ResponseError};
 use crate::http::{request::Request, response::Response};
-use crate::io::{types, Filter, Io};
+use crate::io::{types, Filter, Io, IoRef};
 use crate::service::{IntoServiceFactory, Service, ServiceCtx, ServiceFactory};
-use crate::util::join;
+use crate::{channel::oneshot, util::join, util::HashSet};
 
 use super::control::{Control, ControlAck};
 use super::default::DefaultControlService;
@@ -181,10 +181,14 @@ where
             .await
             .map_err(|e| log::error!("Cannot construct control service: {:?}", e))?;
 
+        let (tx, rx) = oneshot::channel();
         let config = Rc::new(DispatcherConfig::new(self.cfg.clone(), service, control));
 
         Ok(H1ServiceHandler {
             config,
+            inflight: RefCell::new(Default::default()),
+            rx: Cell::new(Some(rx)),
+            tx: Cell::new(Some(tx)),
             _t: marker::PhantomData,
         })
     }
@@ -193,6 +197,9 @@ where
 /// `Service` implementation for HTTP1 transport
 pub struct H1ServiceHandler<F, S, B, C> {
     config: Rc<DispatcherConfig<S, C>>,
+    inflight: RefCell<HashSet<IoRef>>,
+    rx: Cell<Option<oneshot::Receiver<()>>>,
+    tx: Cell<Option<oneshot::Sender<()>>>,
     _t: marker::PhantomData<(F, B)>,
 }
 
@@ -224,18 +231,62 @@ where
     }
 
     async fn shutdown(&self) {
-        self.config.control.shutdown().await;
-        self.config.service.shutdown().await;
+        self.config.shutdown();
+
+        join(
+            self.config.control.shutdown(),
+            self.config.service.shutdown(),
+        )
+        .await;
+
+        // check inflight connections
+        let inflight = {
+            let inflight = self.inflight.borrow();
+            for io in inflight.iter() {
+                io.notify_dispatcher();
+            }
+            inflight.len()
+        };
+        if inflight != 0 {
+            log::trace!("Shutting down service, in-flight connections: {}", inflight);
+
+            if let Some(rx) = self.rx.take() {
+                let _ = rx.await;
+            }
+
+            log::trace!("Shutting down is complected",);
+        }
     }
 
     async fn call(&self, io: Io<F>, _: ServiceCtx<'_, Self>) -> Result<(), Self::Error> {
-        log::trace!(
-            "New http1 connection, peer address {:?}",
-            io.query::<types::PeerAddr>().get()
-        );
+        let inflight = {
+            let mut inflight = self.inflight.borrow_mut();
+            inflight.insert(io.get_ref());
+            inflight.len()
+        };
 
-        Dispatcher::new(io, self.config.clone())
+        log::trace!(
+            "New http1 connection, peer address {:?}, inflight: {}",
+            io.query::<types::PeerAddr>().get(),
+            inflight
+        );
+        let ioref = io.get_ref();
+
+        let result = Dispatcher::new(io, self.config.clone())
             .await
-            .map_err(DispatchError::Control)
+            .map_err(DispatchError::Control);
+
+        {
+            let mut inflight = self.inflight.borrow_mut();
+            inflight.remove(&ioref);
+
+            if inflight.len() == 0 {
+                if let Some(tx) = self.tx.take() {
+                    let _ = tx.send(());
+                }
+            }
+        }
+
+        result
     }
 }

ntex/src/http/service.rs

@@ -1,8 +1,8 @@
-use std::{error, fmt, marker, rc::Rc};
+use std::{cell::Cell, cell::RefCell, error, fmt, marker, rc::Rc};
 
-use crate::io::{types, Filter, Io};
+use crate::io::{types, Filter, Io, IoRef};
 use crate::service::{IntoServiceFactory, Service, ServiceCtx, ServiceFactory};
-use crate::util::join;
+use crate::{channel::oneshot, util::join, util::HashSet};
 
 use super::body::MessageBody;
 use super::builder::HttpServiceBuilder;
@@ -257,11 +257,15 @@ where
             .await
             .map_err(|e| log::error!("Cannot construct control service: {:?}", e))?;
 
+        let (tx, rx) = oneshot::channel();
         let config = DispatcherConfig::new(self.cfg.clone(), service, control);
 
         Ok(HttpServiceHandler {
             config: Rc::new(config),
             h2_control: self.h2_control.clone(),
+            inflight: RefCell::new(HashSet::default()),
+            rx: Cell::new(Some(rx)),
+            tx: Cell::new(Some(tx)),
             _t: marker::PhantomData,
         })
     }
@@ -271,6 +275,9 @@ where
 pub struct HttpServiceHandler<F, S, B, C1, C2> {
     config: Rc<DispatcherConfig<S, C1>>,
     h2_control: Rc<C2>,
+    inflight: RefCell<HashSet<IoRef>>,
+    rx: Cell<Option<oneshot::Receiver<()>>>,
+    tx: Cell<Option<oneshot::Sender<()>>>,
     _t: marker::PhantomData<(F, B)>,
 }
 
@@ -306,8 +313,31 @@ where
 
     #[inline]
     async fn shutdown(&self) {
-        self.config.control.shutdown().await;
-        self.config.service.shutdown().await;
+        self.config.shutdown();
+
+        join(
+            self.config.control.shutdown(),
+            self.config.service.shutdown(),
+        )
+        .await;
+
+        // check inflight connections
+        let inflight = {
+            let inflight = self.inflight.borrow();
+            for io in inflight.iter() {
+                io.notify_dispatcher();
+            }
+            inflight.len()
+        };
+        if inflight != 0 {
+            log::trace!("Shutting down service, in-flight connections: {}", inflight);
+
+            if let Some(rx) = self.rx.take() {
+                let _ = rx.await;
+            }
+
+            log::trace!("Shutting down is complected",);
+        }
     }
 
     async fn call(
@@ -315,12 +345,22 @@ where
         io: Io<F>,
         _: ServiceCtx<'_, Self>,
     ) -> Result<Self::Response, Self::Error> {
-        log::trace!(
-            "New http connection, peer address {:?}",
-            io.query::<types::PeerAddr>().get()
-        );
+        let inflight = {
+            let mut inflight = self.inflight.borrow_mut();
+            inflight.insert(io.get_ref());
+            inflight.len()
+        };
 
-        if io.query::<types::HttpProtocol>().get() == Some(types::HttpProtocol::Http2) {
+        log::trace!(
+            "New http connection, peer address {:?}, in-flight: {}",
+            io.query::<types::PeerAddr>().get(),
+            inflight
+        );
+        let ioref = io.get_ref();
+
+        let result = if io.query::<types::HttpProtocol>().get()
+            == Some(types::HttpProtocol::Http2)
+        {
             let control = self.h2_control.create(()).await.map_err(|e| {
                 DispatchError::Control(
                     format!("Cannot construct control service: {:?}", e).into(),
@@ -331,6 +371,19 @@ where
             h1::Dispatcher::new(io, self.config.clone())
                 .await
                 .map_err(DispatchError::Control)
+        };
+
+        {
+            let mut inflight = self.inflight.borrow_mut();
+            inflight.remove(&ioref);
+
+            if inflight.len() == 0 {
+                if let Some(tx) = self.tx.take() {
+                    let _ = tx.send(());
+                }
+            }
         }
+
+        result
     }
 }

ntex/src/http/test.rs

@@ -7,6 +7,7 @@ use coo_kie::{Cookie, CookieJar};
 #[cfg(feature = "ws")]
 use crate::io::Filter;
 use crate::io::Io;
+use crate::server::Server;
 #[cfg(feature = "ws")]
 use crate::ws::{error::WsClientError, WsClient, WsConnection};
 use crate::{rt::System, service::ServiceFactory};
@@ -237,18 +238,18 @@ where
 
         let system = sys.system();
         sys.run(move || {
-            crate::server::build()
+            let srv = crate::server::build()
                 .listen("test", tcp, move |_| factory())?
                 .set_tag("test", "HTTP-TEST-SRV")
                 .workers(1)
                 .disable_signals()
                 .run();
-            tx.send((system, local_addr)).unwrap();
+            tx.send((system, srv, local_addr)).unwrap();
             Ok(())
         })
     });
 
-    let (system, addr) = rx.recv().unwrap();
+    let (system, server, addr) = rx.recv().unwrap();
 
     let client = {
         let connector = {
@@ -286,6 +287,7 @@ where
         addr,
         client,
         system,
+        server,
     }
 }
 
@@ -295,6 +297,7 @@ pub struct TestServer {
     addr: net::SocketAddr,
     client: Client,
     system: System,
+    server: Server,
 }
 
 impl TestServer {
@@ -402,13 +405,14 @@ impl TestServer {
     }
 
     /// Stop http server
-    fn stop(&mut self) {
+    pub async fn stop(&self) {
+        self.server.stop(true).await;
         self.system.stop();
     }
 }
 
 impl Drop for TestServer {
     fn drop(&mut self) {
-        self.stop()
+        self.system.stop();
     }
 }

ntex/tests/http_server.rs

@@ -9,7 +9,7 @@ use ntex::http::header::{self, HeaderName, HeaderValue};
 use ntex::http::{body, h1::Control, test::server as test_server};
 use ntex::http::{HttpService, KeepAlive, Method, Request, Response, StatusCode, Version};
 use ntex::time::{sleep, timeout, Millis, Seconds};
-use ntex::{service::fn_service, util::Bytes, util::Ready, web::error};
+use ntex::{rt, service::fn_service, util::Bytes, util::Ready, web::error};
 
 #[ntex::test]
 async fn test_h1() {
@@ -724,3 +724,77 @@ async fn test_h1_client_drop() -> io::Result<()> {
     assert_eq!(count.load(Ordering::Relaxed), 1);
     Ok(())
 }
+
+#[ntex::test]
+async fn test_h1_gracefull_shutdown() {
+    let count = Arc::new(AtomicUsize::new(0));
+    let count2 = count.clone();
+
+    let srv = test_server(move || {
+        let count = count2.clone();
+        HttpService::build().h1(move |_: Request| {
+            let count = count.clone();
+            count.fetch_add(1, Ordering::Relaxed);
+            async move {
+                sleep(Millis(1000)).await;
+                count.fetch_sub(1, Ordering::Relaxed);
+                Ok::<_, io::Error>(Response::Ok().finish())
+            }
+        })
+    });
+
+    let mut stream1 = net::TcpStream::connect(srv.addr()).unwrap();
+    let _ = stream1.write_all(b"GET /index.html HTTP/1.1\r\n\r\n");
+
+    let mut stream2 = net::TcpStream::connect(srv.addr()).unwrap();
+    let _ = stream2.write_all(b"GET /index.html HTTP/1.1\r\n\r\n");
+
+    sleep(Millis(150)).await;
+    assert_eq!(count.load(Ordering::Relaxed), 2);
+
+    rt::spawn(async move {
+        srv.stop().await;
+    });
+    sleep(Millis(150)).await;
+    assert_eq!(count.load(Ordering::Relaxed), 2);
+
+    sleep(Millis(1100)).await;
+    assert_eq!(count.load(Ordering::Relaxed), 0);
+}
+
+#[ntex::test]
+async fn test_h1_gracefull_shutdown_2() {
+    let count = Arc::new(AtomicUsize::new(0));
+    let count2 = count.clone();
+
+    let srv = test_server(move || {
+        let count = count2.clone();
+        HttpService::build().finish(move |_: Request| {
+            let count = count.clone();
+            count.fetch_add(1, Ordering::Relaxed);
+            async move {
+                sleep(Millis(1000)).await;
+                count.fetch_sub(1, Ordering::Relaxed);
+                Ok::<_, io::Error>(Response::Ok().finish())
+            }
+        })
+    });
+
+    let mut stream1 = net::TcpStream::connect(srv.addr()).unwrap();
+    let _ = stream1.write_all(b"GET /index.html HTTP/1.1\r\n\r\n");
+
+    let mut stream2 = net::TcpStream::connect(srv.addr()).unwrap();
+    let _ = stream2.write_all(b"GET /index.html HTTP/1.1\r\n\r\n");
+
+    sleep(Millis(150)).await;
+    assert_eq!(count.load(Ordering::Relaxed), 2);
+
+    rt::spawn(async move {
+        srv.stop().await;
+    });
+    sleep(Millis(150)).await;
+    assert_eq!(count.load(Ordering::Relaxed), 2);
+
+    sleep(Millis(1100)).await;
+    assert_eq!(count.load(Ordering::Relaxed), 0);
+}