rocketmq-client-rust 1.0.0

// Copyright 2023 The RocketMQ Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use cheetah_string::CheetahString;
use rocketmq_model::common::message::message_queue::MessageQueue;
use rocketmq_runtime::ChildServiceContext;

use crate::base::client_config::ClientConfig;
use crate::latency::latency_fault_tolerance::LatencyFaultTolerance;
#[cfg(any(test, feature = "test-support"))]
use crate::latency::latency_fault_tolerance_impl::FaultClock;
use crate::latency::latency_fault_tolerance_impl::LatencyFaultToleranceImpl;
use crate::producer::producer_impl::default_mq_producer_impl::DefaultResolver;
use crate::producer::producer_impl::default_mq_producer_impl::DefaultServiceDetector;
use crate::producer::producer_impl::queue_filter::QueueFilter;
use crate::producer::producer_impl::topic_publish_info::TopicPublishInfo;

#[derive(Clone)]
pub struct MQFaultStrategy {
    latency_fault_tolerance: Arc<LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>>,
    send_latency_fault_enable: Arc<AtomicBool>,
    start_detector_enable: Arc<AtomicBool>,
    latency_max: Vec<u64>,
    not_available_duration: Vec<u64>,
}

impl MQFaultStrategy {
    pub const DEFAULT_LATENCY_MAX: [u64; 7] = [50, 100, 550, 1800, 3000, 5000, 15000];
    pub const DEFAULT_NOT_AVAILABLE_DURATION: [u64; 7] = [0, 0, 2000, 5000, 6000, 10000, 30000];
    /// Latency penalty applied when a broker is isolated (10 seconds).
    const ISOLATION_LATENCY_MS: u64 = 10_000;

    pub fn new(service_context: ChildServiceContext, client_config: &ClientConfig) -> Self {
        let tolerance_impl = LatencyFaultToleranceImpl::new(service_context);
        Self::from_tolerance_impl(tolerance_impl, client_config)
    }

    #[cfg(any(test, feature = "test-support"))]
    pub(crate) fn new_with_fault_clock(
        service_context: ChildServiceContext,
        client_config: &ClientConfig,
        clock: Arc<dyn FaultClock>,
    ) -> Self {
        let tolerance_impl = LatencyFaultToleranceImpl::new_with_clock(service_context, clock);
        Self::from_tolerance_impl(tolerance_impl, client_config)
    }

    fn from_tolerance_impl(
        tolerance_impl: LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>,
        client_config: &ClientConfig,
    ) -> Self {
        tolerance_impl.set_detect_interval(client_config.detect_interval);
        tolerance_impl.set_detect_timeout(client_config.detect_timeout);
        tolerance_impl.set_start_detector_enable(client_config.start_detector_enable);
        let latency_fault_tolerance = Arc::new(tolerance_impl);
        Self {
            latency_fault_tolerance: Arc::clone(&latency_fault_tolerance),
            send_latency_fault_enable: Arc::new(AtomicBool::new(client_config.send_latency_enable)),
            start_detector_enable: Arc::new(AtomicBool::new(client_config.start_detector_enable)),
            latency_max: Self::DEFAULT_LATENCY_MAX.to_vec(),
            not_available_duration: Self::DEFAULT_NOT_AVAILABLE_DURATION.to_vec(),
        }
    }

    pub fn start_detector(&self) {
        LatencyFaultTolerance::start_detector(self.latency_fault_tolerance.clone());
    }

    pub(crate) fn set_resolve(&self, resolver: DefaultResolver) {
        self.latency_fault_tolerance.set_resolver(resolver);
    }

    pub(crate) fn set_service_detector(&self, service_detector: DefaultServiceDetector) {
        self.latency_fault_tolerance.set_service_detector(service_detector);
    }

    pub fn shutdown(&self) {
        self.latency_fault_tolerance.shutdown();
    }

    pub async fn shutdown_async(&self) -> bool {
        self.latency_fault_tolerance.shutdown_detector().await
    }

    pub fn is_start_detector_enable(&self) -> bool {
        self.start_detector_enable.load(Ordering::Relaxed)
    }

