rocketmq_client_rust/producer/message_queue_selector.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::sync::Arc;
16
17use rocketmq_model::common::message::message_queue::MessageQueue;
18use rocketmq_model::common::message::MessageTrait;
19
20/// A trait for selecting a message queue from a list of available queues.
21///
22/// This trait provides a zero-cost abstraction for queue selection logic using compile-time
23/// monomorphization. Implement this trait for your custom selector logic, or use closures
24/// directly as they automatically implement this trait.
25///
26/// # Type Parameters
27///
28/// * `M` - Message type that implements `MessageTrait`
29/// * `A` - Argument type for custom selection logic
30///
31/// # Performance
32///
33/// This trait uses compile-time generics, allowing the compiler to:
34/// - Fully inline the selector function
35/// - Eliminate dynamic dispatch overhead
36/// - Perform aggressive optimizations
37///
38/// # Example
39///
40/// ```
41/// use rocketmq_client_rust::MessageQueueSelector;
42/// use rocketmq_model::common::message::message_queue::MessageQueue;
43/// use rocketmq_model::common::message::message_single::Message;
44///
45/// // Closures automatically implement MessageQueueSelector. Order IDs are signed, so guard
46/// // against an empty queue list before indexing and use `rem_euclid` so a negative order ID
47/// // still maps into the valid queue range instead of producing a negative remainder.
48/// let selector = |mqs: &[MessageQueue], _msg: &Message, order_id: &i64| {
49/// if mqs.is_empty() {
50/// return None;
51/// }
52/// let index = order_id.rem_euclid(mqs.len() as i64) as usize;
53/// mqs.get(index).cloned()
54/// };
55///
56/// let msg = Message::builder()
57/// .topic("TopicTest")
58/// .body_slice(b"body")
59/// .build_unchecked();
60///
61/// assert_eq!(selector.select(&[], &msg, &7), None);
62///
63/// let mqs = vec![
64/// MessageQueue::from_parts("TopicTest", "broker-a", 0),
65/// MessageQueue::from_parts("TopicTest", "broker-a", 1),
66/// MessageQueue::from_parts("TopicTest", "broker-a", 2),
67/// ];
68/// assert_eq!(selector.select(&mqs, &msg, &1), Some(mqs[1].clone()));
69/// assert_eq!(selector.select(&mqs, &msg, &-1), Some(mqs[2].clone()));
70///
71/// let first = selector.select(&mqs, &msg, &i64::MIN);
72/// let second = selector.select(&mqs, &msg, &i64::MIN);
73/// assert_eq!(first, second);
74/// assert!(first.is_some_and(|mq| mqs.contains(&mq)));
75/// ```
76pub trait MessageQueueSelector<M: MessageTrait, A>: Send + Sync {
77 /// Selects a message queue from the provided list.
78 ///
79 /// # Arguments
80 ///
81 /// * `mqs` - Available message queues to select from
82 /// * `msg` - The message to be sent
83 /// * `arg` - Custom argument for selection logic
84 ///
85 /// # Returns
86 ///
87 /// Selected `MessageQueue`, or `None` if no suitable queue is found
88 fn select(&self, mqs: &[MessageQueue], msg: &M, arg: &A) -> Option<MessageQueue>;
89}
90
91/// Implement MessageQueueSelector for all compatible closures and functions.
92///
93/// This allows closures to be used directly as selectors without explicit trait implementation.
94impl<F, M, A> MessageQueueSelector<M, A> for F
95where
96 F: Fn(&[MessageQueue], &M, &A) -> Option<MessageQueue> + Send + Sync,
97 M: MessageTrait,
98{
99 fn select(&self, mqs: &[MessageQueue], msg: &M, arg: &A) -> Option<MessageQueue> {
100 self(mqs, msg, arg)
101 }
102}
103
104/// Type-erased message queue selector function for storage and cross-boundary passing.
105///
106/// This type uses dynamic dispatch (`Arc<dyn Fn>`) and is suitable for scenarios where:
107/// - The selector needs to be stored in a struct field
108/// - The selector crosses async boundaries
109/// - The selector type cannot be determined at compile time
110///
111/// # Performance Note
112///
113/// This type incurs runtime overhead due to:
114/// - Dynamic dispatch (~5-10ns per call)
115/// - Arc reference counting
116/// - Type erasure with `dyn Any`
117///
118/// For best performance, prefer using generic parameters with the `MessageQueueSelector` trait
119/// in function signatures instead of this type alias.
120///
121/// # Example
122///
123/// ```
124/// use rocketmq_client_rust::MessageQueueSelectorFn;
125/// use rocketmq_model::common::message::message_queue::MessageQueue;
126/// use rocketmq_model::common::message::message_single::Message;
127/// use std::sync::Arc;
128///
129/// let selector: MessageQueueSelectorFn = Arc::new(|mqs: &[MessageQueue], _msg, arg| {
130/// if mqs.is_empty() {
131/// return None;
132/// }
133/// let order_id = arg.downcast_ref::<i64>()?;
134/// let index = order_id.rem_euclid(mqs.len() as i64) as usize;
135/// mqs.get(index).cloned()
136/// });
137///
138/// let msg = Message::builder()
139/// .topic("TopicTest")
140/// .body_slice(b"body")
141/// .build_unchecked();
142///
143/// assert_eq!(selector(&[], &msg, &7i64), None);
144///
145/// let mqs = vec![
146/// MessageQueue::from_parts("TopicTest", "broker-a", 0),
147/// MessageQueue::from_parts("TopicTest", "broker-a", 1),
148/// ];
149/// assert_eq!(selector(&mqs, &msg, &1i64), Some(mqs[1].clone()));
150///
151/// // An argument that cannot downcast to `i64` still returns `None`.
152/// assert_eq!(selector(&mqs, &msg, &"not-an-i64"), None);
153/// ```
154pub type MessageQueueSelectorFn =
155 Arc<dyn Fn(&[MessageQueue], &dyn MessageTrait, &dyn std::any::Any) -> Option<MessageQueue> + Send + Sync>;