rocketmq_protocol/protocol/static_topic/
topic_queue_mapping_utils.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::cmp::Ordering;
16use std::collections::HashMap;
17use std::collections::HashSet;
18
19use cheetah_string::CheetahString;
20use rand::seq::SliceRandom;
21use rocketmq_model::common::config::TopicConfig;
22use rocketmq_model::common::mix_all;
23use rocketmq_model::time::current_millis;
24
25use rocketmq_error::Error;
26use rocketmq_error::Result;
27
28use crate::protocol::static_topic::logic_queue_mapping_item::LogicQueueMappingItem;
29use crate::protocol::static_topic::topic_config_and_queue_mapping::TopicConfigAndQueueMapping;
30use crate::protocol::static_topic::topic_queue_mapping_detail::TopicQueueMappingDetail;
31use crate::protocol::static_topic::topic_queue_mapping_info::TopicQueueMappingInfo;
32use crate::protocol::static_topic::topic_queue_mapping_one::TopicQueueMappingOne;
33use crate::protocol::static_topic::topic_remapping_detail_wrapper;
34use crate::protocol::static_topic::topic_remapping_detail_wrapper::TopicRemappingDetailWrapper;
35use crate::protocol::RemotingSerializable;
36
37pub struct TopicQueueMappingUtils;
38
39fn static_topic_mapping_invalid_item(reason: impl Into<String>) -> Error {
40    crate::error::invalid_argument(reason)
41}
42
43fn static_topic_mapping_inconsistent(reason: impl Into<String>) -> Error {
44    crate::error::route_inconsistent("static_topic_mapping", reason)
45}
46
47impl TopicQueueMappingUtils {
48    pub fn find_logic_queue_mapping_item(
49        mapping_items: &[LogicQueueMappingItem],
50        logic_offset: i64,
51        ignore_negative: bool,
52    ) -> Option<&LogicQueueMappingItem> {
53        if mapping_items.is_empty() {
54            return None;
55        }
56        // Could use binary search to improve performance
57        for i in (0..mapping_items.len()).rev() {
58            let item = &mapping_items[i];
59            if ignore_negative && item.logic_offset < 0 {
60                continue;
61            }
62            if logic_offset >= item.logic_offset {
63                return Some(item);
64            }
65        }
66        // If not found, maybe out of range, return the first one
67        for item in mapping_items.iter() {
68            if ignore_negative && item.logic_offset < 0 {
69                continue;
70            } else {
71                return Some(item);
72            }
73        }
74        None
75    }
76
77    pub fn find_next<'a>(
78        items: &'a [LogicQueueMappingItem],
79        current_item: Option<&'a LogicQueueMappingItem>,
80        ignore_negative: bool,
81    ) -> Option<&'a LogicQueueMappingItem> {
82        if items.is_empty() || current_item.is_none() {
83            return None;
84        }
85        let current_item = current_item?;
86        for i in 0..items.len() {
87            let item = &items[i];
88            if ignore_negative && item.logic_offset < 0 {
89                continue;
90            }
91            if item.gen == current_item.gen {
92                if i < items.len() - 1 {
93                    let next_item = &items[i + 1];
94                    if ignore_negative && next_item.logic_offset < 0 {
95                        return None;
96                    } else {
97                        return Some(next_item);
98                    }
99                } else {
100                    return None;
101                }
102            }
103        }
104        None
105    }
106
107    pub fn get_mock_broker_name(scope: &str) -> CheetahString {
108        assert!(!scope.is_empty(), "Scope cannot be null");
109
110        if scope == mix_all::METADATA_SCOPE_GLOBAL {
111            format!("{}{}", mix_all::LOGICAL_QUEUE_MOCK_BROKER_PREFIX, &scope[2..]).into()
112        } else {
113            format!("{}{}", mix_all::LOGICAL_QUEUE_MOCK_BROKER_PREFIX, scope).into()
114        }
115    }
116    pub fn get_mapping_detail_from_config(
117        configs: Vec<TopicConfigAndQueueMapping>,
118    ) -> Result<Vec<TopicQueueMappingDetail>> {
119        let mut detail_list = vec![];
120        for config_mapping in &configs {
121            if let Some(detail) = &config_mapping.topic_queue_mapping_detail {
122                detail_list.push(detail.clone());
123            }
124        }
125        Ok(detail_list)
126    }
127    pub fn check_name_epoch_num_consistence(
128        topic: &CheetahString,
129        broker_config_map: &HashMap<CheetahString, TopicConfigAndQueueMapping>,
130    ) -> Result<(i64, i32)> {
131        if broker_config_map.is_empty() {
132            return Err(static_topic_mapping_inconsistent(
133                "check_name_epoch_num_consistence get empty config map",
134            ));
135        }
136        //make sure it is not null
137        let mut max_epoch = -1;
138        let mut max_num = -1;
139        let scope = CheetahString::new();
140        for entry in broker_config_map {
141            let broker = entry.0;
142            let config_mapping = entry.1;
143            if config_mapping.topic_queue_mapping_detail.is_none() {
144                return Err(static_topic_mapping_inconsistent(format!(
145                    "Mapping info should not be null in broker {}",
146                    broker
147                )));
148            }
149            let mapping_detail = &config_mapping.topic_queue_mapping_detail;
150            if let Some(mapping_detail) = &mapping_detail {
151                if let Some(broker_name) = &mapping_detail.topic_queue_mapping_info.bname {
152                    if !broker.eq(broker_name) {
153                        return Err(static_topic_mapping_inconsistent(format!(
154                            "The broker name is not equal {} != {} ",
155                            broker, broker_name
156                        )));
157                    }
158                    if mapping_detail.topic_queue_mapping_info.dirty {
159                        return Err(static_topic_mapping_inconsistent(format!(
160                            "The mapping info is dirty in broker  {}",
161                            broker
162                        )));
163                    }
164                    if let Some(top) = &config_mapping.topic_config.topic_name {
165                        if let Some(mapping_top) = &mapping_detail.topic_queue_mapping_info.topic {
166                            if !top.eq(mapping_top) {
167                                return Err(static_topic_mapping_inconsistent(format!(
168                                    "The topic name is inconsistent in broker  {}",
169                                    broker
170                                )));
171                            }
172                            if !topic.eq(mapping_top) {
173                                return Err(static_topic_mapping_inconsistent(format!(
174                                    "The topic name is not match for broker  {}",
175                                    broker
176                                )));
177                            }
178                            if let Some(m_scope) = &mapping_detail.topic_queue_mapping_info.scope {
179                                if !scope.eq(m_scope) {
180                                    return Err(static_topic_mapping_inconsistent(format!(
181                                        "scope does not match {} != {} in {}",
182                                        m_scope, scope, broker
183                                    )));
184                                }
185                                if max_epoch != -1 && max_epoch != mapping_detail.topic_queue_mapping_info.epoch {
186                                    return Err(static_topic_mapping_inconsistent(format!(
187                                        "epoch does not match {} != {} in {}",
188                                        max_epoch, mapping_detail.topic_queue_mapping_info.epoch, broker_name
189                                    )));
190                                } else {
191                                    max_epoch = mapping_detail.topic_queue_mapping_info.epoch;
192                                }
193                                if max_num != -1 && max_num != mapping_detail.topic_queue_mapping_info.total_queues {
194                                    return Err(static_topic_mapping_inconsistent(format!(
195                                        "total queue number does not match {} != {} in {}",
196                                        max_num, mapping_detail.topic_queue_mapping_info.total_queues, broker_name
197                                    )));
198                                } else {
199                                    max_num = mapping_detail.topic_queue_mapping_info.total_queues;
200                                }
201                                return Ok((max_epoch, max_num));
202                            }
203                        }
204                    }
205                }
206            }
207        }
208        Err(static_topic_mapping_inconsistent(
209            "check_name_epoch_num_consistence err ! maybe some var is none",
210        ))
211    }
212    pub fn check_if_reuse_physical_queue(mapping_ones: &Vec<TopicQueueMappingOne>) -> Result<()> {
213        let mut physical_queue_id_map: HashMap<String, TopicQueueMappingOne> = HashMap::new();
214        for mapping_one in mapping_ones {
215            for item in mapping_one.items() {
216                if let Some(bname) = &item.bname {
217                    let physical_queue_id = format!("{} - {}", bname, item.queue_id);
218                    if let Some(id) = physical_queue_id_map.get(&physical_queue_id) {
219                        return Err(static_topic_mapping_inconsistent(format!(
220                            "Topic {} global queue id {} and {} shared the same physical queue {}",
221                            mapping_one.topic(),
222                            mapping_one.global_id(),
223                            id.global_id(),
224                            physical_queue_id
225                        )));
226                    }
227                    physical_queue_id_map.insert(physical_queue_id, mapping_one.clone());
228                }
229            }
230        }
231        Ok(())
232    }
233
234    pub fn check_logic_queue_mapping_item_offset(items: &[LogicQueueMappingItem]) -> Result<()> {
235        if items.is_empty() {
236            return Ok(());
237        }
238        let mut last_gen = -1;
239        let mut last_offset = -1;
240        for i in (0..items.len()).rev() {
241            let item = &items[i];
242            if item.start_offset < 0 || item.gen < 0 || item.queue_id < 0 {
243                return Err(static_topic_mapping_invalid_item(
244                    "The field is illegal, should not be negative",
245                ));
246            }
247            if items.len() >= 2 && i <= items.len() - 2 && items[i].logic_offset < 0 {
248                return Err(static_topic_mapping_invalid_item(
249                    "The non-latest item has negative logic offset",
250                ));
251            }
252            if last_gen != -1 && item.gen >= last_gen {
253                return Err(static_topic_mapping_invalid_item(
254                    "The gen does not increase monotonically",
255                ));
256            }
257
258            if item.end_offset != -1 && item.end_offset < item.start_offset {
259                return Err(static_topic_mapping_invalid_item(
260                    "The endOffset is smaller than the start offset",
261                ));
262            }
263
264            if last_offset != -1 && item.logic_offset != -1 {
265                if item.logic_offset >= last_offset {
266                    return Err(static_topic_mapping_invalid_item(
267                        "The base logic offset does not increase monotonically",
268                    ));
269                }
270                if item.compute_max_static_queue_offset() >= last_offset {
271                    return Err(static_topic_mapping_invalid_item(
272                        "The max logic offset does not increase monotonically",
273                    ));
274                }
275            }
276            last_gen = item.gen;
277            last_offset = item.logic_offset;
278        }
279        Ok(())
280    }
281
282    pub fn check_if_leader(items: &[LogicQueueMappingItem], mapping_detail: &TopicQueueMappingDetail) -> bool {
283        let Some(item) = items.last() else {
284            return false;
285        };
286        item.bname.as_ref() == mapping_detail.topic_queue_mapping_info.bname.as_ref()
287    }
288
289    pub fn get_leader_item(items: &[LogicQueueMappingItem]) -> Result<LogicQueueMappingItem> {
290        if items.is_empty() {
291            return Err(static_topic_mapping_inconsistent(
292                "get_leader_item failed with empty items",
293            ));
294        }
295        if let Some(i) = items.last() {
296            return Ok(i.clone());
297        }
298        Err(static_topic_mapping_inconsistent(
299            "get_leader_item failed with empty items",
300        ))
301    }
302    pub fn get_leader_broker(items: &[LogicQueueMappingItem]) -> Result<CheetahString> {
303        let item = TopicQueueMappingUtils::get_leader_item(items)?;
304        if let Some(bname) = &item.bname {
305            return Ok(bname.to_string().into());
306        }
307        Err(static_topic_mapping_inconsistent(
308            "get_leader_broker fn get Item with None bname",
309        ))
310    }
311    pub fn check_and_build_mapping_items(
312        mut mapping_detail_list: Vec<TopicQueueMappingDetail>,
313        replace: bool,
314        check_consistence: bool,
315    ) -> Result<HashMap<i32, TopicQueueMappingOne>> {
316        mapping_detail_list.sort_by(|o1, o2| {
317            let sub = o1.topic_queue_mapping_info.epoch - o2.topic_queue_mapping_info.epoch;
318            if sub > 0 {
319                return std::cmp::Ordering::Greater;
320            } else if sub < 0 {
321                return std::cmp::Ordering::Less;
322            }
323            std::cmp::Ordering::Equal
324        });
325
326        let mut max_num = 0;
327        let mut global_id_map = HashMap::new();
328        for mapping_detail in &mapping_detail_list {
329            if mapping_detail.topic_queue_mapping_info.total_queues > max_num {
330                max_num = mapping_detail.topic_queue_mapping_info.total_queues;
331            }
332            if let Some(queue) = &mapping_detail.hosted_queues {
333                for entry in queue {
334                    let global_id = entry.0;
335                    TopicQueueMappingUtils::check_logic_queue_mapping_item_offset(entry.1)?;
336                    let leader_broker_name = TopicQueueMappingUtils::get_leader_broker(entry.1)?;
337                    if let Some(broker_name) = &mapping_detail.topic_queue_mapping_info.bname {
338                        if !leader_broker_name.eq(broker_name) {
339                            //not the leader
340                            continue;
341                        }
342
343                        if global_id_map.contains_key(global_id) {
344                            if !replace {
345                                return Err(static_topic_mapping_inconsistent(format!(
346                                    "The queue id is duplicated in broker {} {}",
347                                    leader_broker_name, broker_name
348                                )));
349                            }
350                        } else if let Some(top) = &mapping_detail.topic_queue_mapping_info.topic {
351                            global_id_map.insert(
352                                *global_id,
353                                TopicQueueMappingOne::new(
354                                    mapping_detail.clone(),
355                                    top.clone().into(),
356                                    broker_name.clone().into(),
357                                    *global_id,
358                                    entry.1.clone(),
359                                ),
360                            );
361                        }
362                    }
363                }
364            }
365        }
366
367        if check_consistence {
368            if max_num as usize != global_id_map.len() {
369                return Err(static_topic_mapping_inconsistent(format!(
370                    "The total queue number in config does not match the real hosted queues {} != {}",
371                    max_num,
372                    global_id_map.len()
373                )));
374            }
375            for i in 0..max_num {
376                if !global_id_map.contains_key(&i) {
377                    return Err(static_topic_mapping_inconsistent(format!(
378                        "The queue number {} is not in globalIdMap",
379                        i
380                    )));
381                }
382            }
383        }
384        let values = global_id_map.values().cloned().collect();
385        TopicQueueMappingUtils::check_if_reuse_physical_queue(&values)?;
386        Ok(global_id_map)
387    }
388    pub fn write_to_temp(wrapper: &TopicRemappingDetailWrapper, after: bool) -> Result<CheetahString> {
389        let topic = wrapper.topic();
390        let data = wrapper.serialize_json()?;
391        let mut suffix = topic_remapping_detail_wrapper::SUFFIX_BEFORE;
392        if after {
393            suffix = topic_remapping_detail_wrapper::SUFFIX_AFTER;
394        }
395        if let Ok(tmpdir) = std::env::var("java.io.tmpdir") {
396            let file_name = format!("{}/{}-{}{}", tmpdir, topic, wrapper.get_epoch(), suffix);
397            std::fs::write(&file_name, data)
398                .map_err(|source| crate::error::io_source("write_static_topic_mapping", source))?;
399            return Ok(file_name.into());
400        }
401
402        Err(crate::error::missing_configuration("java.io.tmpdir"))
403    }
404    pub fn check_target_brokers_complete(
405        target_brokers: &HashSet<CheetahString>,
406        broker_config_map: &HashMap<CheetahString, TopicConfigAndQueueMapping>,
407    ) -> Result<()> {
408        for broker in broker_config_map.keys() {
409            if let Some(mapping) = broker_config_map.get(broker) {
410                if let Some(detail) = mapping.get_topic_queue_mapping_detail() {
411                    if let Some(queues) = &detail.hosted_queues {
412                        if queues.is_empty() {
413                            continue;
414                        }
415                    }
416                }
417            }
418            if !target_brokers.contains(broker) {
419                return Err(static_topic_mapping_inconsistent(format!(
420                    "The existed broker {} does not in target brokers ",
421                    broker
422                )));
423            }
424        }
425
426        Ok(())
427    }
428    pub fn check_physical_queue_consistence(
429        broker_config_map: &HashMap<CheetahString, TopicConfigAndQueueMapping>,
430    ) -> Result<()> {
431        for entry in broker_config_map {
432            let config_mapping = entry.1;
433            if let Some(detail) = config_mapping.get_topic_queue_mapping_detail() {
434                if config_mapping.topic_config.read_queue_nums < config_mapping.topic_config.write_queue_nums {
435                    return Err(static_topic_mapping_inconsistent(
436                        "Read queues is smaller than write queues",
437                    ));
438                }
439                if let Some(queues) = &detail.hosted_queues {
440                    for items in queues.values() {
441                        for item in items {
442                            if item.start_offset != 0 {
443                                return Err(static_topic_mapping_inconsistent(
444                                    "The start offset does not begin from 0",
445                                ));
446                            }
447                            if let Some(bname) = &item.bname {
448                                let topic_config = broker_config_map.get(bname.as_str());
449                                if topic_config.is_none() {
450                                    return Err(static_topic_mapping_inconsistent("The broker of item does not exist"));
451                                } else if let Some(topic_config) = topic_config {
452                                    if item.queue_id >= topic_config.topic_config.write_queue_nums as i32 {
453                                        return Err(static_topic_mapping_inconsistent(
454                                            "The physical queue id is overflow the write queues",
455                                        ));
456                                    }
457                                }
458                            }
459                        }
460                    }
461                }
462            }
463        }
464        Ok(())
465    }
466
467    pub fn create_topic_config_mapping(
468        topic: &str,
469        queue_num: i32,
470        target_brokers: &HashSet<CheetahString>,
471        broker_config_map: &mut HashMap<CheetahString, TopicConfigAndQueueMapping>,
472    ) -> Result<TopicRemappingDetailWrapper> {
473        TopicQueueMappingUtils::check_target_brokers_complete(target_brokers, broker_config_map)?;
474        let mut global_id_map = HashMap::new();
475        let mut max_epoch_and_num = (current_millis(), queue_num);
476        if !broker_config_map.is_empty() {
477            let new_max_epoch_and_num = TopicQueueMappingUtils::check_name_epoch_num_consistence(
478                &CheetahString::from(topic),
479                broker_config_map,
480            )?;
481            max_epoch_and_num.0 = new_max_epoch_and_num.0 as u64;
482            max_epoch_and_num.1 = new_max_epoch_and_num.1;
483            let mut detail_list = vec![];
484            detail_list.extend(TopicQueueMappingUtils::get_mapping_detail_from_config(
485                broker_config_map.values().cloned().collect(),
486            )?);
487            global_id_map = TopicQueueMappingUtils::check_and_build_mapping_items(detail_list, false, true)?;
488            TopicQueueMappingUtils::check_if_reuse_physical_queue(&global_id_map.values().cloned().collect())?;
489            TopicQueueMappingUtils::check_physical_queue_consistence(broker_config_map)?;
490        }
491        if (queue_num as usize) < global_id_map.len() {
492            return Err(static_topic_mapping_inconsistent(format!(
493                "Cannot decrease the queue num for static topic {} < {}",
494                queue_num,
495                global_id_map.len()
496            )));
497        }
498        //check the queue number
499        if (queue_num as usize) == global_id_map.len() {
500            return Err(static_topic_mapping_inconsistent(
501                "The topic queue num is equal the existed queue num, do nothing",
502            ));
503        }
504
505        //the check is ok, now do the mapping allocation
506        let mut broker_num_map = HashMap::new();
507        for broker in target_brokers {
508            broker_num_map.insert(broker.clone(), 0);
509        }
510        let mut old_id_to_broker = HashMap::new();
511        for entry in &global_id_map {
512            let leader_broker = entry.1.bname();
513            old_id_to_broker.insert(*entry.0, CheetahString::from(leader_broker));
514            if !broker_num_map.contains_key(leader_broker) {
515                broker_num_map.insert(leader_broker.into(), 1);
516            } else {
517                broker_num_map.insert(leader_broker.into(), broker_num_map[leader_broker] + 1);
518            }
519        }
520        let mut allocator = MappingAllocator::new(old_id_to_broker, broker_num_map, HashMap::new());
521        allocator.up_to_num(queue_num);
522        let new_id_to_broker = allocator.id_to_broker();
523
524        //construct the topic configAndMapping
525        let new_epoch = (max_epoch_and_num.0 + 1000).max(current_millis());
526        for e in new_id_to_broker {
527            let queue_id = e.0;
528            let broker = e.1;
529            if global_id_map.contains_key(queue_id) {
530                //ignore the exited
531                continue;
532            }
533            let mut config_mapping;
534            if !broker_config_map.contains_key(broker) {
535                config_mapping = TopicConfigAndQueueMapping::new(
536                    TopicConfig::new(topic),
537                    Some(TopicQueueMappingDetail {
538                        topic_queue_mapping_info: TopicQueueMappingInfo::new(
539                            topic.into(),
540                            0,
541                            broker.into(),
542                            current_millis() as i64,
543                        ),
544                        hosted_queues: None,
545                    }),
546                );
547                config_mapping.topic_config.write_queue_nums = 1;
548                config_mapping.topic_config.read_queue_nums = 1;
549                broker_config_map.insert(broker.clone(), config_mapping.clone());
550            } else {
551                config_mapping = broker_config_map[broker].clone();
552                config_mapping.topic_config.write_queue_nums += 1;
553                config_mapping.topic_config.read_queue_nums += 1;
554            }
555            let mapping_item = LogicQueueMappingItem {
556                gen: 0,
557                queue_id: config_mapping.topic_config.write_queue_nums as i32,
558                bname: Some(broker.clone()),
559                logic_offset: 0,
560                start_offset: 0,
561                end_offset: -1,
562                time_of_start: -1,
563                time_of_end: -1,
564            };
565            if let Some(detail) = config_mapping.topic_queue_mapping_detail.as_mut() {
566                TopicQueueMappingDetail::put_mapping_info(detail, *queue_id, vec![mapping_item]);
567            }
568            broker_config_map.insert(broker.clone(), config_mapping);
569        }
570
571        // set the topic config
572        for entry in &mut *broker_config_map {
573            let config_mapping = entry.1;
574            if let Some(detail) = &mut config_mapping.topic_queue_mapping_detail {
575                detail.topic_queue_mapping_info.epoch = new_epoch as i64;
576                detail.topic_queue_mapping_info.total_queues = queue_num;
577            }
578        }
579        //double check the config
580
581        TopicQueueMappingUtils::check_name_epoch_num_consistence(&CheetahString::from(topic), broker_config_map)?;
582        global_id_map = TopicQueueMappingUtils::check_and_build_mapping_items(
583            TopicQueueMappingUtils::get_mapping_detail_from_config(broker_config_map.values().cloned().collect())?,
584            false,
585            true,
586        )?;
587        TopicQueueMappingUtils::check_if_reuse_physical_queue(&global_id_map.values().cloned().collect())?;
588        TopicQueueMappingUtils::check_physical_queue_consistence(broker_config_map)?;
589        let map = broker_config_map.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
590        Ok(TopicRemappingDetailWrapper::new(
591            CheetahString::from_slice(topic),
592            topic_remapping_detail_wrapper::TYPE_CREATE_OR_UPDATE.to_string().into(),
593            new_epoch,
594            map,
595            HashSet::new(),
596            HashSet::new(),
597        ))
598    }
599}
600pub struct MappingAllocator {
601    broker_num_map: HashMap<CheetahString, i32>,
602    id_to_broker: HashMap<i32, CheetahString>,
603    //used for remapping
604    broker_num_map_before_remapping: HashMap<CheetahString, i32>,
605    current_index: i32,
606    least_brokers: Vec<CheetahString>,
607}
608impl MappingAllocator {
609    pub fn new(
610        id_to_broker: HashMap<i32, CheetahString>,
611        broker_num_map: HashMap<CheetahString, i32>,
612        broker_num_map_before_remapping: HashMap<CheetahString, i32>,
613    ) -> Self {
614        Self {
615            id_to_broker,
616            broker_num_map,
617            broker_num_map_before_remapping,
618            current_index: 0,
619            least_brokers: vec![],
620        }
621    }
622
623    fn fresh_state(&mut self) {
624        let mut min_num = i32::MAX;
625        for entry in &self.broker_num_map {
626            if *entry.1 < min_num {
627                self.least_brokers.clear();
628                self.least_brokers.push(entry.0.clone());
629                min_num = *entry.1;
630            } else if *entry.1 == min_num {
631                self.least_brokers.push(entry.0.clone());
632            }
633        }
634        //reduce the remapping
635        if !self.broker_num_map_before_remapping.is_empty() {
636            self.least_brokers.sort_by(|o1, o2| {
637                let mut i1 = 0;
638                let mut i2 = 0;
639                if self.broker_num_map_before_remapping.contains_key(o1) {
640                    if let Some(s) = self.broker_num_map_before_remapping.get(o1) {
641                        i1 = *s;
642                    }
643                }
644                if self.broker_num_map_before_remapping.contains_key(o2) {
645                    if let Some(s) = self.broker_num_map_before_remapping.get(o2) {
646                        i2 = *s;
647                    }
648                }
649                if i1 - i2 > 0 {
650                    return std::cmp::Ordering::Greater;
651                } else if i1 - i2 < 0 {
652                    return Ordering::Less;
653                }
654                Ordering::Equal
655            });
656        } else {
657            //reduce the imbalance
658            let mut rng = rand::rng();
659            self.least_brokers.shuffle(&mut rng);
660        }
661        self.current_index = (self.least_brokers.len() - 1) as i32;
662    }
663    fn next_broker(&mut self) -> CheetahString {
664        if self.least_brokers.is_empty() {
665            self.fresh_state();
666        }
667        let tmp_index = self.current_index as usize % self.least_brokers.len();
668        self.least_brokers.remove(tmp_index)
669    }
670
671    pub fn broker_num_map(&self) -> &HashMap<CheetahString, i32> {
672        &self.broker_num_map
673    }
674
675    pub fn up_to_num(&mut self, max_queue_num: i32) {
676        let curr_size = self.id_to_broker.len();
677        if (max_queue_num as usize) <= curr_size {
678            return;
679        }
680        for i in curr_size..(max_queue_num as usize) {
681            let next_broker = self.next_broker();
682            if self.broker_num_map.contains_key(&next_broker) {
683                self.broker_num_map
684                    .insert(next_broker.clone(), self.broker_num_map[&next_broker] + 1);
685            } else {
686                self.broker_num_map.insert(next_broker.clone(), 1);
687            }
688            self.id_to_broker.insert(i as i32, next_broker);
689        }
690    }
691
692    pub fn id_to_broker(&self) -> &HashMap<i32, CheetahString> {
693        &self.id_to_broker
694    }
695}
696
697impl TopicQueueMappingUtils {
698    pub fn check_leader_in_target_brokers(
699        mapping_ones: &[TopicQueueMappingOne],
700        target_brokers: &HashSet<CheetahString>,
701    ) -> Result<()> {
702        for mapping_one in mapping_ones {
703            if !target_brokers.contains(&CheetahString::from(mapping_one.bname())) {
704                return Err(static_topic_mapping_inconsistent(
705                    "The leader broker is not in target brokers",
706                ));
707            }
708        }
709        Ok(())
710    }
711
712    pub fn block_seq_round_up(offset: i64, block_seq_size: i64) -> i64 {
713        let num = offset / block_seq_size;
714        let left = offset % block_seq_size;
715        if left < block_seq_size / 2 {
716            (num + 1) * block_seq_size
717        } else {
718            (num + 2) * block_seq_size
719        }
720    }
721
722    pub fn remapping_static_topic(
723        topic: &str,
724        broker_config_map: &mut HashMap<CheetahString, TopicConfigAndQueueMapping>,
725        target_brokers: &HashSet<CheetahString>,
726    ) -> Result<TopicRemappingDetailWrapper> {
727        let max_epoch_and_num =
728            TopicQueueMappingUtils::check_name_epoch_num_consistence(&CheetahString::from(topic), broker_config_map)?;
729
730        let mut global_id_map = TopicQueueMappingUtils::check_and_build_mapping_items(
731            TopicQueueMappingUtils::get_mapping_detail_from_config(broker_config_map.values().cloned().collect())?,
732            false,
733            true,
734        )?;
735
736        TopicQueueMappingUtils::check_physical_queue_consistence(broker_config_map)?;
737        TopicQueueMappingUtils::check_if_reuse_physical_queue(&global_id_map.values().cloned().collect())?;
738
739        let max_num = max_epoch_and_num.1;
740        let mut broker_num_map: HashMap<CheetahString, i32> = HashMap::new();
741        for broker in target_brokers {
742            broker_num_map.insert(broker.clone(), 0);
743        }
744
745        let mut broker_num_map_before_remapping: HashMap<CheetahString, i32> = HashMap::new();
746        for mapping_one in global_id_map.values() {
747            let bname = CheetahString::from(mapping_one.bname());
748            if broker_num_map_before_remapping.contains_key(&bname) {
749                broker_num_map_before_remapping.insert(bname.clone(), broker_num_map_before_remapping[&bname] + 1);
750            } else {
751                broker_num_map_before_remapping.insert(bname, 1);
752            }
753        }
754
755        let mut allocator =
756            MappingAllocator::new(HashMap::new(), broker_num_map.clone(), broker_num_map_before_remapping);
757        allocator.up_to_num(max_num);
758        let expected_broker_num_map = allocator.broker_num_map().clone();
759
760        let mut wait_assign_queues = std::collections::VecDeque::new();
761        let mut expected_id_to_broker: HashMap<i32, CheetahString> = HashMap::new();
762        let mut expected_broker_num_map_mut = expected_broker_num_map;
763        for (queue_id, mapping_one) in &global_id_map {
764            let leader_broker = CheetahString::from(mapping_one.bname());
765            if expected_broker_num_map_mut.contains_key(&leader_broker) {
766                if expected_broker_num_map_mut[&leader_broker] > 0 {
767                    expected_id_to_broker.insert(*queue_id, leader_broker.clone());
768                    expected_broker_num_map_mut
769                        .insert(leader_broker.clone(), expected_broker_num_map_mut[&leader_broker] - 1);
770                } else {
771                    wait_assign_queues.push_back(*queue_id);
772                    expected_broker_num_map_mut.remove(&leader_broker);
773                }
774            } else {
775                wait_assign_queues.push_back(*queue_id);
776            }
777        }
778
779        for (broker, queue_num) in &expected_broker_num_map_mut {
780            for _ in 0..*queue_num {
781                if let Some(queue_id) = wait_assign_queues.pop_front() {
782                    expected_id_to_broker.insert(queue_id, broker.clone());
783                }
784            }
785        }
786
787        let new_epoch = ((max_epoch_and_num.0 as u64 + 1000).max(current_millis())) as i64;
788
789        // construct the remapping info
790        let mut brokers_to_map_out: HashSet<CheetahString> = HashSet::new();
791        let mut brokers_to_map_in: HashSet<CheetahString> = HashSet::new();
792
793        for (queue_id, broker) in &expected_id_to_broker {
794            let Some(topic_queue_mapping_one) = global_id_map.get(queue_id) else {
795                continue;
796            };
797            if topic_queue_mapping_one.bname() == broker.as_str() {
798                continue;
799            }
800
801            // remapping
802            let map_in_broker = broker.clone();
803            let map_out_broker = CheetahString::from(topic_queue_mapping_one.bname());
804            brokers_to_map_in.insert(map_in_broker.clone());
805            brokers_to_map_out.insert(map_out_broker.clone());
806
807            let map_in_config = broker_config_map.entry(map_in_broker.clone()).or_insert_with(|| {
808                TopicConfigAndQueueMapping::new(
809                    TopicConfig::with_queues(topic, 0, 0),
810                    Some(TopicQueueMappingDetail {
811                        topic_queue_mapping_info: TopicQueueMappingInfo::new(
812                            topic.into(),
813                            max_num,
814                            map_in_broker.clone(),
815                            new_epoch,
816                        ),
817                        hosted_queues: Some(HashMap::new()),
818                    }),
819                )
820            });
821            map_in_config.topic_config.write_queue_nums += 1;
822            map_in_config.topic_config.read_queue_nums += 1;
823
824            let mut items: Vec<LogicQueueMappingItem> = topic_queue_mapping_one.items().clone();
825            let last = items.last().cloned();
826            if let Some(last) = last {
827                items.push(LogicQueueMappingItem {
828                    gen: last.gen + 1,
829                    queue_id: map_in_config.topic_config.write_queue_nums as i32 - 1,
830                    bname: Some(map_in_broker.clone()),
831                    logic_offset: -1,
832                    start_offset: 0,
833                    end_offset: -1,
834                    time_of_start: -1,
835                    time_of_end: -1,
836                });
837            }
838
839            // use the same object
840            if let Some(detail) = &mut map_in_config.topic_queue_mapping_detail {
841                if detail.hosted_queues.is_none() {
842                    detail.hosted_queues = Some(HashMap::new());
843                }
844                if let Some(queues) = &mut detail.hosted_queues {
845                    queues.insert(*queue_id, items.clone());
846                }
847            }
848
849            let map_out_config = broker_config_map.get_mut(&map_out_broker);
850            if let Some(map_out_config) = map_out_config {
851                if let Some(detail) = &mut map_out_config.topic_queue_mapping_detail {
852                    if detail.hosted_queues.is_none() {
853                        detail.hosted_queues = Some(HashMap::new());
854                    }
855                    if let Some(queues) = &mut detail.hosted_queues {
856                        queues.insert(*queue_id, items);
857                    }
858                }
859            }
860        }
861
862        for config_mapping in broker_config_map.values_mut() {
863            if let Some(detail) = &mut config_mapping.topic_queue_mapping_detail {
864                detail.topic_queue_mapping_info.epoch = new_epoch;
865                detail.topic_queue_mapping_info.total_queues = max_num;
866            }
867        }
868
869        // double check
870        TopicQueueMappingUtils::check_name_epoch_num_consistence(&CheetahString::from(topic), broker_config_map)?;
871        global_id_map = TopicQueueMappingUtils::check_and_build_mapping_items(
872            TopicQueueMappingUtils::get_mapping_detail_from_config(broker_config_map.values().cloned().collect())?,
873            false,
874            true,
875        )?;
876        TopicQueueMappingUtils::check_physical_queue_consistence(broker_config_map)?;
877        TopicQueueMappingUtils::check_if_reuse_physical_queue(&global_id_map.values().cloned().collect())?;
878        TopicQueueMappingUtils::check_leader_in_target_brokers(
879            &global_id_map.values().cloned().collect::<Vec<_>>(),
880            target_brokers,
881        )?;
882
883        let map = broker_config_map.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
884
885        Ok(TopicRemappingDetailWrapper::new(
886            CheetahString::from_slice(topic),
887            topic_remapping_detail_wrapper::TYPE_REMAPPING.to_string().into(),
888            new_epoch as u64,
889            map,
890            brokers_to_map_in,
891            brokers_to_map_out,
892        ))
893    }
894}
895
896#[cfg(test)]
897mod tests {
898    use std::sync::Mutex;
899
900    use rocketmq_error::fields;
901    use rocketmq_error::ViewValueRef;
902
903    use super::*;
904
905    /// Serializes every environment mutation in this module so concurrent
906    /// `cargo test` threads cannot race on the same process-wide env block.
907    static ENV_LOCK: Mutex<()> = Mutex::new(());
908
909    #[test]
910    fn check_logic_queue_mapping_item_offset_accepts_monotonic_items() {
911        let items = vec![
912            LogicQueueMappingItem {
913                gen: 0,
914                logic_offset: 0,
915                start_offset: 0,
916                end_offset: 9,
917                ..Default::default()
918            },
919            LogicQueueMappingItem {
920                gen: 1,
921                logic_offset: 10,
922                start_offset: 0,
923                end_offset: -1,
924                ..Default::default()
925            },
926        ];
927
928        assert!(TopicQueueMappingUtils::check_logic_queue_mapping_item_offset(&items).is_ok());
929    }
930
931    #[test]
932    fn check_logic_queue_mapping_item_offset_rejects_negative_fields_as_illegal_argument() {
933        let items = vec![LogicQueueMappingItem {
934            gen: -1,
935            logic_offset: 0,
936            start_offset: 0,
937            end_offset: -1,
938            queue_id: 0,
939            ..Default::default()
940        }];
941
942        let error = TopicQueueMappingUtils::check_logic_queue_mapping_item_offset(&items)
943            .expect_err("negative gen should be rejected");
944
945        assert_eq!(error.descriptor(), &rocketmq_error::CORE_ARGUMENT_INVALID);
946    }
947
948    #[test]
949    fn check_name_epoch_num_consistence_rejects_empty_config_as_route_inconsistent() {
950        let broker_config_map = HashMap::<CheetahString, TopicConfigAndQueueMapping>::new();
951
952        let error = TopicQueueMappingUtils::check_name_epoch_num_consistence(
953            &CheetahString::from_static_str("TopicA"),
954            &broker_config_map,
955        )
956        .expect_err("empty static topic config map should be rejected");
957
958        assert_eq!(error.descriptor(), &rocketmq_error::ROUTE_TOPIC_INCONSISTENT);
959    }
960
961    #[test]
962    fn check_if_reuse_physical_queue_reports_duplicate_physical_queue_as_route_inconsistent() {
963        let broker = CheetahString::from_static_str("broker-a");
964        let topic = CheetahString::from_static_str("TopicA");
965        let mapping_detail = TopicQueueMappingDetail {
966            topic_queue_mapping_info: TopicQueueMappingInfo::new(topic.clone(), 2, broker.clone(), 1),
967            hosted_queues: None,
968        };
969        let shared_physical_queue = LogicQueueMappingItem {
970            bname: Some(broker.clone()),
971            queue_id: 0,
972            ..Default::default()
973        };
974        let mapping_ones = vec![
975            TopicQueueMappingOne::new(
976                mapping_detail.clone(),
977                topic.to_string(),
978                broker.to_string(),
979                0,
980                vec![shared_physical_queue.clone()],
981            ),
982            TopicQueueMappingOne::new(
983                mapping_detail,
984                topic.to_string(),
985                broker.to_string(),
986                1,
987                vec![shared_physical_queue],
988            ),
989        ];
990
991        let error = TopicQueueMappingUtils::check_if_reuse_physical_queue(&mapping_ones)
992            .expect_err("duplicate physical queue should be rejected");
993
994        assert_eq!(error.descriptor(), &rocketmq_error::ROUTE_TOPIC_INCONSISTENT);
995    }
996
997    #[test]
998    fn write_to_temp_reports_missing_temp_dir_as_config_missing() {
999        struct EnvVarGuard {
1000            key: &'static str,
1001            original: Option<std::ffi::OsString>,
1002        }
1003
1004        impl EnvVarGuard {
1005            fn remove(key: &'static str) -> Self {
1006                let original = std::env::var_os(key);
1007                // SAFETY: ENV_LOCK serializes every environment mutation in
1008                // this module, and the guard restores the original value
1009                // before the lock is released.
1010                unsafe {
1011                    std::env::remove_var(key);
1012                }
1013                Self { key, original }
1014            }
1015        }
1016
1017        impl Drop for EnvVarGuard {
1018            fn drop(&mut self) {
1019                // SAFETY: ENV_LOCK is held by the `_lock` guard in the
1020                // enclosing test for the whole time this guard is alive.
1021                unsafe {
1022                    match &self.original {
1023                        Some(value) => std::env::set_var(self.key, value),
1024                        None => std::env::remove_var(self.key),
1025                    }
1026                }
1027            }
1028        }
1029
1030        let _lock = ENV_LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
1031        let _guard = EnvVarGuard::remove("java.io.tmpdir");
1032
1033        let error = TopicQueueMappingUtils::write_to_temp(&TopicRemappingDetailWrapper::empty(), false)
1034            .expect_err("missing temp directory property should be rejected");
1035
1036        assert_eq!(error.descriptor(), &rocketmq_error::CORE_CONFIGURATION_MISSING);
1037    }
1038
1039    #[test]
1040    fn get_leader_broker_reports_missing_broker_as_route_inconsistent() {
1041        let items = vec![LogicQueueMappingItem {
1042            gen: 0,
1043            logic_offset: 0,
1044            start_offset: 0,
1045            end_offset: -1,
1046            queue_id: 0,
1047            bname: None,
1048            ..Default::default()
1049        }];
1050
1051        let error = TopicQueueMappingUtils::get_leader_broker(&items)
1052            .expect_err("leader item without broker should be rejected");
1053
1054        assert_eq!(error.descriptor(), &rocketmq_error::ROUTE_TOPIC_INCONSISTENT);
1055    }
1056
1057    #[test]
1058    fn check_if_leader_requires_last_item_broker_to_match_detail_broker() {
1059        let broker_a = CheetahString::from_static_str("broker-a");
1060        let broker_b = CheetahString::from_static_str("broker-b");
1061        let detail = TopicQueueMappingDetail {
1062            topic_queue_mapping_info: TopicQueueMappingInfo {
1063                bname: Some(broker_a.clone()),
1064                ..Default::default()
1065            },
1066            hosted_queues: None,
1067        };
1068        let leader_items = vec![LogicQueueMappingItem {
1069            bname: Some(broker_a),
1070            ..Default::default()
1071        }];
1072        let follower_items = vec![LogicQueueMappingItem {
1073            bname: Some(broker_b),
1074            ..Default::default()
1075        }];
1076
1077        assert!(TopicQueueMappingUtils::check_if_leader(&leader_items, &detail));
1078        assert!(!TopicQueueMappingUtils::check_if_leader(&follower_items, &detail));
1079        assert!(!TopicQueueMappingUtils::check_if_leader(&[], &detail));
1080    }
1081
1082    #[test]
1083    fn find_next_returns_none_without_current_item() {
1084        let items = vec![LogicQueueMappingItem {
1085            gen: 1,
1086            ..Default::default()
1087        }];
1088
1089        assert!(TopicQueueMappingUtils::find_next(&items, None, true).is_none());
1090    }
1091
1092    fn item(gen: i32, logic_offset: i64) -> LogicQueueMappingItem {
1093        LogicQueueMappingItem {
1094            gen,
1095            logic_offset,
1096            ..Default::default()
1097        }
1098    }
1099
1100    #[test]
1101    fn find_next_returns_following_item_and_none_after_last() {
1102        let items = vec![item(0, 0), item(1, 10), item(2, -1)];
1103
1104        assert_eq!(
1105            TopicQueueMappingUtils::find_next(&items, Some(&items[0]), true),
1106            Some(&items[1])
1107        );
1108        assert!(TopicQueueMappingUtils::find_next(&items, Some(&items[2]), true).is_none());
1109        assert!(TopicQueueMappingUtils::find_next(&items, Some(&items[1]), true).is_none());
1110        assert_eq!(
1111            TopicQueueMappingUtils::find_next(&items, Some(&items[1]), false),
1112            Some(&items[2])
1113        );
1114    }
1115
1116    #[test]
1117    fn find_logic_queue_mapping_item_scans_from_end_for_covering_item() {
1118        let items = vec![item(0, 0), item(1, 100), item(2, -1)];
1119
1120        assert_eq!(
1121            TopicQueueMappingUtils::find_logic_queue_mapping_item(&items, 150, true),
1122            Some(&items[1])
1123        );
1124        assert_eq!(
1125            TopicQueueMappingUtils::find_logic_queue_mapping_item(&items, 100, true),
1126            Some(&items[1])
1127        );
1128        assert_eq!(
1129            TopicQueueMappingUtils::find_logic_queue_mapping_item(&items, 99, true),
1130            Some(&items[0])
1131        );
1132        assert_eq!(
1133            TopicQueueMappingUtils::find_logic_queue_mapping_item(&items, 150, false),
1134            Some(&items[2])
1135        );
1136    }
1137
1138    #[test]
1139    fn find_logic_queue_mapping_item_falls_back_to_first_non_negative_item() {
1140        let items = vec![item(0, -1), item(1, 10), item(2, 20)];
1141
1142        assert_eq!(
1143            TopicQueueMappingUtils::find_logic_queue_mapping_item(&items, 5, true),
1144            Some(&items[1])
1145        );
1146        assert!(TopicQueueMappingUtils::find_logic_queue_mapping_item(&items[..1], 5, true).is_none());
1147    }
1148
1149    #[test]
1150    fn find_logic_queue_mapping_item_returns_none_for_empty_items() {
1151        assert!(TopicQueueMappingUtils::find_logic_queue_mapping_item(&[], 0, true).is_none());
1152        assert!(TopicQueueMappingUtils::find_logic_queue_mapping_item(&[], 0, false).is_none());
1153    }
1154
1155    #[test]
1156    fn get_leader_item_returns_last_item_and_rejects_empty_items() {
1157        let items = vec![item(0, 0), item(1, 10)];
1158
1159        let leader = TopicQueueMappingUtils::get_leader_item(&items).expect("non-empty items have a leader");
1160        assert_eq!(leader, items[1]);
1161
1162        let error = TopicQueueMappingUtils::get_leader_item(&[]).expect_err("empty items should be rejected");
1163        assert_eq!(error.descriptor(), &rocketmq_error::ROUTE_TOPIC_INCONSISTENT);
1164    }
1165
1166    #[test]
1167    fn check_leader_in_target_brokers_requires_every_leader_broker_in_targets() {
1168        let mapping_one = |bname: &str| {
1169            TopicQueueMappingOne::new(
1170                TopicQueueMappingDetail {
1171                    topic_queue_mapping_info: TopicQueueMappingInfo::default(),
1172                    hosted_queues: None,
1173                },
1174                "TopicA".to_string(),
1175                bname.to_string(),
1176                0,
1177                vec![],
1178            )
1179        };
1180        let mapping_ones = vec![mapping_one("broker-a"), mapping_one("broker-b")];
1181        let broker_a = CheetahString::from_static_str("broker-a");
1182        let broker_b = CheetahString::from_static_str("broker-b");
1183
1184        let complete = HashSet::from([broker_a.clone(), broker_b]);
1185        assert!(TopicQueueMappingUtils::check_leader_in_target_brokers(&mapping_ones, &complete).is_ok());
1186
1187        let missing_b = HashSet::from([broker_a]);
1188        let error = TopicQueueMappingUtils::check_leader_in_target_brokers(&mapping_ones, &missing_b)
1189            .expect_err("leader broker outside target brokers should be rejected");
1190        assert_eq!(error.descriptor(), &rocketmq_error::ROUTE_TOPIC_INCONSISTENT);
1191        // Canonical rendering discards the free-form reason ("The leader broker is not in target
1192        // brokers"): the rejected leader broker only surfaces through the mapping scope the helper
1193        // attaches to the inconsistency.
1194        assert_eq!(
1195            error
1196                .diagnostic_view()
1197                .expect("valid diagnostic error context")
1198                .fields()
1199                .find(|field| field.name() == fields::TOPIC.schema().name())
1200                .map(|field| field.value()),
1201            Some(ViewValueRef::Text("static_topic_mapping"))
1202        );
1203    }
1204
1205    #[test]
1206    fn get_mock_broker_name_prefixes_scope() {
1207        let prefix = mix_all::LOGICAL_QUEUE_MOCK_BROKER_PREFIX;
1208
1209        assert_eq!(
1210            TopicQueueMappingUtils::get_mock_broker_name("scope-a").as_str(),
1211            format!("{prefix}scope-a")
1212        );
1213        assert_eq!(
1214            TopicQueueMappingUtils::get_mock_broker_name(mix_all::METADATA_SCOPE_GLOBAL).as_str(),
1215            format!("{prefix}{}", &mix_all::METADATA_SCOPE_GLOBAL[2..])
1216        );
1217    }
1218
1219    #[test]
1220    #[should_panic(expected = "Scope cannot be null")]
1221    fn get_mock_broker_name_panics_on_empty_scope() {
1222        TopicQueueMappingUtils::get_mock_broker_name("");
1223    }
1224
1225    #[test]
1226    fn block_seq_round_up_rounds_by_half_block() {
1227        for (offset, expected) in [(0, 8), (1, 8), (4, 16), (8, 16), (11, 16), (12, 24)] {
1228            assert_eq!(
1229                TopicQueueMappingUtils::block_seq_round_up(offset, 8),
1230                expected,
1231                "offset {offset}"
1232            );
1233        }
1234    }
1235}
rocketmq_protocol/protocol/static_topic/topic_queue_mapping_utils.rs

rocketmq_protocol/protocol/static_topic/
topic_queue_mapping_utils.rs
