implement with_timeout()/with_deadline() method style call on Future

2 files changed, 31 insertions, 1 deletions

diff --git a/embassy-time/src/lib.rs b/embassy-time/src/lib.rs
index 24ee51be7..8d0648ce5 100644
--- a/embassy-time/src/lib.rs
+++ b/embassy-time/src/lib.rs
@@ -32,7 +32,7 @@ pub use delay::{block_for, Delay};
 pub use duration::Duration;
 pub use embassy_time_driver::TICK_HZ;
 pub use instant::Instant;
-pub use timer::{with_deadline, with_timeout, Ticker, TimeoutError, Timer};
+pub use timer::{with_deadline, with_timeout, Ticker, TimeoutError, Timer, WithTimeout};
 
 const fn gcd(a: u64, b: u64) -> u64 {
     if b == 0 {
diff --git a/embassy-time/src/timer.rs b/embassy-time/src/timer.rs
index bc39d8bc7..4d7194b20 100644
--- a/embassy-time/src/timer.rs
+++ b/embassy-time/src/timer.rs
@@ -37,6 +37,36 @@ pub async fn with_deadline<F: Future>(at: Instant, fut: F) -> Result<F::Output,
     }
 }
 
+/// Provides functions to run a given future with a timeout or a deadline.
+pub trait WithTimeout {
+    /// Output type of the future.
+    type Output;
+
+    /// Runs a given future with a timeout.
+    ///
+    /// If the future completes before the timeout, its output is returned. Otherwise, on timeout,
+    /// work on the future is stopped (`poll` is no longer called), the future is dropped and `Err(TimeoutError)` is returned.
+    async fn with_timeout(self, timeout: Duration) -> Result<Self::Output, TimeoutError>;
+
+    /// Runs a given future with a deadline time.
+    ///
+    /// If the future completes before the deadline, its output is returned. Otherwise, on timeout,
+    /// work on the future is stopped (`poll` is no longer called), the future is dropped and `Err(TimeoutError)` is returned.
+    async fn with_deadline(self, at: Instant) -> Result<Self::Output, TimeoutError>;
+}
+
+impl<F: Future> WithTimeout for F {
+    type Output = F::Output;
+
+    async fn with_timeout(self, timeout: Duration) -> Result<Self::Output, TimeoutError> {
+        with_timeout(timeout, self).await
+    }
+
+    async fn with_deadline(self, at: Instant) -> Result<Self::Output, TimeoutError> {
+        with_deadline(at, self).await
+    }
+}
+
 /// A future that completes at a specified [Instant](struct.Instant.html).
 #[must_use = "futures do nothing unless you `.await` or poll them"]
 pub struct Timer {