rocketmq_client_rust/consumer/
mq_consumer_inner.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::HashMap;
16use std::collections::HashSet;
17use std::sync::Arc;
18
19use cheetah_string::CheetahString;
20use rocketmq_model::common::consumer::consume_from_where::ConsumeFromWhere;
21use rocketmq_model::common::message::message_ext::MessageExt;
22use rocketmq_model::common::message::message_queue::MessageQueue;
23use rocketmq_protocol::protocol::body::consume_message_directly_result::ConsumeMessageDirectlyResult;
24use rocketmq_protocol::protocol::body::consumer_running_info::ConsumerRunningInfo;
25use rocketmq_protocol::protocol::heartbeat::consume_type::ConsumeType;
26use rocketmq_protocol::protocol::heartbeat::message_model::MessageModel;
27use rocketmq_protocol::protocol::heartbeat::subscription_data::SubscriptionData;
28
29use crate::consumer::consumer_impl::default_lite_pull_consumer_impl::DefaultLitePullConsumerImpl;
30use crate::consumer::consumer_impl::default_mq_push_consumer_impl::DefaultMQPushConsumerImpl;
31use crate::consumer::consumer_impl::pop_request::PopRequest;
32use crate::consumer::consumer_impl::pull_request::PullRequest;
33
34/// The `MQConsumerInnerLocal` trait defines the core functionalities required for a local MQ
35/// consumer. It extends the `MQConsumerInnerAny` trait and requires implementations to be `Sync`
36/// and `'static`.
37
38#[trait_variant::make(MQConsumerInner: Send)]
39pub trait MQConsumerInnerLocal: MQConsumerInnerAny + Sync + 'static {
40    /// Returns the group name of the consumer.
41    fn group_name(&self) -> CheetahString;
42
43    /// Returns the message model used by the consumer.
44    fn message_model(&self) -> MessageModel;
45
46    /// Returns the type of consumption (e.g., push or pull).
47    fn consume_type(&self) -> ConsumeType;
48
49    /// Returns the point from where the consumer should start consuming messages.
50    fn consume_from_where(&self) -> ConsumeFromWhere;
51
52    /// Returns the set of subscriptions for the consumer.
53    fn subscriptions(&self) -> HashSet<SubscriptionData>;
54
55    /// Performs the rebalancing of the consumer.
56    async fn do_rebalance(&self);
57
58    /// Attempts to perform rebalancing asynchronously and returns a `Result` indicating success or
59    /// failure.
60    async fn try_rebalance(&self) -> crate::ClientResult<bool>;
61
62    /// Persists the consumer offset asynchronously.
63    async fn persist_consumer_offset(&self);
64
65    /// Updates the subscription information for a given topic asynchronously.
66    ///
67    /// # Arguments
68    ///
69    /// * `topic` - A string slice that holds the name of the topic.
70    /// * `info` - A reference to a `HashSet` containing `MessageQueue` information.
71    async fn update_topic_subscribe_info(&self, topic: CheetahString, info: &HashSet<MessageQueue>);
72
73    /// Checks if the subscription information for a given topic needs to be updated asynchronously.
74    ///
75    /// # Arguments
76    ///
77    /// * `topic` - A string slice that holds the name of the topic.
78    ///
79    /// # Returns
80    ///
81    /// * `bool` - `true` if the subscription information needs to be updated, `false` otherwise.
82    async fn is_subscribe_topic_need_update(&self, topic: &str) -> bool;
83
84    /// Returns whether the consumer is in unit mode.
85    fn is_unit_mode(&self) -> bool;
86
87    /// Replaces the in-memory offsets for queues belonging to the given topic.
88    async fn reset_offsets(&self, topic: &CheetahString, offsets: HashMap<MessageQueue, i64>);
89
90    /// Returns a snapshot of the consumer offsets for the given topic.
91    async fn consumer_status(&self, topic: &CheetahString) -> HashMap<MessageQueue, i64>;
92
93    /// Returns the running information of the consumer.
94    async fn consumer_running_info(&self) -> ConsumerRunningInfo;
95}
96
97pub trait MQConsumerInnerAny: std::any::Any {
98    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
99
100    fn as_any(&self) -> &dyn std::any::Any;
101}
102
103impl<T: MQConsumerInner> MQConsumerInnerAny for T {
104    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
105        self
106    }
107
108    fn as_any(&self) -> &dyn std::any::Any {
109        self
110    }
111}
112
113#[derive(Clone)]
114pub struct MQConsumerInnerImpl {
115    inner: MQConsumerInnerImplKind,
116}
117
118#[derive(Clone)]
119enum MQConsumerInnerImplKind {
120    Push(Arc<DefaultMQPushConsumerImpl>),
121    LitePull(Arc<DefaultLitePullConsumerImpl>),
122}
123
124impl MQConsumerInnerImpl {
125    pub(crate) fn from_push(default_mqpush_consumer_impl: Arc<DefaultMQPushConsumerImpl>) -> Self {
126        Self {
127            inner: MQConsumerInnerImplKind::Push(default_mqpush_consumer_impl),
128        }
129    }
130
131    pub(crate) fn from_lite_pull(default_lite_pull_consumer_impl: Arc<DefaultLitePullConsumerImpl>) -> Self {
132        Self {
133            inner: MQConsumerInnerImplKind::LitePull(default_lite_pull_consumer_impl),
134        }
135    }
136
137    pub(crate) fn is_push_consumer(&self) -> bool {
138        matches!(self.inner, MQConsumerInnerImplKind::Push(_))
139    }
140
141    pub(crate) async fn pop_message(&mut self, pop_request: PopRequest) {
142        match &mut self.inner {
143            MQConsumerInnerImplKind::Push(consumer) => consumer.pop_message(pop_request).await,
144            MQConsumerInnerImplKind::LitePull(consumer) => {
145                tracing::warn!(
146                    "Ignoring broker pop request for lite pull consumer group={}",
147                    MQConsumerInner::group_name(consumer.as_ref())
148                );
149            }
150        }
151    }
152
153    pub(crate) async fn pull_message(&mut self, pull_request: PullRequest) {
154        match &mut self.inner {
155            MQConsumerInnerImplKind::Push(consumer) => consumer.pull_message(pull_request).await,
156            MQConsumerInnerImplKind::LitePull(consumer) => {
157                tracing::warn!(
158                    "Ignoring broker pull request for lite pull consumer group={}",
159                    MQConsumerInner::group_name(consumer.as_ref())
160                );
161            }
162        }
163    }
164
165    pub(crate) async fn consume_message_directly(
166        &self,
167        msg: MessageExt,
168        broker_name: Option<CheetahString>,
169    ) -> Option<ConsumeMessageDirectlyResult> {
170        match &self.inner {
171            MQConsumerInnerImplKind::Push(consumer) => consumer.consume_message_directly(msg, broker_name).await,
172            MQConsumerInnerImplKind::LitePull(consumer) => {
173                tracing::warn!(
174                    "consumeMessageDirectly is not supported for lite pull consumer group={}",
175                    MQConsumerInner::group_name(consumer.as_ref())
176                );
177                None
178            }
179        }
180    }
181}
182
183impl MQConsumerInner for MQConsumerInnerImpl {
184    #[inline]
185    fn group_name(&self) -> CheetahString {
186        match &self.inner {
187            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::group_name(consumer.as_ref()),
188            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::group_name(consumer.as_ref()),
189        }
190    }
191
192    #[inline]
193    fn message_model(&self) -> MessageModel {
194        match &self.inner {
195            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::message_model(consumer.as_ref()),
196            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::message_model(consumer.as_ref()),
197        }
198    }
199
200    #[inline]
201    fn consume_type(&self) -> ConsumeType {
202        match &self.inner {
203            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::consume_type(consumer.as_ref()),
204            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::consume_type(consumer.as_ref()),
205        }
206    }
207
208    #[inline]
209    fn consume_from_where(&self) -> ConsumeFromWhere {
210        match &self.inner {
211            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::consume_from_where(consumer.as_ref()),
212            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::consume_from_where(consumer.as_ref()),
213        }
214    }
215
216    #[inline]
217    fn subscriptions(&self) -> HashSet<SubscriptionData> {
218        match &self.inner {
219            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::subscriptions(consumer.as_ref()),
220            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::subscriptions(consumer.as_ref()),
221        }
222    }
223
224    #[inline]
225    async fn do_rebalance(&self) {
226        match &self.inner {
227            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::do_rebalance(consumer.as_ref()).await,
228            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::do_rebalance(consumer.as_ref()).await,
229        }
230    }
231
232    #[inline]
233    async fn try_rebalance(&self) -> crate::ClientResult<bool> {
234        match &self.inner {
235            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::try_rebalance(consumer.as_ref()).await,
236            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::try_rebalance(consumer.as_ref()).await,
237        }
238    }
239
240    #[inline]
241    async fn persist_consumer_offset(&self) {
242        match &self.inner {
243            MQConsumerInnerImplKind::Push(consumer) => {
244                MQConsumerInner::persist_consumer_offset(consumer.as_ref()).await
245            }
246            MQConsumerInnerImplKind::LitePull(consumer) => {
247                MQConsumerInner::persist_consumer_offset(consumer.as_ref()).await
248            }
249        }
250    }
251
252    #[inline]
253    async fn update_topic_subscribe_info(&self, topic: CheetahString, info: &HashSet<MessageQueue>) {
254        match &self.inner {
255            MQConsumerInnerImplKind::Push(consumer) => {
256                MQConsumerInner::update_topic_subscribe_info(consumer.as_ref(), topic, info).await
257            }
258            MQConsumerInnerImplKind::LitePull(consumer) => {
259                MQConsumerInner::update_topic_subscribe_info(consumer.as_ref(), topic, info).await
260            }
261        }
262    }
263
264    #[inline]
265    async fn is_subscribe_topic_need_update(&self, topic: &str) -> bool {
266        match &self.inner {
267            MQConsumerInnerImplKind::Push(consumer) => {
268                MQConsumerInner::is_subscribe_topic_need_update(consumer.as_ref(), topic).await
269            }
270            MQConsumerInnerImplKind::LitePull(consumer) => {
271                MQConsumerInner::is_subscribe_topic_need_update(consumer.as_ref(), topic).await
272            }
273        }
274    }
275
276    #[inline]
277    fn is_unit_mode(&self) -> bool {
278        match &self.inner {
279            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::is_unit_mode(consumer.as_ref()),
280            MQConsumerInnerImplKind::LitePull(consumer) => MQConsumerInner::is_unit_mode(consumer.as_ref()),
281        }
282    }
283
284    #[inline]
285    async fn reset_offsets(&self, topic: &CheetahString, offsets: HashMap<MessageQueue, i64>) {
286        match &self.inner {
287            MQConsumerInnerImplKind::Push(consumer) => {
288                MQConsumerInner::reset_offsets(consumer.as_ref(), topic, offsets).await
289            }
290            MQConsumerInnerImplKind::LitePull(consumer) => {
291                MQConsumerInner::reset_offsets(consumer.as_ref(), topic, offsets).await
292            }
293        }
294    }
295
296    #[inline]
297    async fn consumer_status(&self, topic: &CheetahString) -> HashMap<MessageQueue, i64> {
298        match &self.inner {
299            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::consumer_status(consumer.as_ref(), topic).await,
300            MQConsumerInnerImplKind::LitePull(consumer) => {
301                MQConsumerInner::consumer_status(consumer.as_ref(), topic).await
302            }
303        }
304    }
305
306    #[inline]
307    async fn consumer_running_info(&self) -> ConsumerRunningInfo {
308        match &self.inner {
309            MQConsumerInnerImplKind::Push(consumer) => MQConsumerInner::consumer_running_info(consumer.as_ref()).await,
310            MQConsumerInnerImplKind::LitePull(consumer) => {
311                MQConsumerInner::consumer_running_info(consumer.as_ref()).await
312            }
313        }
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use rocketmq_protocol::protocol::heartbeat::consume_type::ConsumeType;
320
321    use super::*;
322    use crate::base::client_config::ClientConfig;
323    use crate::consumer::consumer_impl::default_lite_pull_consumer_impl::LitePullConsumerConfig;
324
325    #[test]
326    fn lite_pull_wrapper_delegates_heartbeat_fields() {
327        let consumer_group = CheetahString::from_static_str("lite_pull_wrapper_group");
328        let consumer_config = Arc::new(LitePullConsumerConfig {
329            consumer_group: consumer_group.clone(),
330            ..Default::default()
331        });
332        let impl_ = Arc::new(DefaultLitePullConsumerImpl::new(
333            crate::runtime::test_client_runtime("mq-consumer-inner-test"),
334            Arc::new(ClientConfig::default()),
335            consumer_config,
336        ));
337        let wrapper = MQConsumerInnerImpl::from_lite_pull(impl_);
338
339        assert_eq!(MQConsumerInner::group_name(&wrapper), consumer_group);
340        assert_eq!(MQConsumerInner::message_model(&wrapper), MessageModel::Clustering);
341        assert_eq!(MQConsumerInner::consume_type(&wrapper), ConsumeType::ConsumeActively);
342        assert_eq!(
343            MQConsumerInner::consume_from_where(&wrapper),
344            ConsumeFromWhere::ConsumeFromLastOffset
345        );
346        assert!(!MQConsumerInner::is_unit_mode(&wrapper));
347        assert!(MQConsumerInner::subscriptions(&wrapper).is_empty());
348    }
349}
rocketmq_client_rust/consumer/mq_consumer_inner.rs

rocketmq_client_rust/consumer/
mq_consumer_inner.rs
