From 108a74f7eb9ac9358165703afa0bdb70bfb605e5 Mon Sep 17 00:00:00 2001
From: Cameron Bytheway <bytheway.cameron@gmail.com>
Date: Fri, 16 Feb 2024 11:38:03 -0700
Subject: [PATCH] feat(s2n-quic-core): add Cached clock implementation

---
 quic/s2n-quic-core/src/time/clock.rs | 36 ++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/quic/s2n-quic-core/src/time/clock.rs b/quic/s2n-quic-core/src/time/clock.rs
index cf31a46d9b..2ecf3fc03b 100644
--- a/quic/s2n-quic-core/src/time/clock.rs
+++ b/quic/s2n-quic-core/src/time/clock.rs
@@ -48,3 +48,39 @@ impl Clock for NoopClock {
         unsafe { Timestamp::from_duration(Duration::from_micros(1)) }
     }
 }
+
+impl Clock for Timestamp {
+    #[inline]
+    fn get_time(&self) -> Timestamp {
+        *self
+    }
+}
+
+/// A clock that caches the time query for the inner clock
+pub struct Cached<'a, C: Clock> {
+    clock: &'a C,
+    cached_value: core::cell::Cell<Option<Timestamp>>,
+}
+
+impl<'a, C: Clock> Cached<'a, C> {
+    #[inline]
+    pub fn new(clock: &'a C) -> Self {
+        Self {
+            clock,
+            cached_value: Default::default(),
+        }
+    }
+}
+
+impl<'a, C: Clock> Clock for Cached<'a, C> {
+    #[inline]
+    fn get_time(&self) -> Timestamp {
+        if let Some(time) = self.cached_value.get() {
+            return time;
+        }
+
+        let now = self.clock.get_time();
+        self.cached_value.set(Some(now));
+        now
+    }
+}