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

rocketmq_client_rust/latency/
mq_fault_strategy.rs
