rocketmq_protocol/protocol/route/

route_data_view.rs

// 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::cmp::Ordering;
use std::collections::HashMap;

use crate::common::wire_constants::MASTER_ID;
use cheetah_string::CheetahString;
use serde::Deserialize;
use serde::Serialize;

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct BrokerData {
    cluster: CheetahString,
    #[serde(rename = "brokerName")]
    broker_name: CheetahString,
    #[serde(rename = "brokerAddrs")]
    broker_addrs: HashMap<u64 /* broker id */, CheetahString /* broker ip */>,
    #[serde(rename = "zoneName")]
    zone_name: Option<CheetahString>,
    #[serde(rename = "enableActingMaster")]
    enable_acting_master: bool,
}

impl PartialEq for BrokerData {
    fn eq(&self, other: &Self) -> bool {
        self.broker_name == other.broker_name && self.broker_addrs == other.broker_addrs
    }
}

impl Eq for BrokerData {}

impl PartialOrd for BrokerData {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for BrokerData {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        self.broker_name.cmp(&other.broker_name)
    }
}

impl BrokerData {
    pub fn new(
        cluster: CheetahString,
        broker_name: CheetahString,
        broker_addrs: HashMap<u64, CheetahString>,
        zone_name: Option<CheetahString>,
    ) -> BrokerData {
        BrokerData {
            cluster,
            broker_name,
            broker_addrs,
            zone_name,
            enable_acting_master: false,
        }
    }

    #[inline]
    pub fn set_cluster(&mut self, cluster: CheetahString) {
        self.cluster = cluster;
    }

    #[inline]
    pub fn set_broker_name(&mut self, broker_name: CheetahString) {
        self.broker_name = broker_name;
    }

    #[inline]
    pub fn set_broker_addrs(&mut self, broker_addrs: HashMap<u64, CheetahString>) {
        self.broker_addrs = broker_addrs;
    }

    #[inline]
    pub fn set_zone_name(&mut self, zone_name: Option<CheetahString>) {
        self.zone_name = zone_name;
    }

    #[inline]
    pub fn set_enable_acting_master(&mut self, enable_acting_master: bool) {
        self.enable_acting_master = enable_acting_master;
    }

    #[inline]
    pub fn cluster(&self) -> &str {
        &self.cluster
    }

    #[inline]
    pub fn broker_name(&self) -> &CheetahString {
        &self.broker_name
    }

    #[inline]
    pub fn broker_addrs(&self) -> &HashMap<u64, CheetahString> {
        &self.broker_addrs
    }

    #[inline]
    pub fn broker_addrs_mut(&mut self) -> &mut HashMap<u64, CheetahString> {
        &mut self.broker_addrs
    }

    #[inline]
    pub fn remove_broker_by_addr(&mut self, broker_id: u64, broker_addr: &CheetahString) {
        self.broker_addrs
            .retain(|key, value| value != broker_addr || *key == broker_id);
    }

    #[inline]
    pub fn zone_name(&self) -> Option<&CheetahString> {
        self.zone_name.as_ref()
    }

    #[inline]
    pub fn enable_acting_master(&self) -> bool {
        self.enable_acting_master
    }

    pub fn select_broker_addr_with_index(&self, index: usize) -> Option<CheetahString> {
        if let Some(master_address) = self.broker_addrs.get(&MASTER_ID) {
            return Some(master_address.clone());
        }
        let mut slaves = self.broker_addrs.iter().collect::<Vec<_>>();
        slaves.sort_unstable_by_key(|(id, _)| **id);
        if slaves.is_empty() {
            return None;
        }
        slaves.get(index % slaves.len()).map(|(_, value)| (*value).clone())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd, Default)]
pub struct QueueData {
    #[serde(rename = "brokerName")]
    pub broker_name: CheetahString,
    #[serde(rename = "readQueueNums")]
    pub read_queue_nums: u32,
    #[serde(rename = "writeQueueNums")]
    pub write_queue_nums: u32,
    pub perm: u32,
    #[serde(rename = "topicSysFlag")]
    pub topic_sys_flag: u32,
}

impl QueueData {
    pub fn new(
        broker_name: CheetahString,
        read_queue_nums: u32,
        write_queue_nums: u32,
        perm: u32,
        topic_sys_flag: u32,
    ) -> Self {
        Self {
            broker_name,
            read_queue_nums,
            write_queue_nums,
            perm,
            topic_sys_flag,
        }
    }

