From c96b729207448ccae09176e06b35b60c7215e833 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Wed, 1 May 2024 17:26:05 +0500
Subject: [PATCH] Fix KeepAlive timeout handling in default dispatcher

---
 ntex-io/CHANGES.md        |  6 ++++++
 ntex-io/Cargo.toml        |  2 +-
 ntex-io/src/dispatcher.rs | 18 +++++++++++-------
 ntex-io/src/io.rs         |  3 +--
 ntex-io/src/ioref.rs      |  6 ++++++
 ntex/Cargo.toml           |  2 +-
 ntex/examples/client.rs   | 18 +++++-------------
 7 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/ntex-io/CHANGES.md b/ntex-io/CHANGES.md
index ff95c38ea..34b537ad5 100644
--- a/ntex-io/CHANGES.md
+++ b/ntex-io/CHANGES.md
@@ -1,5 +1,11 @@
 # Changes
 
+## [1.1.0] - 2024-05-01
+
+* Add IoRef::notify_timeout() helper method
+
+* Fix KeepAlive timeout handling in default dispatcher
+
 ## [1.0.2] - 2024-03-31
 
 * Add IoRef::is_wr_backpressure() method
diff --git a/ntex-io/Cargo.toml b/ntex-io/Cargo.toml
index a75780cff..b0e51e784 100644
--- a/ntex-io/Cargo.toml
+++ b/ntex-io/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ntex-io"
-version = "1.0.2"
+version = "1.1.0"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Utilities for encoding and decoding frames"
 keywords = ["network", "framework", "async", "futures"]
diff --git a/ntex-io/src/dispatcher.rs b/ntex-io/src/dispatcher.rs
index 84f40c12d..ede2af34b 100644
--- a/ntex-io/src/dispatcher.rs
+++ b/ntex-io/src/dispatcher.rs
@@ -575,15 +575,19 @@ where
                         self.shared.io.tag()
                     );
                 }
-                return Err(DispatchItem::ReadTimeout);
+                Err(DispatchItem::ReadTimeout)
+            } else {
+                Ok(())
             }
+        } else if self.flags.contains(Flags::KA_TIMEOUT) {
+            log::trace!(
+                "{}: Keep-alive error, stopping dispatcher",
+                self.shared.io.tag()
+            );
+            Err(DispatchItem::KeepAliveTimeout)
+        } else {
+            Ok(())
         }
-
-        log::trace!(
-            "{}: Keep-alive error, stopping dispatcher",
-            self.shared.io.tag()
-        );
-        Err(DispatchItem::KeepAliveTimeout)
     }
 }
 
diff --git a/ntex-io/src/io.rs b/ntex-io/src/io.rs
index eaa20f0a5..9db8ae6d7 100644
--- a/ntex-io/src/io.rs
+++ b/ntex-io/src/io.rs
@@ -93,13 +93,12 @@ impl IoState {
     }
 
     pub(super) fn notify_timeout(&self) {
-        log::trace!("{}: Timeout, notify dispatcher", self.tag.get());
-
         let mut flags = self.flags.get();
         if !flags.contains(Flags::DSP_TIMEOUT) {
             flags.insert(Flags::DSP_TIMEOUT);
             self.flags.set(flags);
             self.dispatch_task.wake();
+            log::trace!("{}: Timer, notify dispatcher", self.tag.get());
         }
     }
 
diff --git a/ntex-io/src/ioref.rs b/ntex-io/src/ioref.rs
index f938137ca..9a25fb543 100644
--- a/ntex-io/src/ioref.rs
+++ b/ntex-io/src/ioref.rs
@@ -228,6 +228,12 @@ impl IoRef {
         self.0.timeout.get()
     }
 
+    #[inline]
+    /// wakeup dispatcher and send keep-alive error
+    pub fn notify_timeout(&self) {
+        self.0.notify_timeout()
+    }
+
     #[inline]
     /// Start timer
     pub fn start_timer(&self, timeout: Seconds) -> timer::TimerHandle {
diff --git a/ntex/Cargo.toml b/ntex/Cargo.toml
index bb36997a8..b9b1503cb 100644
--- a/ntex/Cargo.toml
+++ b/ntex/Cargo.toml
@@ -67,7 +67,7 @@ ntex-bytes = "0.1.25"
 ntex-server = "1.0.5"
 ntex-h2 = "0.5.4"
 ntex-rt = "0.4.12"
-ntex-io = "1.0.1"
+ntex-io = "1.1.0"
 ntex-net = "1.0.1"
 ntex-tls = "1.1.0"
 
diff --git a/ntex/examples/client.rs b/ntex/examples/client.rs
index e5ecbf139..790d69124 100644
--- a/ntex/examples/client.rs
+++ b/ntex/examples/client.rs
@@ -2,24 +2,16 @@ use ntex::http::client::{error::SendRequestError, Client};
 
 #[ntex::main]
 async fn main() -> Result<(), SendRequestError> {
-    std::env::set_var("RUST_LOG", "ntex=trace");
+    std::env::set_var("RUST_LOG", "trace");
     env_logger::init();
 
     let client = Client::new();
 
     // Create request builder, configure request and send
-    let mut response = client
-        .get("https://www.rust-lang.org/")
-        .header("User-Agent", "ntex")
-        .send()
-        .await?;
-
-    // server http response
-    println!("Response: {:?}", response);
-
-    // read response body
-    let body = response.body().await.unwrap();
-    println!("Downloaded: {:?} bytes", body.len());
+    loop {
+        let _response = client.get("https://www.google.com").send().await.unwrap();
+        ntex::time::sleep(std::time::Duration::from_secs(10)).await;
+    }
 
     Ok(())
 }