rocketmq_protocol/protocol/
route_facade.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
9use crate::protocol::route::route_data_view::BrokerData;
10use cheetah_string::CheetahString;
11use rand::RngExt;
12
13pub fn select_broker_addr(broker_data: &BrokerData) -> Option<CheetahString> {
14    let index = if broker_data.broker_addrs().is_empty() {
15        0
16    } else {
17        rand::rng().random_range(0..broker_data.broker_addrs().len())
18    };
19    broker_data.select_broker_addr_with_index(index)
20}
21
22pub trait BrokerDataExt {
23    fn select_broker_addr(&self) -> Option<CheetahString>;
24}
25
26impl BrokerDataExt for BrokerData {
27    fn select_broker_addr(&self) -> Option<CheetahString> {
28        select_broker_addr(self)
29    }
30}
31
32#[cfg(test)]
33mod tests {
34    use super::*;
35    use crate::common::wire_constants::MASTER_ID;
36    use std::collections::HashMap;
37
38    fn create_test_broker_data(addrs: HashMap<u64, CheetahString>) -> BrokerData {
39        BrokerData::new(
40            CheetahString::from("test_cluster"),
41            CheetahString::from("test_broker"),
42            addrs,
43            None,
44        )
45    }
46
47    fn make_addr(id: u64, addr: &str) -> (u64, CheetahString) {
48        (id, CheetahString::from(addr))
49    }
50
51    #[test]
52    fn test_empty_broker_address_map() {
53        let broker_data = create_test_broker_data(HashMap::new());
54        assert!(select_broker_addr(&broker_data).is_none());
55        assert!(broker_data.select_broker_addr().is_none())
56    }
57
58    #[test]
59    fn test_map_containing_only_one_slave() {
60        let slave_id = MASTER_ID + 1;
61        let slave_addr = CheetahString::from("192.168.1.101");
62        let mut broker_addrs = HashMap::new();
63        broker_addrs.insert(slave_id, slave_addr.clone());
64        let broker_data = create_test_broker_data(broker_addrs);
65        assert_eq!(broker_data.select_broker_addr(), Some(slave_addr.clone()));
66        assert_eq!(select_broker_addr(&broker_data), Some(slave_addr));
67    }
68
69    #[test]
70    fn test_map_containing_master_and_slaves() {
71        let master_addr = CheetahString::from("192.168.1.1");
72        let broker_addrs = HashMap::from([
73            make_addr(MASTER_ID + 1, "192.168.1.101"),
74            make_addr(MASTER_ID + 2, "192.168.1.102"),
75            (MASTER_ID, master_addr.clone()),
76        ]);
77        let broker_data = create_test_broker_data(broker_addrs);
78        assert_eq!(broker_data.select_broker_addr(), Some(master_addr.clone()));
79        assert_eq!(select_broker_addr(&broker_data), Some(master_addr));
80    }
81
82    #[test]
83    fn test_map_containing_multiple_slaves() {
84        let broker_addrs: HashMap<u64, CheetahString> = (1..=5)
85            .map(|id| (id, CheetahString::from(format!("192.168.1.10{}", id))))
86            .collect();
87        let broker_data = create_test_broker_data(broker_addrs.clone());
88        let selected = broker_data.select_broker_addr();
89        assert!(selected.is_some());
90        assert!(broker_addrs.values().any(|v| v == selected.as_ref().unwrap()));
91        let selected_free = select_broker_addr(&broker_data);
92        assert!(selected_free.is_some());
93        assert!(broker_addrs.values().any(|v| v == selected_free.as_ref().unwrap()));
94    }
95}
rocketmq_protocol/protocol/route_facade.rs

rocketmq_protocol/protocol/
route_facade.rs
