rocketmq-client-rust 0.9.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::collections::HashSet;

use dashmap::DashMap;
use tokio_util::sync::CancellationToken;

use crate::latency::latency_fault_tolerance::LatencyFaultTolerance;
use crate::latency::resolver::Resolver;
use crate::latency::service_detector::ServiceDetector;

pub struct LatencyFaultToleranceImpl<R, S> {
    fault_item_table: DashMap<CheetahString, FaultItem>,
    detect_timeout: u32,
    detect_interval: u32,
    start_detector_enable: AtomicBool,
    resolver: Option<R>,
    service_detector: Option<S>,
    cancel_token: CancellationToken,
}

impl<R, S> LatencyFaultToleranceImpl<R, S> {
    pub fn new() -> Self {
        Self {
            resolver: None,
            service_detector: None,
            fault_item_table: Default::default(),
            detect_timeout: 200,
            detect_interval: 2000,
            start_detector_enable: AtomicBool::new(false),
            cancel_token: CancellationToken::new(),
        }
    }

    pub fn set_resolver(&mut self, resolver: R) {
        self.resolver = Some(resolver);
    }

    pub fn set_service_detector(&mut self, service_detector: S) {
        self.service_detector = Some(service_detector);
    }

    #[cfg(test)]
    pub(crate) fn detector_config_for_test(&self) -> (u32, u32) {
        (self.detect_interval, self.detect_timeout)
    }
}

impl<R, S> LatencyFaultTolerance<CheetahString> for LatencyFaultToleranceImpl<R, S>
where
    R: Resolver,
    S: ServiceDetector,
{
    async fn update_fault_item(
        &self,
        name: CheetahString,
        current_latency: u64,
        not_available_duration: u64,
        reachable: bool,
    ) {
        let entry = self
            .fault_item_table
            .entry(name.clone())
            .or_insert_with(|| FaultItem::new(name.clone()));

        entry.set_current_latency(current_latency);
        entry.update_not_available_duration(not_available_duration);
        entry.set_reachable(reachable);

        if !reachable {
            info!("{} is unreachable, it will not be used until it's reachable", name);
        }
    }

    #[inline]
    fn is_available(&self, name: &CheetahString) -> bool {
        self.fault_item_table.get(name).is_none_or(|item| item.is_available())
    }

    #[inline]
    fn is_reachable(&self, name: &CheetahString) -> bool {
        self.fault_item_table.get(name).is_none_or(|item| item.is_reachable())
    }

    async fn remove(&self, name: &CheetahString) {
        self.fault_item_table.remove(name);
    }

    async fn pick_one_at_least(&self) -> Option<CheetahString> {
        use smallvec::SmallVec;

        let mut reachable_names: SmallVec<[CheetahString; 4]> = SmallVec::new();

        for entry in self.fault_item_table.iter() {
            if entry.value().is_reachable() {
                reachable_names.push(entry.key().clone());
            }
        }

        if !reachable_names.is_empty() {
            use rand::seq::SliceRandom;
            reachable_names.shuffle(&mut rand::rng());
            return Some(reachable_names[0].clone());
        }
        None
    }

    fn start_detector(this: ArcMut<Self>) {
        let token = this.cancel_token.clone();
        spawn_detector_loop(this, token);
    }

    fn shutdown(&self) {
        self.cancel_token.cancel();
    }

    async fn detect_by_one_round(&self) {
        let mut remove_set = HashSet::new();

        for entry in self.fault_item_table.iter() {
            let (name, fault_item) = (entry.key(), entry.value());
            if current_millis() as i64 - (fault_item.check_stamp.load(std::sync::atomic::Ordering::Relaxed) as i64) < 0
            {
                continue;
            }
            fault_item.check_stamp.store(
                current_millis() + self.detect_interval as u64,
                std::sync::atomic::Ordering::Release,
            );

            let broker_addr = match self.resolver.as_ref() {
                Some(resolver) => match resolver.resolve(fault_item.name.as_ref()).await {
                    Some(addr) => addr,
                    None => {
                        remove_set.insert(name.clone());
                        continue;
                    }
                },
                None => continue,
            };

            let service_detector = match self.service_detector.as_ref() {
                Some(d) => d,
                None => continue,
            };

            let service_ok = service_detector
                .detect(broker_addr.as_str(), self.detect_timeout as u64)
                .await;

            if service_ok && !fault_item.is_reachable() {
                info!("{} is reachable now, then it can be used.", name);
                fault_item.set_reachable(true);
            }
        }

        for name in remove_set {
            self.fault_item_table.remove(&name);
        }
    }

    fn set_detect_timeout(&mut self, detect_timeout: u32) {
        self.detect_timeout = detect_timeout;
    }

    fn set_detect_interval(&mut self, detect_interval: u32) {
        self.detect_interval = detect_interval;
    }

    fn set_start_detector_enable(&mut self, start_detector_enable: bool) {
        self.start_detector_enable
            .store(start_detector_enable, std::sync::atomic::Ordering::Relaxed);
    }

    fn is_start_detector_enable(&self) -> bool {
        self.start_detector_enable.load(std::sync::atomic::Ordering::Acquire)
    }
}

