log_store/kafka/consumer/
range.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::cmp::min;
16use std::ops::Range;
17
18use store_api::logstore::EntryId;
19
20#[derive(Debug, PartialEq, Eq)]
21pub(crate) struct NextBatchHint {
22    pub(crate) bytes: usize,
23    pub(crate) len: usize,
24}
25
26/// Represents a range [next_entry_id, end_entry_id) of WAL entries for a region.
27#[derive(Debug)]
28pub struct RegionWalRange {
29    current_entry_id: EntryId,
30    end_entry_id: EntryId,
31    max_batch_size: usize,
32}
33
34impl RegionWalRange {
35    pub fn new(range: Range<EntryId>, max_batch_size: usize) -> Self {
36        Self {
37            current_entry_id: range.start,
38            end_entry_id: range.end,
39            max_batch_size,
40        }
41    }
42
43    fn next_batch_size(&self) -> Option<u64> {
44        if self.current_entry_id < self.end_entry_id {
45            Some(self.end_entry_id.saturating_sub(self.current_entry_id))
46        } else {
47            None
48        }
49    }
50
51    pub(crate) fn next_batch_hint(&self, avg_size: usize) -> Option<NextBatchHint> {
52        if let Some(size) = self.next_batch_size() {
53            let bytes = min(size as usize * avg_size, self.max_batch_size);
54            let len = bytes / avg_size;
55
56            return Some(NextBatchHint { bytes, len });
57        }
58
59        None
60    }
61
62    pub(crate) fn peek(&self) -> Option<EntryId> {
63        if self.current_entry_id < self.end_entry_id {
64            Some(self.current_entry_id)
65        } else {
66            None
67        }
68    }
69
70    pub(crate) fn next(&mut self) -> Option<EntryId> {
71        if self.current_entry_id < self.end_entry_id {
72            let next = self.current_entry_id;
73            self.current_entry_id += 1;
74            Some(next)
75        } else {
76            None
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn test_region_wal_range() {
87        let range = RegionWalRange::new(0..1024, 1024);
88        assert_eq!(
89            range.next_batch_hint(10),
90            Some(NextBatchHint {
91                bytes: 1024,
92                len: 102
93            })
94        );
95
96        let mut range = RegionWalRange::new(0..1, 1024);
97
98        assert_eq!(range.next_batch_size(), Some(1));
99        assert_eq!(range.peek(), Some(0));
100
101        // Advance 1 step
102        assert_eq!(range.next(), Some(0));
103        assert_eq!(range.next_batch_size(), None);
104
105        // Advance 1 step
106        assert_eq!(range.next(), None);
107        assert_eq!(range.next_batch_size(), None);
108        // No effect
109        assert_eq!(range.next(), None);
110        assert_eq!(range.next_batch_size(), None);
111
112        let mut range = RegionWalRange::new(0..0, 1024);
113        assert_eq!(range.next_batch_size(), None);
114        // No effect
115        assert_eq!(range.next(), None);
116        assert_eq!(range.next_batch_size(), None);
117    }
118}
log_store/kafka/consumer/range.rs

log_store/kafka/consumer/
range.rs
