rocketmq_client_rust/consumer/consumer_impl/
default_mq_pull_consumer_impl.rs

1// Copyright 2026 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;
18use std::sync::RwLock as StdRwLock;
19use std::time::Duration;
20
21use crate::ClientError;
22use crate::ClientResult;
23use cheetah_string::CheetahString;
24use rocketmq_model::common::message::message_queue::MessageQueue;
25use rocketmq_protocol::protocol::namespace_util::NamespaceUtil;
26use tokio::sync::Mutex;
27
28use crate::base::client_config::ClientConfig;
29use crate::consumer::consumer_impl::default_lite_pull_consumer_impl::ServiceState;
30use crate::consumer::default_lite_pull_consumer::DefaultLitePullConsumer;
31use crate::consumer::default_mq_pull_consumer::ClassicPullCallback;
32use crate::consumer::default_mq_pull_consumer::PullOptions;
33use crate::consumer::message_queue_listener::ArcMessageQueueListener;
34use crate::consumer::pull_result::PullResult;
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37enum ClassicPullServiceState {
38    Created,
39    Running,
40    Failed,
41    Shutdown,
42}
43
44struct ClassicPullCore {
45    lite_consumer: DefaultLitePullConsumer,
46    consumer_pull_timeout: Duration,
47    broker_suspend_timeout: Duration,
48    consumer_timeout_when_suspend: Duration,
49    state: Mutex<ClassicPullServiceState>,
50    listeners: Arc<StdRwLock<HashMap<CheetahString, ArcMessageQueueListener>>>,
51}
52
53struct ClassicPullQueueListener {
54    listeners: Arc<StdRwLock<HashMap<CheetahString, ArcMessageQueueListener>>>,
55    namespace: Option<CheetahString>,
56}
57
58impl ClassicPullQueueListener {
59    fn new(listeners: Arc<StdRwLock<HashMap<CheetahString, ArcMessageQueueListener>>>) -> Self {
60        Self {
61            listeners,
62            namespace: None,
63        }
64    }
65
66    fn with_namespace(mut self, namespace: Option<CheetahString>) -> Self {
67        self.namespace = namespace.filter(|namespace| !namespace.is_empty());
68        self
69    }
70
71    fn without_namespace(&self, resource: &str) -> CheetahString {
72        match &self.namespace {
73            Some(namespace) => NamespaceUtil::without_namespace_with_namespace(resource, namespace.as_str()).into(),
74            None => NamespaceUtil::without_namespace(resource).into(),
75        }
76    }
77
78    fn queues_without_namespace(&self, queues: &HashSet<MessageQueue>) -> HashSet<MessageQueue> {
79        queues
80            .iter()
81            .map(|queue| {
82                MessageQueue::from_parts(
83                    self.without_namespace(queue.topic()),
84                    queue.broker_name().clone(),
85                    queue.queue_id(),
86                )
87            })
88            .collect()
89    }
90}
91
92impl crate::consumer::message_queue_listener::MessageQueueListener for ClassicPullQueueListener {
93    fn message_queue_changed(&self, topic: &str, mq_all: &HashSet<MessageQueue>, mq_divided: &HashSet<MessageQueue>) {
94        let topic = self.without_namespace(topic);
95        let listener = self
96            .listeners
97            .read()
98            .unwrap_or_else(|poisoned| poisoned.into_inner())
99            .get(&topic)
100            .cloned();
101        if let Some(listener) = listener {
102            let mq_all = self.queues_without_namespace(mq_all);
103            let mq_divided = self.queues_without_namespace(mq_divided);
104            listener.message_queue_changed(topic.as_str(), &mq_all, &mq_divided);
105        }
106    }
107}
108
109impl ClassicPullCore {
110    async fn start(self: &Arc<Self>) -> ClientResult<()> {
111        let mut state = self.state.lock().await;
112        match *state {
113            ClassicPullServiceState::Created => {}
114            ClassicPullServiceState::Running => {
115                return Err(crate::mq_client_err!("DefaultMQPullConsumer already started"));
116            }
117            ClassicPullServiceState::Failed => {
118                return Err(crate::mq_client_err!(
119                    "DefaultMQPullConsumer start previously failed; create a new consumer"
120                ));
121            }
122            ClassicPullServiceState::Shutdown => {
123                return Err(crate::mq_client_err!(
124                    "DefaultMQPullConsumer has been shut down; create a new consumer"
125                ));
126            }
127        }
128
129        let topics = self.listener_snapshot().keys().cloned().collect::<Vec<_>>();
130        for topic in topics {
131            let topic = self.lite_consumer.classic_topic_with_namespace(topic.as_str());
132            self.lite_consumer.register_classic_pull_subscription(&topic).await?;
133        }
134        if let Err(error) = self.lite_consumer.start().await {
135            *state = ClassicPullServiceState::Failed;
136            return Err(error);
137        }
138        *state = ClassicPullServiceState::Running;
139        Ok(())
140    }
141
142    fn listener_snapshot(&self) -> HashMap<CheetahString, ArcMessageQueueListener> {
143        self.listeners
144            .read()
145            .unwrap_or_else(|poisoned| poisoned.into_inner())
146            .clone()
147    }
148
149    async fn shutdown(&self) -> ClientResult<()> {
150        let mut state = self.state.lock().await;
151        match *state {
152            ClassicPullServiceState::Shutdown => return Ok(()),
153            ClassicPullServiceState::Created => {
154                *state = ClassicPullServiceState::Shutdown;
155                return Ok(());
156            }
157            ClassicPullServiceState::Running | ClassicPullServiceState::Failed => {}
158        }
159        self.lite_consumer.shutdown().await;
160        *state = ClassicPullServiceState::Shutdown;
161        Ok(())
162    }
163
164    async fn ensure_running(&self) -> ClientResult<()> {
165        if *self.state.lock().await != ClassicPullServiceState::Running {
166            return Err(ClientError::not_initialized(
167                "DefaultMQPullConsumer not started. Call start() first.",
168            ));
169        }
170        self.lite_consumer.try_impl_()?.ensure_classic_pull_running()
171    }
172}
173
174/// Functional implementation handle retained under the Java compatibility name.
175#[deprecated(
176    since = "0.9.0",
177    note = "Classic Pull is retained for compatibility; prefer DefaultLitePullConsumer for new applications"
178)]
179#[derive(Clone, Default)]
180pub struct DefaultMQPullConsumerImpl {
181    core: Option<Arc<ClassicPullCore>>,
182}
183
184#[allow(deprecated)]
185impl DefaultMQPullConsumerImpl {
186    /// Creates a detached implementation marker for source compatibility.
187    pub fn new() -> ClientResult<Self> {
188        Ok(Self::default())
189    }
190
191    pub(crate) fn from_lite_consumer(
192        lite_consumer: DefaultLitePullConsumer,
193        consumer_pull_timeout: Duration,
194        broker_suspend_timeout: Duration,
195        consumer_timeout_when_suspend: Duration,
196    ) -> ClientResult<Self> {
197        let listeners = Arc::new(StdRwLock::new(HashMap::new()));
198        let namespace = lite_consumer.client_config().resolved_namespace();
199        lite_consumer.set_classic_pull_message_queue_listener(Some(Arc::new(
200            ClassicPullQueueListener::new(listeners.clone()).with_namespace(namespace),
201        )))?;
202        Ok(Self {
203            core: Some(Arc::new(ClassicPullCore {
204                lite_consumer,
205                consumer_pull_timeout,
206                broker_suspend_timeout,
207                consumer_timeout_when_suspend,
208                state: Mutex::new(ClassicPullServiceState::Created),
209                listeners,
210            })),
211        })
212    }
213
214    fn core(&self) -> ClientResult<&Arc<ClassicPullCore>> {
215        self.core.as_ref().ok_or_else(|| {
216            ClientError::not_initialized(
217                "DefaultMQPullConsumerImpl is detached; create a consumer with DefaultMQPullConsumer::builder",
218            )
219        })
220    }
221
222    /// Starts the implementation.
223    ///
224    /// # Errors
225    ///
226    /// Returns an initialization, lifecycle, or underlying client startup error.
227    pub async fn start(&self) -> ClientResult<()> {
228        self.core()?.start().await
229    }
230
231    /// Shuts the implementation down and awaits its runtime-owned client tasks.
232    ///
233    /// # Errors
234    ///
235    /// Returns an initialization or underlying client shutdown error.
236    pub async fn shutdown(&self) -> ClientResult<()> {
237        self.core()?.shutdown().await
238    }
239
240    /// Returns whether the implementation is running.
241    pub async fn is_running(&self) -> bool {
242        let Ok(core) = self.core() else {
243            return false;
244        };
245        *core.state.lock().await == ClassicPullServiceState::Running
246            && core
247                .lite_consumer
248                .try_impl_()
249                .is_ok_and(|implementation| implementation.service_state() == ServiceState::Running)
250    }
251
252    pub(crate) fn consumer_pull_timeout(&self) -> Duration {
253        self.core
254            .as_ref()
255            .map_or(Duration::from_secs(10), |core| core.consumer_pull_timeout)
256    }
257
258    pub(crate) fn client_config(&self) -> ClientResult<Arc<ClientConfig>> {
259        Ok(self.core()?.lite_consumer.client_config())
260    }
261
262    pub(crate) fn broker_suspend_timeout(&self) -> Duration {
263        self.core
264            .as_ref()
265            .map_or(Duration::from_secs(20), |core| core.broker_suspend_timeout)
266    }
267
268    pub(crate) fn consumer_timeout_when_suspend(&self) -> Duration {
269        self.core
270            .as_ref()
271            .map_or(Duration::from_secs(30), |core| core.consumer_timeout_when_suspend)
272    }
273
274    pub(crate) async fn pull_with_options(&self, options: PullOptions) -> ClientResult<PullResult> {
275        let core = self.core()?;
276        core.ensure_running().await?;
277        let message_queue = core.lite_consumer.queue_with_namespace(options.message_queue());
278        let options = options.with_message_queue(message_queue);
279        core.lite_consumer.try_impl_()?.classic_pull(&options).await
280    }
281
282    pub(crate) async fn pull_async_with_options<C>(&self, options: PullOptions, callback: C) -> ClientResult<()>
283    where
284        C: ClassicPullCallback,
285    {
286        let core = self.core()?;
287        core.ensure_running().await?;
288        let message_queue = core.lite_consumer.queue_with_namespace(options.message_queue());
289        let options = options.with_message_queue(message_queue);
290        core.lite_consumer
291            .try_impl_()?
292            .classic_pull_async(options, callback)
293            .await
294    }
295
296    pub(crate) async fn fetch_subscribe_message_queues(&self, topic: &str) -> ClientResult<Vec<MessageQueue>> {
297        let core = self.core()?;
298        core.ensure_running().await?;
299        core.lite_consumer.fetch_message_queues(topic).await
300    }
301
302    pub(crate) async fn register_message_queue_listener(
303        &self,
304        topic: &str,
305        listener: ArcMessageQueueListener,
306    ) -> ClientResult<()> {
307        if topic.trim().is_empty() {
308            return Err(crate::mq_client_err!("topic is blank"));
309        }
310        let core = self.core()?;
311        let topic = CheetahString::from_slice(topic);
312        core.listeners
313            .write()
314            .unwrap_or_else(|poisoned| poisoned.into_inner())
315            .insert(topic.clone(), listener);
316        if *core.state.lock().await == ClassicPullServiceState::Running {
317            let wrapped_topic = core.lite_consumer.classic_topic_with_namespace(topic.as_str());
318            core.lite_consumer
319                .register_classic_pull_subscription(&wrapped_topic)
320                .await?;
321        }
322        Ok(())
323    }
324
325    pub(crate) async fn update_consume_offset(&self, message_queue: &MessageQueue, offset: i64) -> ClientResult<()> {
326        let core = self.core()?;
327        core.ensure_running().await?;
328        let queue = core.lite_consumer.queue_with_namespace(message_queue);
329        core.lite_consumer
330            .try_impl_()?
331            .update_classic_pull_offset(&queue, offset)
332            .await
333    }
334
335    pub(crate) async fn fetch_consume_offset(
336        &self,
337        message_queue: &MessageQueue,
338        from_store: bool,
339    ) -> ClientResult<i64> {
340        let core = self.core()?;
341        core.ensure_running().await?;
342        let queue = core.lite_consumer.queue_with_namespace(message_queue);
343        core.lite_consumer
344            .try_impl_()?
345            .fetch_classic_pull_offset(&queue, from_store)
346            .await
347    }
348
349    pub(crate) async fn search_offset(&self, message_queue: &MessageQueue, timestamp: u64) -> ClientResult<i64> {
350        let core = self.core()?;
351        core.ensure_running().await?;
352        core.lite_consumer.offset_for_timestamp(message_queue, timestamp).await
353    }
354
355    pub(crate) async fn max_offset(&self, message_queue: &MessageQueue) -> ClientResult<i64> {
356        let core = self.core()?;
357        core.ensure_running().await?;
358        core.lite_consumer.max_offset(message_queue).await
359    }
360
361    pub(crate) async fn min_offset(&self, message_queue: &MessageQueue) -> ClientResult<i64> {
362        let core = self.core()?;
363        core.ensure_running().await?;
364        core.lite_consumer.min_offset(message_queue).await
365    }
366
367    /// Returns a live rebalance compatibility marker.
368    ///
369    /// # Errors
370    ///
371    /// Returns an initialization error for a detached implementation marker.
372    pub fn rebalance_impl(&self) -> ClientResult<crate::legacy::RebalancePullImpl> {
373        self.core()?;
374        Ok(crate::legacy::RebalancePullImpl)
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use std::sync::atomic::AtomicUsize;
381    use std::sync::atomic::Ordering;
382
383    use super::*;
384    use crate::consumer::message_queue_listener::MessageQueueListener;
385
386    struct AssignmentListener {
387        calls: Arc<AtomicUsize>,
388        assigned: Arc<AtomicUsize>,
389    }
390
391    impl MessageQueueListener for AssignmentListener {
392        fn message_queue_changed(
393            &self,
394            _topic: &str,
395            _mq_all: &HashSet<MessageQueue>,
396            mq_divided: &HashSet<MessageQueue>,
397        ) {
398            self.calls.fetch_add(1, Ordering::SeqCst);
399            self.assigned.store(mq_divided.len(), Ordering::SeqCst);
400        }
401    }
402
403    #[test]
404    fn classic_queue_listener_forwards_the_rebalance_assignment_for_its_topic() {
405        let listeners = Arc::new(StdRwLock::new(HashMap::new()));
406        let calls = Arc::new(AtomicUsize::new(0));
407        let assigned = Arc::new(AtomicUsize::new(0));
408        listeners.write().expect("listener registry should be writable").insert(
409            CheetahString::from_static_str("TopicA"),
410            Arc::new(AssignmentListener {
411                calls: calls.clone(),
412                assigned: assigned.clone(),
413            }) as ArcMessageQueueListener,
414        );
415        let dispatcher = ClassicPullQueueListener::new(listeners);
416        let all = HashSet::from([
417            MessageQueue::from_parts("TopicA", "broker-a", 0),
418            MessageQueue::from_parts("TopicA", "broker-a", 1),
419        ]);
420        let divided = HashSet::from([MessageQueue::from_parts("TopicA", "broker-a", 1)]);
421
422        dispatcher.message_queue_changed("TopicA", &all, &divided);
423        dispatcher.message_queue_changed("OtherTopic", &all, &divided);
424
425        assert_eq!(calls.load(Ordering::SeqCst), 1);
426        assert_eq!(assigned.load(Ordering::SeqCst), 1);
427    }
428}
rocketmq_client_rust/consumer/consumer_impl/default_mq_pull_consumer_impl.rs

rocketmq_client_rust/consumer/consumer_impl/
default_mq_pull_consumer_impl.rs
