rocketmq_model/
allocation.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::collections::BTreeMap;
16use std::collections::HashMap;
17use std::collections::HashSet;
18use std::fmt;
19use std::sync::Arc;
20use std::sync::Mutex;
21use std::sync::OnceLock;
22
23use cheetah_string::CheetahString;
24use rocketmq_error::Result;
25use tracing::info;
26
27use crate::consistent_hash::ConsistentHashRouter;
28use crate::consistent_hash::HashFunction;
29use crate::consistent_hash::Node;
30use crate::message::MessageQueue;
31
32/// Runtime-neutral queue allocation contract.
33pub trait AllocateMessageQueueStrategy: Send + Sync {
34    fn allocate(
35        &self,
36        consumer_group: &CheetahString,
37        current_cid: &CheetahString,
38        mq_all: &[MessageQueue],
39        cid_all: &[CheetahString],
40    ) -> Result<Vec<MessageQueue>>;
41
42    fn get_name(&self) -> &'static str;
43}
44
45pub type AbstractAllocateMessageQueueStrategy = dyn AllocateMessageQueueStrategy;
46
47/// Validates the shared preconditions used by allocation strategies.
48pub fn check(
49    consumer_group: &CheetahString,
50    current_cid: &CheetahString,
51    mq_all: &[MessageQueue],
52    cid_all: &[CheetahString],
53) -> Result<bool> {
54    if current_cid.is_empty() {
55        return Err(crate::error::invalid_argument("currentCID is empty"));
56    }
57    if mq_all.is_empty() {
58        return Err(crate::error::invalid_argument("mqAll is null or mqAll empty"));
59    }
60    if cid_all.is_empty() {
61        return Err(crate::error::invalid_argument("cidAll is null or cidAll empty"));
62    }
63    if !cid_all.iter().any(|cid| cid == current_cid) {
64        info!(
65            consumer_group = %consumer_group,
66            current_cid = %current_cid,
67            "consumer id is not present in the allocation set"
68        );
69        return Ok(false);
70    }
71    Ok(true)
72}
73
74/// Allocates contiguous queue ranges as evenly as possible.
75pub struct AllocateMessageQueueAveragely;
76
77impl AllocateMessageQueueStrategy for AllocateMessageQueueAveragely {
78    fn allocate(
79        &self,
80        consumer_group: &CheetahString,
81        current_cid: &CheetahString,
82        mq_all: &[MessageQueue],
83        cid_all: &[CheetahString],
84    ) -> Result<Vec<MessageQueue>> {
85        if !check(consumer_group, current_cid, mq_all, cid_all)? {
86            return Ok(Vec::new());
87        }
88        let index = cid_all.iter().position(|cid| cid == current_cid).unwrap_or(0);
89        let remainder = mq_all.len() % cid_all.len();
90        let average = if mq_all.len() <= cid_all.len() {
91            1
92        } else if remainder > 0 && index < remainder {
93            mq_all.len() / cid_all.len() + 1
94        } else {
95            mq_all.len() / cid_all.len()
96        };
97        let start = if remainder > 0 && index < remainder {
98            index * average
99        } else {
100            index * average + remainder
101        };
102        Ok(mq_all.iter().skip(start).take(average).cloned().collect())
103    }
104
105    fn get_name(&self) -> &'static str {
106        "AVG"
107    }
108}
109
110/// Allocates every Nth queue to each consumer.
111pub struct AllocateMessageQueueAveragelyByCircle;
112
113impl AllocateMessageQueueStrategy for AllocateMessageQueueAveragelyByCircle {
114    fn allocate(
115        &self,
116        consumer_group: &CheetahString,
117        current_cid: &CheetahString,
118        mq_all: &[MessageQueue],
119        cid_all: &[CheetahString],
120    ) -> Result<Vec<MessageQueue>> {
121        if !check(consumer_group, current_cid, mq_all, cid_all)? {
122            return Ok(Vec::new());
123        }
124        let index = cid_all.iter().position(|cid| cid == current_cid).unwrap_or(0);
125        Ok(mq_all
126            .iter()
127            .enumerate()
128            .filter(|(queue_index, _)| queue_index % cid_all.len() == index)
129            .map(|(_, queue)| queue.clone())
130            .collect())
131    }
132
133    fn get_name(&self) -> &'static str {
134        "AVG_BY_CIRCLE"
135    }
136}
137
138/// Returns a statically configured set of queues.
139pub struct AllocateMessageQueueByConfig {
140    message_queue_list: Vec<MessageQueue>,
141}
142
143impl AllocateMessageQueueByConfig {
144    pub fn new(message_queue_list: Vec<MessageQueue>) -> Self {
145        Self { message_queue_list }
146    }
147}
148
149impl AllocateMessageQueueStrategy for AllocateMessageQueueByConfig {
150    fn allocate(
151        &self,
152        _consumer_group: &CheetahString,
153        _current_cid: &CheetahString,
154        _mq_all: &[MessageQueue],
155        _cid_all: &[CheetahString],
156    ) -> Result<Vec<MessageQueue>> {
157        Ok(self.message_queue_list.clone())
158    }
159
160    fn get_name(&self) -> &'static str {
161        "CONFIG"
162    }
163}
164
165/// Limits allocation to selected broker machine rooms.
166pub struct AllocateMessageQueueByMachineRoom {
167    consumer_idcs: HashSet<CheetahString>,
168}
169
170impl AllocateMessageQueueByMachineRoom {
171    pub fn new(consumer_idcs: HashSet<CheetahString>) -> Self {
172        Self { consumer_idcs }
173    }
174}
175
176impl AllocateMessageQueueStrategy for AllocateMessageQueueByMachineRoom {
177    fn allocate(
178        &self,
179        consumer_group: &CheetahString,
180        current_cid: &CheetahString,
181        mq_all: &[MessageQueue],
182        cid_all: &[CheetahString],
183    ) -> Result<Vec<MessageQueue>> {
184        if !check(consumer_group, current_cid, mq_all, cid_all)? {
185            return Ok(Vec::new());
186        }
187        let index = cid_all.iter().position(|cid| cid == current_cid).unwrap_or(0);
188        let queues: Vec<_> = mq_all
189            .iter()
190            .filter(|queue| {
191                queue.broker_name().split_once('@').is_some_and(|(room, broker)| {
192                    !broker.is_empty() && !broker.contains('@') && self.consumer_idcs.contains(room)
193                })
194            })
195            .cloned()
196            .collect();
197        let per_consumer = queues.len() / cid_all.len();
198        let remainder = queues.len() % cid_all.len();
199        let mut result: Vec<_> = queues
200            .iter()
201            .skip(per_consumer * index)
202            .take(per_consumer)
203            .cloned()
204            .collect();
205        if remainder > index {
206            let extra = index + per_consumer * cid_all.len();
207            if let Some(queue) = queues.get(extra) {
208                result.push(queue.clone());
209            }
210        }
211        Ok(result)
212    }
213
214    fn get_name(&self) -> &'static str {
215        "MACHINE_ROOM"
216    }
217}
218
219/// Resolves the deployment room of brokers and consumers.
220pub trait MachineRoomResolver: Send + Sync {
221    fn broker_deploy_in(&self, message_queue: &MessageQueue) -> Option<CheetahString>;
222    fn consumer_deploy_in(&self, client_id: &CheetahString) -> Option<CheetahString>;
223}
224
225/// Prefers queues local to a consumer and shares queues from rooms without consumers.
226pub struct AllocateMessageQueueByMachineRoomNearby {
227    strategy: Box<dyn AllocateMessageQueueStrategy>,
228    resolver: Box<dyn MachineRoomResolver>,
229    name: &'static str,
230}
231
232pub type AllocateMachineRoomNearby = AllocateMessageQueueByMachineRoomNearby;
233
234impl AllocateMessageQueueByMachineRoomNearby {
235    pub fn new(strategy: Box<dyn AllocateMessageQueueStrategy>, resolver: Box<dyn MachineRoomResolver>) -> Self {
236        let name = cached_strategy_name(strategy.get_name());
237        Self {
238            strategy,
239            resolver,
240            name,
241        }
242    }
243}
244
245impl AllocateMessageQueueStrategy for AllocateMessageQueueByMachineRoomNearby {
246    fn allocate(
247        &self,
248        consumer_group: &CheetahString,
249        current_cid: &CheetahString,
250        mq_all: &[MessageQueue],
251        cid_all: &[CheetahString],
252    ) -> Result<Vec<MessageQueue>> {
253        if !check(consumer_group, current_cid, mq_all, cid_all)? {
254            return Ok(Vec::new());
255        }
256        let mut room_queues: BTreeMap<CheetahString, Vec<MessageQueue>> = BTreeMap::new();
257        for queue in mq_all {
258            let room = self
259                .resolver
260                .broker_deploy_in(queue)
261                .filter(|room| !room.is_empty())
262                .ok_or_else(|| crate::error::invalid_argument(format!("Machine room is null for mq {queue}")))?;
263            room_queues.entry(room).or_default().push(queue.clone());
264        }
265        let mut room_consumers: BTreeMap<CheetahString, Vec<CheetahString>> = BTreeMap::new();
266        for cid in cid_all {
267            let room = self
268                .resolver
269                .consumer_deploy_in(cid)
270                .filter(|room| !room.is_empty())
271                .ok_or_else(|| crate::error::invalid_argument(format!("Machine room is null for consumer id {cid}")))?;
272            room_consumers.entry(room).or_default().push(cid.clone());
273        }
274        let Some(current_room) = self.resolver.consumer_deploy_in(current_cid) else {
275            return Ok(Vec::new());
276        };
277        let mut result = Vec::new();
278        if let Some(local_queues) = room_queues.remove(&current_room) {
279            if let Some(local_consumers) = room_consumers.get(&current_room) {
280                result.extend(
281                    self.strategy
282                        .allocate(consumer_group, current_cid, &local_queues, local_consumers)?,
283                );
284            }
285        }
286        for (room, queues) in room_queues {
287            if !room_consumers.contains_key(&room) {
288                result.extend(self.strategy.allocate(consumer_group, current_cid, &queues, cid_all)?);
289            }
290        }
291        Ok(result)
292    }
293
294    fn get_name(&self) -> &'static str {
295        self.name
296    }
297}
298
299fn cached_strategy_name(inner: &'static str) -> &'static str {
300    static NAMES: OnceLock<Mutex<HashMap<&'static str, &'static str>>> = OnceLock::new();
301    let mut names = NAMES
302        .get_or_init(|| Mutex::new(HashMap::new()))
303        .lock()
304        .unwrap_or_else(std::sync::PoisonError::into_inner);
305    names
306        .entry(inner)
307        .or_insert_with(|| Box::leak(format!("MACHINE_ROOM_NEARBY-{inner}").into_boxed_str()))
308}
309
310/// Assigns queues through a consistent hash ring of consumer ids.
311#[derive(Clone)]
312pub struct AllocateMessageQueueConsistentHash {
313    virtual_node_count: i32,
314    hash_function: Option<Arc<dyn HashFunction + Send + Sync>>,
315}
316
317impl fmt::Debug for AllocateMessageQueueConsistentHash {
318    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
319        f.debug_struct("AllocateMessageQueueConsistentHash")
320            .field("virtual_node_count", &self.virtual_node_count)
321            .field("custom_hash_function", &self.hash_function.is_some())
322            .finish()
323    }
324}
325
326impl Default for AllocateMessageQueueConsistentHash {
327    fn default() -> Self {
328        Self::new(10)
329    }
330}
331
332impl AllocateMessageQueueConsistentHash {
333    pub fn new(virtual_node_count: i32) -> Self {
334        Self {
335            virtual_node_count,
336            hash_function: None,
337        }
338    }
339
340    pub fn try_new(virtual_node_count: i32) -> Result<Self> {
341        if virtual_node_count < 0 {
342            return Err(crate::error::invalid_argument(format!(
343                "illegal virtualNodeCnt :{virtual_node_count}"
344            )));
345        }
346        Ok(Self::new(virtual_node_count))
347    }
348
349    pub fn with_hash_function(
350        virtual_node_count: i32,
351        hash_function: Arc<dyn HashFunction + Send + Sync>,
352    ) -> Result<Self> {
353        if virtual_node_count < 0 {
354            return Err(crate::error::invalid_argument(format!(
355                "illegal virtualNodeCnt :{virtual_node_count}"
356            )));
357        }
358        Ok(Self {
359            virtual_node_count,
360            hash_function: Some(hash_function),
361        })
362    }
363}
364
365#[derive(Clone)]
366struct ClientNode {
367    client_id: CheetahString,
368}
369
370impl Node for ClientNode {
371    fn get_key(&self) -> &str {
372        &self.client_id
373    }
374}
375
376impl AllocateMessageQueueStrategy for AllocateMessageQueueConsistentHash {
377    fn allocate(
378        &self,
379        consumer_group: &CheetahString,
380        current_cid: &CheetahString,
381        mq_all: &[MessageQueue],
382        cid_all: &[CheetahString],
383    ) -> Result<Vec<MessageQueue>> {
384        if !check(consumer_group, current_cid, mq_all, cid_all)? {
385            return Ok(Vec::new());
386        }
387        if self.virtual_node_count < 0 {
388            return Err(crate::error::invalid_argument(format!(
389                "illegal virtualNodeCnt :{}",
390                self.virtual_node_count
391            )));
392        }
393        let nodes = cid_all
394            .iter()
395            .cloned()
396            .map(|client_id| ClientNode { client_id })
397            .collect();
398        let router = if let Some(hash_function) = &self.hash_function {
399            let hash_function: Arc<dyn HashFunction> = hash_function.clone();
400            ConsistentHashRouter::try_new_with_hash_function(nodes, self.virtual_node_count, hash_function)?
401        } else {
402            ConsistentHashRouter::try_new(nodes, self.virtual_node_count)?
403        };
404        Ok(mq_all
405            .iter()
406            .filter(|queue| {
407                router
408                    .route_node_ref(&queue.to_string())
409                    .is_some_and(|node| node.client_id == *current_cid)
410            })
411            .cloned()
412            .collect())
413    }
414
415    fn get_name(&self) -> &'static str {
416        "CONSISTENT_HASH"
417    }
418}
rocketmq_model/allocation.rs

rocketmq_model/
allocation.rs
