common_batcher/
pending_batch.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 tokio::time::Instant;
16
17/// Accumulates complete submissions in arrival order, without interpreting their data.
18#[derive(Debug)]
19pub struct PendingBatch<T> {
20    items: Vec<T>,
21    total_rows: usize,
22    first_submitted_at: Option<Instant>,
23}
24
25impl<T> Default for PendingBatch<T> {
26    fn default() -> Self {
27        Self {
28            items: Vec::new(),
29            total_rows: 0,
30            first_submitted_at: None,
31        }
32    }
33}
34
35impl<T> PendingBatch<T> {
36    /// Creates an empty batch.
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Appends a complete submission without moving the first submission's timestamp.
42    pub fn push(&mut self, item: T, total_rows: usize, now: Instant) {
43        self.first_submitted_at.get_or_insert(now);
44        self.items.push(item);
45        // Saturation preserves threshold decisions even for caller-supplied oversized weights.
46        self.total_rows = self.total_rows.saturating_add(total_rows);
47    }
48
49    /// Returns the accumulated row count, saturated at `usize::MAX`.
50    pub fn total_rows(&self) -> usize {
51        self.total_rows
52    }
53
54    /// Returns the number of complete submissions.
55    pub fn len(&self) -> usize {
56        self.items.len()
57    }
58
59    /// Returns whether there are no submissions, including zero-row submissions.
60    pub fn is_empty(&self) -> bool {
61        self.items.is_empty()
62    }
63
64    /// Returns the time the first submission was appended to the current batch.
65    pub fn first_submitted_at(&self) -> Option<Instant> {
66        self.first_submitted_at
67    }
68
69    /// Takes all submissions in arrival order and resets the batch's accounting.
70    pub fn take(&mut self) -> Vec<T> {
71        self.total_rows = 0;
72        self.first_submitted_at = None;
73        std::mem::take(&mut self.items)
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use std::time::Duration;
80
81    use super::*;
82
83    #[test]
84    fn test_complete_submissions_and_reset() {
85        let mut batch = PendingBatch::new();
86        let first = Instant::now();
87        assert!(batch.is_empty());
88        assert_eq!(batch.first_submitted_at(), None);
89        batch.push(vec![1, 2], 2, first);
90        batch.push(vec![3, 4, 5], 3, first + Duration::from_secs(1));
91        assert_eq!(batch.len(), 2);
92        assert_eq!(batch.total_rows(), 5);
93        assert_eq!(batch.first_submitted_at(), Some(first));
94        assert_eq!(batch.take(), vec![vec![1, 2], vec![3, 4, 5]]);
95        assert!(batch.is_empty());
96        assert_eq!(batch.total_rows(), 0);
97        assert_eq!(batch.first_submitted_at(), None);
98        let next = first + Duration::from_secs(2);
99        batch.push(vec![6], 1, next);
100        assert_eq!(batch.first_submitted_at(), Some(next));
101    }
102
103    #[test]
104    fn test_zero_rows_and_saturating_count() {
105        let mut batch = PendingBatch::new();
106        let now = Instant::now();
107        batch.push(1, 0, now);
108        assert!(!batch.is_empty());
109        assert_eq!(batch.total_rows(), 0);
110        batch.push(2, usize::MAX, now);
111        batch.push(3, 1, now);
112        assert_eq!(batch.total_rows(), usize::MAX);
113        assert_eq!(batch.take(), vec![1, 2, 3]);
114    }
115}
common_batcher/pending_batch.rs

common_batcher/
pending_batch.rs
