s2n_quic_dc/
clock.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use core::{fmt, pin::Pin, task::Poll, time::Duration};
5use s2n_quic_core::{
6    ensure, time,
7    time::{timer, timer::Provider},
8};
9use tracing::trace;
10
11#[macro_use]
12mod macros;
13
14#[cfg(any(test, feature = "testing"))]
15pub mod bach;
16#[cfg(feature = "tokio")]
17pub mod tokio;
18pub use time::clock::Cached;
19
20pub use time::Timestamp;
21
22use crate::either::Either;
23pub type SleepHandle = Pin<Box<dyn Sleep>>;
24
25pub trait Clock: 'static + Send + Sync + fmt::Debug + time::Clock {
26    fn sleep(&self, amount: Duration) -> (SleepHandle, Timestamp);
27
28    fn timer(&self) -> Timer
29    where
30        Self: Sized,
31    {
32        Timer::new(self)
33    }
34}
35
36impl<A, B> Clock for Either<A, B>
37where
38    A: Clock,
39    B: Clock,
40{
41    fn sleep(&self, amount: Duration) -> (SleepHandle, Timestamp) {
42        match self {
43            Either::A(a) => a.sleep(amount),
44            Either::B(b) => b.sleep(amount),
45        }
46    }
47}
48
49impl<A, B> time::Clock for Either<A, B>
50where
51    A: time::Clock,
52    B: time::Clock,
53{
54    fn get_time(&self) -> Timestamp {
55        match self {
56            Either::A(a) => a.get_time(),
57            Either::B(b) => b.get_time(),
58        }
59    }
60}
61
62pub trait Sleep: Clock + core::future::Future<Output = ()> {
63    fn update(self: Pin<&mut Self>, target: Timestamp);
64}
65
66pub struct Timer {
67    /// The `Instant` at which the timer should expire
68    target: timer::Timer,
69    /// The handle to the timer entry in the tokio runtime
70    sleep: Pin<Box<dyn Sleep>>,
71}
72
73impl fmt::Debug for Timer {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.debug_struct("Timer")
76            .field("target", &self.target)
77            .finish()
78    }
79}
80
81impl Timer {
82    #[inline]
83    pub fn new(clock: &dyn Clock) -> Self {
84        /// We can't create a timer without first arming it to something, so just set it to 1s in
85        /// the future.
86        const INITIAL_TIMEOUT: Duration = Duration::from_secs(1);
87
88        Self::new_with_timeout(clock, INITIAL_TIMEOUT)
89    }
90
91    #[inline]
92    pub fn new_with_timeout(clock: &dyn Clock, timeout: Duration) -> Self {
93        let (sleep, target) = clock.sleep(timeout);
94        let mut timer = timer::Timer::default();
95        timer.set(target);
96        Self {
97            target: timer,
98            sleep,
99        }
100    }
101
102    #[inline]
103    pub fn cancel(&mut self) {
104        trace!(cancel = ?self.target);
105        self.target.cancel();
106    }
107
108    pub async fn sleep(&mut self, target: Timestamp) {
109        use time::clock::Timer;
110        self.update(target);
111        core::future::poll_fn(|cx| self.poll_ready(cx)).await
112    }
113}
114
115impl time::clock::Timer for Timer {
116    #[inline]
117    fn poll_ready(&mut self, cx: &mut core::task::Context) -> Poll<()> {
118        ensure!(self.target.is_armed(), Poll::Ready(()));
119
120        let res = self.sleep.as_mut().poll(cx);
121
122        if res.is_ready() {
123            // clear the target after it fires, otherwise we'll endlessly wake up the task
124            self.target.cancel();
125        }
126
127        res
128    }
129
130    #[inline]
131    fn update(&mut self, target: Timestamp) {
132        // no need to update if it hasn't changed
133        ensure!(self.target.next_expiration() != Some(target));
134
135        self.sleep.as_mut().update(target);
136        self.target.set(target);
137    }
138}
139
140impl timer::Provider for Timer {
141    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
142        self.target.timers(query)
143    }
144}
s2n_quic_dc/clock.rs

s2n_quic_dc/
clock.rs
