From e1b6de06abbe6f11402a1549d623cac47e82336b Mon Sep 17 00:00:00 2001
From: Yiannis Marangos <yiannis@eiger.co>
Date: Mon, 7 Aug 2023 14:53:41 +0300
Subject: [PATCH] feat(timers): implement a drop-in replacement for
 `futures_timer::Delay`

---
 src/ic-cdk-timers/CHANGELOG.md |  4 ++
 src/ic-cdk-timers/src/delay.rs | 87 ++++++++++++++++++++++++++++++++++
 src/ic-cdk-timers/src/lib.rs   |  4 ++
 3 files changed, 95 insertions(+)
 create mode 100644 src/ic-cdk-timers/src/delay.rs

diff --git a/src/ic-cdk-timers/CHANGELOG.md b/src/ic-cdk-timers/CHANGELOG.md
index 341af576b..7823a1e37 100644
--- a/src/ic-cdk-timers/CHANGELOG.md
+++ b/src/ic-cdk-timers/CHANGELOG.md
@@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased]
 
+### Added
+
+- Implemented a drop-in replacement for `futures_timer::Delay`.
+
 ## [0.4.0] - 2023-07-13
 
 ### Changed
diff --git a/src/ic-cdk-timers/src/delay.rs b/src/ic-cdk-timers/src/delay.rs
new file mode 100644
index 000000000..aeb07ba55
--- /dev/null
+++ b/src/ic-cdk-timers/src/delay.rs
@@ -0,0 +1,87 @@
+use std::fmt;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use std::time::Duration;
+
+use futures_util::task::AtomicWaker;
+
+use crate::{clear_timer, set_timer, TimerId};
+
+/// A future representing the notification that an elapsed duration has
+/// occurred.
+///
+/// This is created through the `Delay::new` method indicating when the future should fire.
+/// Note that these futures are not intended for high resolution timers.
+pub struct Delay {
+    timer_id: Option<TimerId>,
+    waker: Arc<AtomicWaker>,
+    at: Duration,
+}
+
+impl Delay {
+    /// Creates a new future which will fire at `dur` time into the future.
+    pub fn new(dur: Duration) -> Delay {
+        let now = duration_since_epoch();
+
+        Delay {
+            timer_id: None,
+            waker: Arc::new(AtomicWaker::new()),
+            at: now + dur,
+        }
+    }
+
+    /// Resets this timeout to an new timeout which will fire at the time
+    /// specified by `at`.
+    pub fn reset(&mut self, dur: Duration) {
+        let now = duration_since_epoch();
+        self.at = now + dur;
+
+        if let Some(id) = self.timer_id.take() {
+            clear_timer(id);
+        }
+    }
+}
+
+impl Future for Delay {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let now = duration_since_epoch();
+
+        if now >= self.at {
+            Poll::Ready(())
+        } else {
+            // Register the latest waker
+            self.waker.register(cx.waker());
+
+            // Register to global timer
+            if self.timer_id.is_none() {
+                let waker = self.waker.clone();
+                let id = set_timer(self.at - now, move || waker.wake());
+                self.timer_id = Some(id);
+            }
+
+            Poll::Pending
+        }
+    }
+}
+
+impl Drop for Delay {
+    fn drop(&mut self) {
+        if let Some(id) = self.timer_id.take() {
+            clear_timer(id);
+        }
+    }
+}
+
+impl fmt::Debug for Delay {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        f.debug_struct("Delay").finish()
+    }
+}
+
+fn duration_since_epoch() -> Duration {
+    Duration::from_nanos(ic_cdk::api::time())
+}
diff --git a/src/ic-cdk-timers/src/lib.rs b/src/ic-cdk-timers/src/lib.rs
index 79a5eac71..599bb887c 100644
--- a/src/ic-cdk-timers/src/lib.rs
+++ b/src/ic-cdk-timers/src/lib.rs
@@ -31,6 +31,10 @@ use slotmap::{new_key_type, KeyData, SlotMap};
 
 use ic_cdk::api::call::RejectionCode;
 
+mod delay;
+
+pub use delay::Delay;
+
 // To ensure that tasks are removable seamlessly, there are two separate concepts here: tasks, for the actual function being called,
 // and timers, the scheduled execution of tasks. As this is an implementation detail, this does not affect the exported name TimerId,
 // which is more accurately a task ID. (The obvious solution to this, `pub use`, invokes a very silly compiler error.)