    #[inline]
    pub fn broker_name(&self) -> &CheetahString {
        &self.broker_name
    }

    #[inline]
    pub fn read_queue_nums(&self) -> u32 {
        self.read_queue_nums
    }

    #[inline]
    pub fn write_queue_nums(&self) -> u32 {
        self.write_queue_nums
    }

    #[inline]
    pub fn perm(&self) -> u32 {
        self.perm
    }

    #[inline]
    pub fn topic_sys_flag(&self) -> u32 {
        self.topic_sys_flag
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;

    #[test]
    fn broker_data_methods_preserve_route_fields() {
        let mut broker_data = BrokerData::new(
            CheetahString::from("cluster1"),
            CheetahString::from("broker1"),
            HashMap::new(),
            None,
        );
        assert_eq!(broker_data.cluster(), "cluster1");
        assert_eq!(broker_data.broker_name(), "broker1");
        assert!(broker_data.broker_addrs().is_empty());
        assert_eq!(broker_data.zone_name(), None);
        assert!(!broker_data.enable_acting_master());

        broker_data.set_cluster(CheetahString::from("cluster2"));
        broker_data.set_broker_name(CheetahString::from("broker2"));
        broker_data.set_broker_addrs(HashMap::from([(1, CheetahString::from("127.0.0.1"))]));
        broker_data
            .broker_addrs_mut()
            .insert(2, CheetahString::from("127.0.0.2"));
        broker_data.set_zone_name(Some(CheetahString::from("zone1")));
        broker_data.set_enable_acting_master(true);

        assert_eq!(broker_data.cluster(), "cluster2");
        assert_eq!(broker_data.broker_name(), "broker2");
        assert_eq!(
            broker_data.broker_addrs(),
            &HashMap::from([
                (1, CheetahString::from("127.0.0.1")),
                (2, CheetahString::from("127.0.0.2")),
            ])
        );
        assert_eq!(broker_data.zone_name(), Some(&CheetahString::from("zone1")));
        assert!(broker_data.enable_acting_master());
    }

    #[test]
    fn remove_broker_by_addr_keeps_the_current_id_only() {
        let mut broker_data = BrokerData::new(
            CheetahString::from("cluster1"),
            CheetahString::from("broker1"),
            HashMap::from([
                (1, CheetahString::from("127.0.0.1")),
                (2, CheetahString::from("127.0.0.1")),
                (3, CheetahString::from("127.0.0.3")),
            ]),
            None,
        );

        broker_data.remove_broker_by_addr(1, &"127.0.0.1".into());
        assert_eq!(
            broker_data.broker_addrs(),
            &HashMap::from([
                (1, CheetahString::from("127.0.0.1")),
                (3, CheetahString::from("127.0.0.3")),
            ])
        );
    }

    #[test]
    fn select_broker_addr_prefers_master_and_indexes_slaves() {
        let mut broker_data = BrokerData::new(
            CheetahString::from("cluster1"),
            CheetahString::from("broker1"),
            HashMap::new(),
            None,
        );
        assert_eq!(broker_data.select_broker_addr_with_index(0), None);

        broker_data.set_broker_addrs(HashMap::from([
            (4, CheetahString::from("127.0.0.4")),
            (2, CheetahString::from("127.0.0.2")),
        ]));
        assert_eq!(
            broker_data.select_broker_addr_with_index(0),
            Some(CheetahString::from("127.0.0.2"))
        );
        assert_eq!(
            broker_data.select_broker_addr_with_index(1),
            Some(CheetahString::from("127.0.0.4"))
        );
        assert_eq!(
            broker_data.select_broker_addr_with_index(2),
            Some(CheetahString::from("127.0.0.2"))
        );

        broker_data
            .broker_addrs_mut()
            .insert(MASTER_ID, CheetahString::from("127.0.0.1"));
        assert_eq!(
            broker_data.select_broker_addr_with_index(1),
            Some(CheetahString::from("127.0.0.1"))
        );
    }

    #[test]
    fn queue_data_methods_preserve_route_fields() {
        let queue_data = QueueData::new(CheetahString::from("broker1"), 4, 4, 6, 0);

        assert_eq!(queue_data.broker_name(), "broker1");
        assert_eq!(queue_data.read_queue_nums(), 4);
        assert_eq!(queue_data.write_queue_nums(), 4);
        assert_eq!(queue_data.perm(), 6);
        assert_eq!(queue_data.topic_sys_flag(), 0);
    }
}