common_batcher/flush_policy/
timing.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::time::Duration;
16
17use tokio::time::Instant;
18
19use crate::flush_policy::{FlushPolicy, FlushTrigger};
20use crate::pending_batch::PendingBatch;
21
22/// Decides whether the first-submission deadline or row threshold requires a flush.
23///
24/// The row threshold is checked after appending a complete submission. It is not a
25/// hard batch-size limit and never splits a submission.
26#[derive(Debug, Clone, Copy)]
27pub struct TimingFlushPolicy {
28    flush_interval: Duration,
29    max_batch_rows: usize,
30}
31
32impl TimingFlushPolicy {
33    /// Creates a policy, rejecting zero limits, zero intervals, and unrepresentable deadlines.
34    pub fn try_new(flush_interval: Duration, max_batch_rows: usize) -> Option<Self> {
35        if max_batch_rows == 0 || !Self::validate(flush_interval) {
36            return None;
37        }
38        Some(Self {
39            flush_interval,
40            max_batch_rows,
41        })
42    }
43
44    /// Checks that the interval is nonzero and its deadline is representable.
45    pub fn validate(flush_interval: Duration) -> bool {
46        !flush_interval.is_zero() && Instant::now().checked_add(flush_interval).is_some()
47    }
48
49    /// Checks only the row trigger when a caller drives deadline events separately.
50    pub fn reached_row_threshold<T>(&self, batch: &PendingBatch<T>) -> bool {
51        !batch.is_empty() && batch.total_rows() >= self.max_batch_rows
52    }
53}
54
55impl FlushPolicy for TimingFlushPolicy {
56    /// Returns the deadline anchored to the first submission, or `None` for an empty batch.
57    ///
58    /// An unrepresentable deadline falls back to the first submission time so that
59    /// extreme timestamps cannot leave a nonempty batch waiting indefinitely.
60    fn deadline<T>(&self, batch: &PendingBatch<T>) -> Option<Instant> {
61        batch
62            .first_submitted_at()
63            .map(|first| first.checked_add(self.flush_interval).unwrap_or(first))
64    }
65
66    /// Returns whether a nonempty batch has reached either trigger.
67    fn should_flush<T>(
68        &self,
69        batch: &PendingBatch<T>,
70        now: Instant,
71        trigger: FlushTrigger,
72    ) -> bool {
73        !batch.is_empty()
74            && (self.reached_row_threshold(batch)
75                || (trigger == FlushTrigger::Deadline
76                    && self.deadline(batch).is_some_and(|deadline| now >= deadline)))
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_validate_matches_construction() {
86        for interval in [Duration::ZERO, Duration::from_secs(1), Duration::MAX] {
87            assert_eq!(
88                TimingFlushPolicy::validate(interval),
89                TimingFlushPolicy::try_new(interval, 1).is_some()
90            );
91        }
92    }
93
94    #[test]
95    fn test_first_submission_deadline() {
96        let policy = TimingFlushPolicy::try_new(Duration::from_millis(10), 100).unwrap();
97        let mut batch = PendingBatch::new();
98        let first = Instant::now();
99        assert_eq!(policy.deadline(&batch), None);
100        assert!(!policy.should_flush(&batch, first, FlushTrigger::Deadline));
101        batch.push(1, 1, first);
102        batch.push(2, 1, first + Duration::from_millis(9));
103        assert_eq!(
104            policy.deadline(&batch),
105            Some(first + Duration::from_millis(10))
106        );
107        assert!(!policy.should_flush(
108            &batch,
109            first + Duration::from_millis(9),
110            FlushTrigger::Deadline
111        ));
112        assert!(policy.should_flush(
113            &batch,
114            first + Duration::from_millis(10),
115            FlushTrigger::Deadline
116        ));
117        assert!(!policy.reached_row_threshold(&batch));
118        assert!(!policy.should_flush(
119            &batch,
120            first + Duration::from_millis(10),
121            FlushTrigger::Submission,
122        ));
123        assert_eq!(batch.take(), vec![1, 2]);
124        let next = first + Duration::from_millis(20);
125        batch.push(3, 0, next);
126        assert!(!policy.should_flush(&batch, next, FlushTrigger::Deadline));
127        assert!(policy.should_flush(
128            &batch,
129            next + Duration::from_millis(10),
130            FlushTrigger::Deadline
131        ));
132    }
133
134    #[test]
135    fn test_row_threshold_preserves_complete_submissions() {
136        let policy = TimingFlushPolicy::try_new(Duration::from_secs(1), 3).unwrap();
137        let now = Instant::now();
138        for rows in [2, 3, 4] {
139            let mut batch = PendingBatch::new();
140            batch.push(vec![0; rows], rows, now);
141            assert_eq!(
142                policy.should_flush(&batch, now, FlushTrigger::Submission),
143                rows >= 3
144            );
145            assert_eq!(policy.reached_row_threshold(&batch), rows >= 3);
146            assert_eq!(batch.take(), vec![vec![0; rows]]);
147        }
148    }
149
150    #[test]
151    fn test_invalid_interval() {
152        assert!(TimingFlushPolicy::try_new(Duration::from_secs(1), 0).is_none());
153        let max_batch_rows = 1;
154        assert!(TimingFlushPolicy::try_new(Duration::ZERO, max_batch_rows).is_none());
155        assert!(TimingFlushPolicy::try_new(Duration::MAX, max_batch_rows).is_none());
156    }
157}
common_batcher/flush_policy/timing.rs

common_batcher/flush_policy/
timing.rs