use std::cmp::Ordering;
use std::hash::Hash;
use std::sync::atomic::AtomicBool;

use cheetah_string::CheetahString;
use rocketmq_common::TimeUtils::current_millis;
use rocketmq_rust::ArcMut;
use tracing::error;
use tracing::info;
use tracing::warn;

async fn run_detector_loop<R, S>(this: ArcMut<LatencyFaultToleranceImpl<R, S>>, token: CancellationToken)
where
    R: Resolver,
    S: ServiceDetector,
{
    let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(3));
    // First tick completes immediately; skip it to match Java's 3s initial delay.
    interval.tick().await;

    loop {
        tokio::select! {
            _ = token.cancelled() => {
                tracing::info!("Fault tolerance detector stopped");
                break;
            }
            _ = interval.tick() => {}
        }
        if !this.start_detector_enable.load(std::sync::atomic::Ordering::Relaxed) {
            continue;
        }

        this.detect_by_one_round().await;
    }
}

fn spawn_detector_loop<R, S>(this: ArcMut<LatencyFaultToleranceImpl<R, S>>, token: CancellationToken)
where
    R: Resolver,
    S: ServiceDetector,
{
    match tokio::runtime::Handle::try_current() {
        Ok(handle) => {
            let task = handle.spawn(run_detector_loop(this, token));
            drop(task);
        }
        Err(error) => {
            warn!(
                "No Tokio runtime available while starting fault tolerance detector; using a dedicated detector \
                 thread. error={}",
                error
            );
            let thread = std::thread::Builder::new().name("rocketmq-client-fault-detector".to_string());
            if let Err(error) = thread.spawn(move || {
                let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
                    Ok(runtime) => runtime,
                    Err(error) => {
                        error!("Failed to create fault tolerance detector runtime: {}", error);
                        return;
                    }
                };
                runtime.block_on(run_detector_loop(this, token));
            }) {
                error!("Failed to spawn fault tolerance detector thread: {}", error);
            }
        }
    }
}

#[repr(C)]
#[derive(Debug)]
pub struct FaultItem {
    reachable_flag: std::sync::atomic::AtomicBool,
    _padding1: [u8; 7],
    start_timestamp: std::sync::atomic::AtomicU64,
    current_latency: std::sync::atomic::AtomicU64,

    check_stamp: std::sync::atomic::AtomicU64,
    name: CheetahString,
}

impl FaultItem {
    pub fn new(name: CheetahString) -> Self {
        FaultItem {
            reachable_flag: std::sync::atomic::AtomicBool::new(true),
            _padding1: [0; 7],
            start_timestamp: std::sync::atomic::AtomicU64::new(0),
            current_latency: std::sync::atomic::AtomicU64::new(0),
            check_stamp: std::sync::atomic::AtomicU64::new(0),
            name,
        }
    }

    pub fn update_not_available_duration(&self, not_available_duration: u64) {
        let now = current_millis();
        if not_available_duration > 0
            && now + not_available_duration > self.start_timestamp.load(std::sync::atomic::Ordering::Relaxed)
        {
            self.start_timestamp
                .store(now + not_available_duration, std::sync::atomic::Ordering::Release);
            info!("{} will be isolated for {} ms.", self.name, not_available_duration);
        }
    }