    pub fn set_start_detector_enable(&self, start_detector_enable: bool) {
        self.start_detector_enable
            .store(start_detector_enable, Ordering::Relaxed);
        self.latency_fault_tolerance
            .set_start_detector_enable(start_detector_enable);
    }

    pub fn is_send_latency_fault_enable(&self) -> bool {
        self.send_latency_fault_enable.load(Ordering::Relaxed)
    }

    pub fn select_one_message_queue(
        &self,
        tp_info: &TopicPublishInfo,
        last_broker_name: Option<&CheetahString>,
        reset_index: bool,
    ) -> Option<MessageQueue> {
        let broker_filter = BrokerFilter::with_last_broker_name(last_broker_name.cloned());

        if self.send_latency_fault_enable.load(Ordering::Relaxed) {
            if reset_index {
                tp_info.reset_index();
            }

            let available_filter = AvailableFilter {
                latency_fault_tolerance: Arc::clone(&self.latency_fault_tolerance),
            };
            let filter = &[
                &available_filter as &dyn QueueFilter,
                &broker_filter as &dyn QueueFilter,
            ];
            if let Some(mq) = tp_info.select_one_message_queue_filters(filter) {
                return Some(mq);
            }

            let reachable_filter = ReachableFilter {
                latency_fault_tolerance: Arc::clone(&self.latency_fault_tolerance),
            };
            let filter = &[
                &reachable_filter as &dyn QueueFilter,
                &broker_filter as &dyn QueueFilter,
            ];
            if let Some(mq) = tp_info.select_one_message_queue_filters(filter) {
                return Some(mq);
            }

            return tp_info.select_one_message_queue_filters(&[]);
        }

        if let Some(mq) = tp_info.select_one_message_queue_filters(&[&broker_filter]) {
            return Some(mq);
        }
        tp_info.select_one_message_queue_filters(&[])
    }

    pub fn get_latency_max(&self) -> &[u64] {
        &self.latency_max
    }

    pub fn set_latency_max(&mut self, latency_max: impl Into<Vec<u64>>) {
        self.latency_max = latency_max.into();
    }

    pub fn get_not_available_duration(&self) -> &[u64] {
        &self.not_available_duration
    }

    pub fn set_not_available_duration(&mut self, not_available_duration: impl Into<Vec<u64>>) {
        self.not_available_duration = not_available_duration.into();
    }

    pub async fn update_fault_item(
        &self,
        broker_name: CheetahString,
        current_latency: u64,
        isolation: bool,
        reachable: bool,
    ) {
        if self.send_latency_fault_enable.load(Ordering::Relaxed) {
            let effective_latency = if isolation {
                Self::ISOLATION_LATENCY_MS
            } else {
                current_latency
            };
            let duration = self.compute_not_available_duration(effective_latency);
            self.latency_fault_tolerance
                .update_fault_item(broker_name, current_latency, duration, reachable)
                .await;
        }
    }

    fn compute_not_available_duration(&self, current_latency: u64) -> u64 {
        let len = self.latency_max.len().min(self.not_available_duration.len());
        for i in (0..len).rev() {
            if current_latency >= self.latency_max[i] {
                return self.not_available_duration[i];
            }
        }
        0
    }

    #[cfg(any(test, feature = "test-support"))]
    pub(crate) fn not_available_duration_for_test(&self, current_latency: u64, isolation: bool) -> u64 {
        self.compute_not_available_duration(if isolation {
            Self::ISOLATION_LATENCY_MS
        } else {
            current_latency
        })
    }

    #[cfg(any(test, feature = "test-support"))]
    pub(crate) fn fault_item_state_for_test(&self, broker_name: &CheetahString) -> Option<(bool, bool, u64, u64)> {
        self.latency_fault_tolerance.fault_item_state_for_test(broker_name)
    }

