From 9cb52bd37b4511a5d629b099a191f2f0f771576f Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Tue, 4 Jan 2022 16:04:04 +0600
Subject: [PATCH] Add time::timeout_checked, if duration is zero then timeout
 is disabled

---
 ntex-util/CHANGES.md      |  4 +++
 ntex-util/Cargo.toml      |  4 +--
 ntex-util/src/time/mod.rs | 73 +++++++++++++++++++++++++++++++++++++++
 ntex/Cargo.toml           |  2 +-
 4 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/ntex-util/CHANGES.md b/ntex-util/CHANGES.md
index 8e3e3531..c0774247 100644
--- a/ntex-util/CHANGES.md
+++ b/ntex-util/CHANGES.md
@@ -1,5 +1,9 @@
 # Changes
 
+## [0.1.7] - 2022-01-04
+
+* Add time::timeout_checked, if duration is zero then timeout is disabled
+
 ## [0.1.6] - 2022-01-03
 
 * Use ntex-rt::spawn
diff --git a/ntex-util/Cargo.toml b/ntex-util/Cargo.toml
index d18aca5a..434e5ed1 100644
--- a/ntex-util/Cargo.toml
+++ b/ntex-util/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ntex-util"
-version = "0.1.6"
+version = "0.1.7"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Utilities for ntex framework"
 keywords = ["network", "framework", "async", "futures"]
@@ -28,7 +28,7 @@ futures-sink = { version = "0.3", default-features = false, features = ["alloc"]
 pin-project-lite = "0.2.6"
 
 [dev-dependencies]
-ntex = "0.5.4"
+ntex = { version = "0.5.6", features = ["tokio"] }
 ntex-bytes = "0.1.8"
 ntex-macros = "0.1.3"
 derive_more = "0.99"
diff --git a/ntex-util/src/time/mod.rs b/ntex-util/src/time/mod.rs
index 701aa30a..ee292932 100644
--- a/ntex-util/src/time/mod.rs
+++ b/ntex-util/src/time/mod.rs
@@ -40,6 +40,20 @@ where
     Timeout::new_with_delay(future, Sleep::new(dur.into()))
 }
 
+/// Require a `Future` to complete before the specified duration has elapsed.
+///
+/// If the future completes before the duration has elapsed, then the completed
+/// value is returned. Otherwise, an error is returned and the future is
+/// canceled. If duration value is zero then timeout is disabled.
+#[inline]
+pub fn timeout_checked<T, U>(dur: U, future: T) -> TimeoutChecked<T>
+where
+    T: Future,
+    U: Into<Millis>,
+{
+    TimeoutChecked::new_with_delay(future, dur.into())
+}
+
 /// Future returned by [`sleep`](sleep).
 ///
 /// # Examples
@@ -141,6 +155,53 @@ where
     }
 }
 
+pin_project_lite::pin_project! {
+    /// Future returned by [`timeout_checked`](timeout_checked).
+    #[must_use = "futures do nothing unless you `.await` or poll them"]
+    pub struct TimeoutChecked<T> {
+        #[pin]
+        state: TimeoutCheckedState<T>,
+    }
+}
+
+pin_project_lite::pin_project! {
+    #[project = TimeoutCheckedStateProject]
+    enum TimeoutCheckedState<T> {
+        Timeout{ #[pin] fut: Timeout<T> },
+        NoTimeout{ #[pin] fut: T },
+    }
+}
+
+impl<T> TimeoutChecked<T> {
+    pub(crate) fn new_with_delay(value: T, delay: Millis) -> TimeoutChecked<T> {
+        if delay.is_zero() {
+            TimeoutChecked {
+                state: TimeoutCheckedState::NoTimeout { fut: value },
+            }
+        } else {
+            TimeoutChecked {
+                state: TimeoutCheckedState::Timeout {
+                    fut: Timeout::new_with_delay(value, sleep(delay)),
+                },
+            }
+        }
+    }
+}
+
+impl<T> Future for TimeoutChecked<T>
+where
+    T: Future,
+{
+    type Output = Result<T::Output, ()>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
+        match self.project().state.as_mut().project() {
+            TimeoutCheckedStateProject::Timeout { fut } => fut.poll(cx),
+            TimeoutCheckedStateProject::NoTimeout { fut } => fut.poll(cx).map(Result::Ok),
+        }
+    }
+}
+
 /// Interval returned by [`interval`]
 ///
 /// This type allows you to wait on a sequence of instants with a certain
@@ -288,4 +349,16 @@ mod tests {
             );
         }
     }
+
+    #[ntex_macros::rt_test2]
+    async fn test_timeout_checked() {
+        let result = timeout_checked(Millis(200), sleep(Millis(100))).await;
+        assert!(result.is_ok());
+
+        let result = timeout_checked(Millis(5), sleep(Millis(100))).await;
+        assert!(result.is_err());
+
+        let result = timeout_checked(Millis(0), sleep(Millis(100))).await;
+        assert!(result.is_ok());
+    }
 }
diff --git a/ntex/Cargo.toml b/ntex/Cargo.toml
index 7bf3341a..ee3e8669 100644
--- a/ntex/Cargo.toml
+++ b/ntex/Cargo.toml
@@ -49,7 +49,7 @@ ntex-codec = "0.6.0"
 ntex-router = "0.5.1"
 ntex-service = "0.3.1"
 ntex-macros = "0.1.3"
-ntex-util = "0.1.6"
+ntex-util = "0.1.7"
 ntex-bytes = "0.1.8"
 ntex-tls = "0.1.0"
 ntex-rt = "0.4.1"