Skip to main content

rocketmq_client_rust/latency/
mq_fault_strategy.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::Ordering;
17
18use cheetah_string::CheetahString;
19use rocketmq_common::common::message::message_queue::MessageQueue;
20use rocketmq_rust::ArcMut;
21
22use crate::base::client_config::ClientConfig;
23use crate::latency::latency_fault_tolerance::LatencyFaultTolerance;
24use crate::latency::latency_fault_tolerance_impl::LatencyFaultToleranceImpl;
25use crate::producer::producer_impl::default_mq_producer_impl::DefaultResolver;
26use crate::producer::producer_impl::default_mq_producer_impl::DefaultServiceDetector;
27use crate::producer::producer_impl::queue_filter::QueueFilter;
28use crate::producer::producer_impl::topic_publish_info::TopicPublishInfo;
29
30pub struct MQFaultStrategy {
31    latency_fault_tolerance: ArcMut<LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>>,
32    send_latency_fault_enable: AtomicBool,
33    start_detector_enable: AtomicBool,
34    latency_max: Vec<u64>,
35    not_available_duration: Vec<u64>,
36    reachable_filter: Box<dyn QueueFilter>,
37    available_filter: Box<dyn QueueFilter>,
38}
39
40impl MQFaultStrategy {
41    pub const DEFAULT_LATENCY_MAX: [u64; 7] = [50, 100, 550, 1800, 3000, 5000, 15000];
42    pub const DEFAULT_NOT_AVAILABLE_DURATION: [u64; 7] = [0, 0, 2000, 5000, 6000, 10000, 30000];
43    /// Latency penalty applied when a broker is isolated (10 seconds).
44    const ISOLATION_LATENCY_MS: u64 = 10_000;
45
46    pub fn new(client_config: &ClientConfig) -> Self {
47        let mut tolerance_impl = LatencyFaultToleranceImpl::new();
48        tolerance_impl.set_detect_interval(client_config.detect_interval);
49        tolerance_impl.set_detect_timeout(client_config.detect_timeout);
50        tolerance_impl.set_start_detector_enable(client_config.start_detector_enable);
51        let latency_fault_tolerance = ArcMut::new(tolerance_impl);
52        Self {
53            latency_fault_tolerance: ArcMut::clone(&latency_fault_tolerance),
54            send_latency_fault_enable: AtomicBool::new(client_config.send_latency_enable),
55            start_detector_enable: AtomicBool::new(client_config.start_detector_enable),
56            latency_max: Self::DEFAULT_LATENCY_MAX.to_vec(),
57            not_available_duration: Self::DEFAULT_NOT_AVAILABLE_DURATION.to_vec(),
58            reachable_filter: Box::new(ReachableFilter {
59                latency_fault_tolerance: ArcMut::clone(&latency_fault_tolerance),
60            }),
61            available_filter: Box::new(AvailableFilter {
62                latency_fault_tolerance,
63            }),
64        }
65    }
66
67    pub fn start_detector(&mut self) {
68        LatencyFaultTolerance::start_detector(self.latency_fault_tolerance.clone());
69    }
70
71    pub(crate) fn set_resolve(&mut self, resolver: DefaultResolver) {
72        self.latency_fault_tolerance.set_resolver(resolver);
73    }
74
75    pub(crate) fn set_service_detector(&mut self, service_detector: DefaultServiceDetector) {
76        self.latency_fault_tolerance.set_service_detector(service_detector);
77    }
78
79    pub fn shutdown(&mut self) {
80        self.latency_fault_tolerance.shutdown();
81    }
82
83    pub fn is_start_detector_enable(&self) -> bool {
84        self.start_detector_enable.load(Ordering::Relaxed)
85    }
86
87    pub fn set_start_detector_enable(&mut self, start_detector_enable: bool) {
88        self.start_detector_enable
89            .store(start_detector_enable, Ordering::Relaxed);
90        self.latency_fault_tolerance
91            .set_start_detector_enable(start_detector_enable);
92    }
93
94    pub fn is_send_latency_fault_enable(&self) -> bool {
95        self.send_latency_fault_enable.load(Ordering::Relaxed)
96    }
97
98    pub fn select_one_message_queue(
99        &self,
100        tp_info: &TopicPublishInfo,
101        last_broker_name: Option<&CheetahString>,
102        reset_index: bool,
103    ) -> Option<MessageQueue> {
104        let broker_filter = BrokerFilter::with_last_broker_name(last_broker_name.cloned());
105
106        if self.send_latency_fault_enable.load(Ordering::Relaxed) {
107            if reset_index {
108                tp_info.reset_index();
109            }
110
111            let filter = &[self.available_filter.as_ref(), &broker_filter as &dyn QueueFilter];
112            if let Some(mq) = tp_info.select_one_message_queue_filters(filter) {
113                return Some(mq);
114            }
115
116            let filter = &[self.reachable_filter.as_ref(), &broker_filter as &dyn QueueFilter];
117            if let Some(mq) = tp_info.select_one_message_queue_filters(filter) {
118                return Some(mq);
119            }
120
121            return tp_info.select_one_message_queue_filters(&[]);
122        }
123
124        if let Some(mq) = tp_info.select_one_message_queue_filters(&[&broker_filter]) {
125            return Some(mq);
126        }
127        tp_info.select_one_message_queue_filters(&[])
128    }
129
130    pub fn get_latency_max(&self) -> &[u64] {
131        &self.latency_max
132    }
133
134    pub fn set_latency_max(&mut self, latency_max: impl Into<Vec<u64>>) {
135        self.latency_max = latency_max.into();
136    }
137
138    pub fn get_not_available_duration(&self) -> &[u64] {
139        &self.not_available_duration
140    }
141
142    pub fn set_not_available_duration(&mut self, not_available_duration: impl Into<Vec<u64>>) {
143        self.not_available_duration = not_available_duration.into();
144    }
145
146    pub async fn update_fault_item(
147        &self,
148        broker_name: CheetahString,
149        current_latency: u64,
150        isolation: bool,
151        reachable: bool,
152    ) {
153        if self.send_latency_fault_enable.load(Ordering::Relaxed) {
154            let effective_latency = if isolation {
155                Self::ISOLATION_LATENCY_MS
156            } else {
157                current_latency
158            };
159            let duration = self.compute_not_available_duration(effective_latency);
160            self.latency_fault_tolerance
161                .update_fault_item(broker_name, current_latency, duration, reachable)
162                .await;
163        }
164    }
165
166    fn compute_not_available_duration(&self, current_latency: u64) -> u64 {
167        let len = self.latency_max.len().min(self.not_available_duration.len());
168        for i in (0..len).rev() {
169            if current_latency >= self.latency_max[i] {
170                return self.not_available_duration[i];
171            }
172        }
173        0
174    }
175
176    #[inline]
177    pub fn set_send_latency_fault_enable(&mut self, send_latency_fault_enable: bool) {
178        self.send_latency_fault_enable
179            .store(send_latency_fault_enable, Ordering::Relaxed);
180    }
181}
182
183#[derive(Clone, Debug, Default, Eq, PartialEq)]
184pub struct BrokerFilter {
185    last_broker_name: Option<CheetahString>,
186}
187
188impl BrokerFilter {
189    pub fn new() -> Self {
190        Self::default()
191    }
192
193    pub fn with_last_broker_name(last_broker_name: Option<CheetahString>) -> Self {
194        Self { last_broker_name }
195    }
196
197    pub fn last_broker_name(&self) -> Option<&CheetahString> {
198        self.last_broker_name.as_ref()
199    }
200
201    pub fn get_last_broker_name(&self) -> Option<&CheetahString> {
202        self.last_broker_name()
203    }
204
205    pub fn set_last_broker_name(&mut self, last_broker_name: impl Into<CheetahString>) {
206        self.last_broker_name = Some(last_broker_name.into());
207    }
208
209    pub fn clear_last_broker_name(&mut self) {
210        self.last_broker_name = None;
211    }
212}
213
214impl QueueFilter for BrokerFilter {
215    fn filter(&self, message_queue: &MessageQueue) -> bool {
216        if let Some(last_broker_name) = &self.last_broker_name {
217            message_queue.broker_name() != last_broker_name
218        } else {
219            true
220        }
221    }
222}
223
224struct ReachableFilter {
225    latency_fault_tolerance: ArcMut<LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>>,
226}
227
228impl QueueFilter for ReachableFilter {
229    fn filter(&self, message_queue: &MessageQueue) -> bool {
230        self.latency_fault_tolerance.is_reachable(message_queue.broker_name())
231    }
232}
233
234struct AvailableFilter {
235    latency_fault_tolerance: ArcMut<LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>>,
236}
237
238impl QueueFilter for AvailableFilter {
239    fn filter(&self, message_queue: &MessageQueue) -> bool {
240        self.latency_fault_tolerance.is_available(message_queue.broker_name())
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247
248    #[test]
249    fn broker_filter_matches_java_last_broker_name_rule() {
250        let mq = MessageQueue::from_parts("topic", "broker-a", 0);
251        let mut filter = BrokerFilter::new();
252
253        assert!(filter.filter(&mq));
254        assert!(filter.get_last_broker_name().is_none());
255
256        filter.set_last_broker_name("broker-a");
257        assert_eq!(
258            filter.get_last_broker_name(),
259            Some(&CheetahString::from_static_str("broker-a"))
260        );
261        assert!(!filter.filter(&mq));
262
263        filter.set_last_broker_name("broker-b");
264        assert!(filter.filter(&mq));
265
266        filter.clear_last_broker_name();
267        assert!(filter.filter(&mq));
268    }
269
270    #[test]
271    fn mq_fault_strategy_copies_detector_config_like_java_constructor() {
272        let mut client_config = ClientConfig::default();
273        client_config.set_detect_interval(12_345);
274        client_config.set_detect_timeout(321);
275        client_config.set_start_detector_enable(true);
276
277        let strategy = MQFaultStrategy::new(&client_config);
278        let (detect_interval, detect_timeout) = strategy.latency_fault_tolerance.detector_config_for_test();
279
280        assert_eq!(detect_interval, 12_345);
281        assert_eq!(detect_timeout, 321);
282        assert!(strategy.is_start_detector_enable());
283        assert!(strategy.latency_fault_tolerance.is_start_detector_enable());
284    }
285}