    #[inline]
    pub fn set_send_latency_fault_enable(&self, send_latency_fault_enable: bool) {
        self.send_latency_fault_enable
            .store(send_latency_fault_enable, Ordering::Relaxed);
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct BrokerFilter {
    last_broker_name: Option<CheetahString>,
}

impl BrokerFilter {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_last_broker_name(last_broker_name: Option<CheetahString>) -> Self {
        Self { last_broker_name }
    }

    pub fn last_broker_name(&self) -> Option<&CheetahString> {
        self.last_broker_name.as_ref()
    }

    pub fn get_last_broker_name(&self) -> Option<&CheetahString> {
        self.last_broker_name()
    }

    pub fn set_last_broker_name(&mut self, last_broker_name: impl Into<CheetahString>) {
        self.last_broker_name = Some(last_broker_name.into());
    }

    pub fn clear_last_broker_name(&mut self) {
        self.last_broker_name = None;
    }
}

impl QueueFilter for BrokerFilter {
    fn filter(&self, message_queue: &MessageQueue) -> bool {
        if let Some(last_broker_name) = &self.last_broker_name {
            message_queue.broker_name() != last_broker_name
        } else {
            true
        }
    }
}

struct ReachableFilter {
    latency_fault_tolerance: Arc<LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>>,
}

impl QueueFilter for ReachableFilter {
    fn filter(&self, message_queue: &MessageQueue) -> bool {
        self.latency_fault_tolerance.is_reachable(message_queue.broker_name())
    }
}

struct AvailableFilter {
    latency_fault_tolerance: Arc<LatencyFaultToleranceImpl<DefaultResolver, DefaultServiceDetector>>,
}

impl QueueFilter for AvailableFilter {
    fn filter(&self, message_queue: &MessageQueue) -> bool {
        self.latency_fault_tolerance.is_available(message_queue.broker_name())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn broker_filter_matches_java_last_broker_name_rule() {
        let mq = MessageQueue::from_parts("topic", "broker-a", 0);
        let mut filter = BrokerFilter::new();

        assert!(filter.filter(&mq));
        assert!(filter.get_last_broker_name().is_none());

        filter.set_last_broker_name("broker-a");
        assert_eq!(
            filter.get_last_broker_name(),
            Some(&CheetahString::from_static_str("broker-a"))
        );
        assert!(!filter.filter(&mq));

        filter.set_last_broker_name("broker-b");
        assert!(filter.filter(&mq));

        filter.clear_last_broker_name();
        assert!(filter.filter(&mq));
    }

    #[test]
    fn mq_fault_strategy_copies_detector_config_like_java_constructor() {
        let mut client_config = ClientConfig::default();
        client_config.set_detect_interval(12_345);
        client_config.set_detect_timeout(321);
        client_config.set_start_detector_enable(true);

        let strategy = MQFaultStrategy::new(
            crate::runtime::test_service_context("mq-fault-strategy-test"),
            &client_config,
        );
        let (detect_interval, detect_timeout) = strategy.latency_fault_tolerance.detector_config_for_test();

        assert_eq!(detect_interval, 12_345);
        assert_eq!(detect_timeout, 321);
        assert!(strategy.is_start_detector_enable());
        assert!(strategy.latency_fault_tolerance.is_start_detector_enable());
    }

    #[test]
    fn cloned_strategy_shares_runtime_switches_and_snapshots_thresholds() {
        let mut strategy = MQFaultStrategy::new(
            crate::runtime::test_service_context("mq-fault-strategy-clone-test"),
            &ClientConfig::default(),
        );
        strategy.set_latency_max(vec![10, 20, 30]);
        strategy.set_not_available_duration(vec![0, 100, 200]);
        strategy.set_send_latency_fault_enable(true);
        strategy.set_start_detector_enable(true);

        let snapshot = strategy.clone();

        strategy.set_latency_max(vec![40, 50, 60]);
        strategy.set_not_available_duration(vec![300, 400, 500]);
        strategy.set_send_latency_fault_enable(false);
        strategy.set_start_detector_enable(false);

        assert_eq!(snapshot.get_latency_max(), &[10, 20, 30]);
        assert_eq!(snapshot.get_not_available_duration(), &[0, 100, 200]);
        assert!(!snapshot.is_send_latency_fault_enable());
        assert!(!snapshot.is_start_detector_enable());
    }
}