common_batcher/
flush_timer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::pin::Pin;
16
17use tokio::time::{Instant, Sleep, sleep_until};
18
19/// Maintains one reusable timer for the current flush deadline.
20///
21/// Creating a timer does not require a Tokio runtime. Arming it requires a Tokio
22/// runtime with time enabled. Disabling retains the allocation for the next arm.
23#[derive(Default)]
24pub struct FlushTimer {
25    deadline: Option<Instant>,
26    sleep: Option<Pin<Box<Sleep>>>,
27}
28
29impl FlushTimer {
30    /// Creates a disabled timer without allocating a Tokio sleep.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Arms or disables the timer, resetting its sleep only when the deadline changes.
36    ///
37    /// # Panics
38    ///
39    /// Arming for the first time panics outside a Tokio runtime with time enabled.
40    pub fn set_deadline(&mut self, deadline: Option<Instant>) {
41        if self.deadline == deadline {
42            return;
43        }
44        self.deadline = deadline;
45        if let Some(deadline) = deadline {
46            match &mut self.sleep {
47                Some(sleep) => sleep.as_mut().reset(deadline),
48                None => self.sleep = Some(Box::pin(sleep_until(deadline))),
49            }
50        }
51    }
52
53    /// Waits until the armed deadline, or indefinitely while disabled.
54    ///
55    /// Cancelling this wait does not change the deadline. An elapsed timer stays
56    /// ready until the caller changes or disables its deadline.
57    pub async fn wait(&mut self) {
58        if self.deadline.is_some()
59            && let Some(sleep) = &mut self.sleep
60        {
61            sleep.as_mut().await;
62        } else {
63            std::future::pending::<()>().await;
64        }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use std::future::{Future, poll_fn};
71    use std::task::Poll;
72    use std::time::Duration;
73
74    use super::*;
75
76    async fn is_ready(timer: &mut FlushTimer) -> bool {
77        let wait = timer.wait();
78        tokio::pin!(wait);
79        poll_fn(|cx| Poll::Ready(wait.as_mut().poll(cx).is_ready())).await
80    }
81
82    #[test]
83    fn test_create_without_runtime() {
84        let mut timer = FlushTimer::new();
85        timer.set_deadline(None);
86        assert!(timer.sleep.is_none());
87    }
88
89    #[tokio::test(start_paused = true)]
90    async fn test_deadline_and_cancelled_wait() {
91        let mut timer = FlushTimer::new();
92        assert!(!is_ready(&mut timer).await);
93        let deadline = Instant::now() + Duration::from_millis(10);
94        timer.set_deadline(Some(deadline));
95        assert!(!is_ready(&mut timer).await);
96        tokio::time::advance(Duration::from_millis(9)).await;
97        // Each poll drops its wait future. Repeating the same deadline must not
98        // move the timer forward or replace its sleep allocation.
99        let sleep = timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep;
100        timer.set_deadline(Some(deadline));
101        assert_eq!(
102            sleep,
103            timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep
104        );
105        assert!(!is_ready(&mut timer).await);
106        tokio::time::advance(Duration::from_millis(1)).await;
107        assert!(is_ready(&mut timer).await);
108    }
109
110    #[tokio::test(start_paused = true)]
111    async fn test_disable_and_rearm_reuses_sleep() {
112        let mut timer = FlushTimer::new();
113        timer.set_deadline(Some(Instant::now()));
114        assert!(is_ready(&mut timer).await);
115        let sleep = timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep;
116        timer.set_deadline(None);
117        assert!(!is_ready(&mut timer).await);
118        assert!(!is_ready(&mut timer).await);
119        timer.set_deadline(Some(Instant::now() + Duration::from_millis(5)));
120        assert_eq!(
121            sleep,
122            timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep
123        );
124        assert!(!is_ready(&mut timer).await);
125        tokio::time::advance(Duration::from_millis(5)).await;
126        assert!(is_ready(&mut timer).await);
127    }
128
129    #[tokio::test(start_paused = true)]
130    async fn test_change_armed_deadline() {
131        let mut timer = FlushTimer::new();
132        let now = Instant::now();
133        timer.set_deadline(Some(now + Duration::from_millis(5)));
134        assert!(!is_ready(&mut timer).await);
135        timer.set_deadline(Some(now + Duration::from_millis(10)));
136        tokio::time::advance(Duration::from_millis(5)).await;
137        assert!(!is_ready(&mut timer).await);
138        timer.set_deadline(Some(Instant::now()));
139        assert!(is_ready(&mut timer).await);
140    }
141}
common_batcher/flush_timer.rs

common_batcher/
flush_timer.rs
