rocketmq_client_rust/consumer/store/
controllable_offset.rs

1// Copyright 2023 The RocketMQ Rust Authors
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::sync::atomic::AtomicBool;
16use std::sync::atomic::AtomicI64;
17use std::sync::atomic::Ordering;
18use std::sync::Arc;
19
20pub struct ControllableOffset {
21    value: Arc<AtomicI64>,
22    allow_to_update: Arc<AtomicBool>,
23}
24
25impl ControllableOffset {
26    pub fn new(value: i64) -> Self {
27        Self {
28            value: Arc::new(AtomicI64::new(value)),
29            allow_to_update: Arc::new(AtomicBool::new(true)),
30        }
31    }
32
33    pub fn new_frozen(value: i64) -> Self {
34        Self {
35            value: Arc::new(AtomicI64::new(value)),
36            allow_to_update: Arc::new(AtomicBool::new(false)),
37        }
38    }
39
40    pub fn new_atomic(value: AtomicI64) -> Self {
41        Self {
42            value: Arc::new(value),
43            allow_to_update: Arc::new(AtomicBool::new(true)),
44        }
45    }
46
47    pub fn update(&self, target: i64, increase_only: bool) {
48        if self.allow_to_update.load(Ordering::SeqCst) {
49            self.value
50                .try_update(Ordering::SeqCst, Ordering::SeqCst, |val| {
51                    if self.allow_to_update.load(Ordering::SeqCst) {
52                        if increase_only {
53                            Some(std::cmp::max(target, val))
54                        } else {
55                            Some(target)
56                        }
57                    } else {
58                        Some(val)
59                    }
60                })
61                .ok();
62        }
63    }
64
65    pub fn update_unconditionally(&self, target: i64) {
66        self.update(target, false);
67    }
68
69    pub fn update_and_freeze(&self, target: i64) {
70        self.value
71            .try_update(Ordering::SeqCst, Ordering::SeqCst, |_val| {
72                self.allow_to_update.store(false, Ordering::SeqCst);
73                Some(target)
74            })
75            .ok();
76    }
77
78    pub fn get_offset(&self) -> i64 {
79        self.value.load(Ordering::SeqCst)
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn unfrozen_offset_accepts_lower_value_only_without_increase_only() {
89        let offset = ControllableOffset::new(10);
90
91        offset.update(5, true);
92        assert_eq!(offset.get_offset(), 10);
93
94        offset.update(5, false);
95        assert_eq!(offset.get_offset(), 5);
96    }
97
98    #[test]
99    fn update_and_freeze_locks_the_value_against_later_updates() {
100        let offset = ControllableOffset::new(10);
101
102        offset.update_and_freeze(20);
103        assert_eq!(offset.get_offset(), 20);
104
105        offset.update(30, true);
106        assert_eq!(offset.get_offset(), 20);
107
108        offset.update(5, false);
109        assert_eq!(offset.get_offset(), 20);
110
111        offset.update_unconditionally(40);
112        assert_eq!(offset.get_offset(), 20);
113    }
114
115    #[test]
116    fn new_frozen_rejects_both_update_modes_immediately() {
117        let offset = ControllableOffset::new_frozen(10);
118
119        offset.update(20, true);
120        assert_eq!(offset.get_offset(), 10);
121
122        offset.update(5, false);
123        assert_eq!(offset.get_offset(), 10);
124    }
125
126    #[test]
127    fn second_update_and_freeze_changes_the_frozen_value() {
128        let offset = ControllableOffset::new(10);
129
130        offset.update_and_freeze(20);
131        offset.update_and_freeze(7);
132        assert_eq!(offset.get_offset(), 7);
133
134        offset.update_unconditionally(100);
135        assert_eq!(offset.get_offset(), 7);
136    }
137}
rocketmq_client_rust/consumer/store/controllable_offset.rs

rocketmq_client_rust/consumer/store/
controllable_offset.rs