    pub fn set_reachable(&self, reachable_flag: bool) {
        self.reachable_flag
            .store(reachable_flag, std::sync::atomic::Ordering::Release);
    }

    pub fn set_check_stamp(&self, check_stamp: u64) {
        self.check_stamp
            .store(check_stamp, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn is_available(&self) -> bool {
        let now = current_millis();
        now >= self.start_timestamp.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn is_reachable(&self) -> bool {
        self.reachable_flag.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn set_current_latency(&self, latency: u64) {
        self.current_latency
            .store(latency, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn get_current_latency(&self) -> u64 {
        self.current_latency.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_start_timestamp(&self) -> u64 {
        self.start_timestamp.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn compare_to(&self, other: &Self) -> i32 {
        match self.cmp(other) {
            Ordering::Less => -1,
            Ordering::Equal => 0,
            Ordering::Greater => 1,
        }
    }
}

impl Eq for FaultItem {}

impl Ord for FaultItem {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.is_available() != other.is_available() {
            if self.is_available() {
                return Ordering::Less;
            }
            if other.is_available() {
                return Ordering::Greater;
            }
        }

        match self
            .current_latency
            .load(std::sync::atomic::Ordering::Relaxed)
            .cmp(&other.current_latency.load(std::sync::atomic::Ordering::Relaxed))
        {
            Ordering::Equal => (),
            ord => return ord,
        }

        match self
            .start_timestamp
            .load(std::sync::atomic::Ordering::Relaxed)
            .cmp(&other.start_timestamp.load(std::sync::atomic::Ordering::Relaxed))
        {
            Ordering::Equal => (),
            ord => return ord,
        }

        Ordering::Equal
    }
}

impl PartialEq<Self> for FaultItem {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
            && self.current_latency.load(std::sync::atomic::Ordering::Relaxed)
                == other.current_latency.load(std::sync::atomic::Ordering::Relaxed)
            && self.start_timestamp.load(std::sync::atomic::Ordering::Relaxed)
                == other.start_timestamp.load(std::sync::atomic::Ordering::Relaxed)
    }
}

impl PartialOrd for FaultItem {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Hash for FaultItem {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        self.current_latency
            .load(std::sync::atomic::Ordering::Relaxed)
            .hash(state);
        self.start_timestamp
            .load(std::sync::atomic::Ordering::Relaxed)
            .hash(state);
    }
}

impl std::fmt::Display for FaultItem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "FaultItem{{ name='{}', current_latency={}, start_timestamp={}, reachable_flag={} }}",
            self.name,
            self.get_current_latency(),
            self.get_start_timestamp(),
            self.is_reachable()
        )
    }
}

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

    struct NoopResolver;

    impl Resolver for NoopResolver {
        async fn resolve(&self, _name: &CheetahString) -> Option<CheetahString> {
            None
        }
    }

    struct NoopServiceDetector;

    impl ServiceDetector for NoopServiceDetector {
        async fn detect(&self, _endpoint: &str, _timeout_millis: u64) -> bool {
            true
        }
    }

    #[test]
    fn fault_item_compare_to_matches_java_comparable_shape() {
        let fast = FaultItem::new(CheetahString::from("fast"));
        let slow = FaultItem::new(CheetahString::from("slow"));
        fast.set_current_latency(10);
        slow.set_current_latency(20);

        assert_eq!(fast.compare_to(&slow), -1);
        assert_eq!(slow.compare_to(&fast), 1);
        assert_eq!(fast.compare_to(&fast), 0);
    }

    #[test]
    fn start_detector_without_tokio_runtime_does_not_spawn_panic() {
        let detector = ArcMut::new(LatencyFaultToleranceImpl::<NoopResolver, NoopServiceDetector>::new());

        LatencyFaultTolerance::start_detector(detector.clone());
        detector.shutdown();

        std::thread::sleep(tokio::time::Duration::from_millis(20));
    }
}