Skip to main content

rocketmq_client_rust/consumer/
pull_callback.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 crate::ClientError;
18use rocketmq_model::common::message::message_queue::MessageQueue;
19use rocketmq_model::common::mix_all;
20use rocketmq_protocol::code::response_code::ResponseCode;
21use rocketmq_protocol::protocol::heartbeat::subscription_data::SubscriptionData;
22use tracing::warn;
23
24use crate::consumer::consumer_impl::default_mq_push_consumer_impl::DefaultMQPushConsumerImpl;
25use crate::consumer::consumer_impl::default_mq_push_consumer_impl::PULL_TIME_DELAY_MILLS_WHEN_BROKER_FLOW_CONTROL;
26use crate::consumer::consumer_impl::pull_request::PullRequest;
27use crate::consumer::consumer_impl::pull_request_ext::PullResultExt;
28use crate::consumer::consumer_impl::re_balance::Rebalance;
29use crate::consumer::pull_status::PullStatus;
30
31pub type PullCallbackFn = Arc<dyn FnOnce(Option<PullResultExt>, Option<ClientError>) + Send + Sync>;
32
33fn broker_response_code(error: &ClientError) -> Option<ResponseCode> {
34    error.broker_response_code().map(ResponseCode::from)
35}
36
37#[trait_variant::make(PullCallback: Send)]
38pub trait PullCallbackLocal: Sync {
39    async fn on_success(&mut self, pull_result: PullResultExt);
40    fn on_exception(&mut self, e: ClientError);
41}
42
43pub(crate) struct DefaultPullCallback {
44    pub(crate) push_consumer_impl: Arc<DefaultMQPushConsumerImpl>,
45    pub(crate) message_queue_inner: Option<MessageQueue>,
46    pub(crate) subscription_data: Option<SubscriptionData>,
47    pub(crate) pull_request: Option<PullRequest>,
48}
49
50impl PullCallback for DefaultPullCallback {
51    async fn on_success(&mut self, mut pull_result_ext: PullResultExt) {
52        let push_consumer_impl = self.push_consumer_impl.clone();
53
54        let Some(message_queue_inner) = self.message_queue_inner.take() else {
55            warn!("pull callback success ignored: message queue is missing");
56            return;
57        };
58        let Some(subscription_data) = self.subscription_data.take() else {
59            warn!(
60                "pull callback success ignored: subscription data is missing, mq={}",
61                message_queue_inner
62            );
63            if let Some(pull_request) = self.pull_request.take() {
64                let delay = push_consumer_impl.pull_time_delay_mills_when_exception;
65                push_consumer_impl.execute_pull_request_later(pull_request, delay);
66            }
67            return;
68        };
69        let Some(mut pull_request) = self.pull_request.take() else {
70            warn!(
71                "pull callback success ignored: pull request is missing, mq={}",
72                message_queue_inner
73            );
74            return;
75        };
76
77        let Some(pull_api_wrapper) = push_consumer_impl.pull_api_wrapper() else {
78            warn!(
79                "pull callback success ignored: PullAPIWrapper is not initialized, mq={}",
80                message_queue_inner
81            );
82            let delay = push_consumer_impl.pull_time_delay_mills_when_exception;
83            push_consumer_impl.execute_pull_request_later(pull_request, delay);
84            return;
85        };
86        pull_api_wrapper.process_pull_result(&message_queue_inner, &mut pull_result_ext, &subscription_data);
87        match pull_result_ext.pull_result.pull_status {
88            PullStatus::Found => {
89                let prev_request_offset = pull_request.next_offset;
90                pull_request.set_next_offset(pull_result_ext.pull_result.next_begin_offset as i64);
91                let mut first_msg_offset = i64::MAX;
92                if pull_result_ext
93                    .pull_result
94                    .msg_found_list
95                    .as_ref()
96                    .is_none_or(|v| v.is_empty())
97                {
98                    push_consumer_impl.execute_pull_request_immediately(pull_request).await;
99                } else {
100                    let msg_found_list = pull_result_ext
101                        .pull_result
102                        .msg_found_list
103                        .take()
104                        .unwrap_or_default()
105                        .into_iter()
106                        .map(Arc::new)
107                        .collect::<Vec<_>>();
108                    if msg_found_list.is_empty() {
109                        push_consumer_impl.execute_pull_request_immediately(pull_request).await;
110                        return;
111                    }
112                    first_msg_offset = msg_found_list.first().map_or(i64::MAX, |msg| msg.queue_offset);
113                    if push_consumer_impl.consume_message_service().is_none() {
114                        warn!(
115                            "pull callback found messages but ConsumeMessageService is not initialized, mq={}",
116                            message_queue_inner
117                        );
118                        let delay = push_consumer_impl.pull_time_delay_mills_when_exception;
119                        push_consumer_impl.execute_pull_request_later(pull_request, delay);
120                        return;
121                    }
122                    let dispatch_to_consume = pull_request.process_queue.put_message(&msg_found_list).await;
123                    if let Some(consume_message_service) = push_consumer_impl.consume_message_service() {
124                        consume_message_service
125                            .submit_consume_request(
126                                msg_found_list,
127                                pull_request.get_process_queue().clone(),
128                                pull_request.get_message_queue().clone(),
129                                dispatch_to_consume,
130                            )
131                            .await;
132                    }
133                    let pull_interval = push_consumer_impl.consumer_config_snapshot().pull_interval;
134                    if pull_interval > 0 {
135                        push_consumer_impl.execute_pull_request_later(pull_request, pull_interval);
136                    } else {
137                        push_consumer_impl.execute_pull_request_immediately(pull_request).await;
138                    }
139                }
140                if pull_result_ext.pull_result.next_begin_offset < prev_request_offset as u64
141                    || first_msg_offset < prev_request_offset
142                {
143                    warn!(
144                        "[BUG] pull message result maybe data wrong, nextBeginOffset: {} firstMsgOffset: {} \
145                         prevRequestOffset: {}",
146                        pull_result_ext.pull_result.next_begin_offset, prev_request_offset, prev_request_offset
147                    );
148                }
149            }
150            PullStatus::NoNewMsg | PullStatus::NoMatchedMsg => {
151                pull_request.next_offset = pull_result_ext.pull_result.next_begin_offset as i64;
152                push_consumer_impl.correct_tags_offset(&pull_request).await;
153                push_consumer_impl.execute_pull_request_immediately(pull_request).await;
154            }
155
156            PullStatus::OffsetIllegal => {
157                warn!(
158                    "the pull request offset illegal, {},{}",
159                    pull_result_ext.pull_result, pull_result_ext.pull_result.pull_status
160                );
161                pull_request.next_offset = pull_result_ext.pull_result.next_begin_offset as i64;
162                pull_request.process_queue.set_dropped(true);
163
164                let Some(offset_store) = push_consumer_impl.offset_store() else {
165                    warn!(
166                        "offset illegal repair skipped: OffsetStore is not initialized, {}",
167                        pull_request
168                    );
169                    return;
170                };
171                offset_store
172                    .update_and_freeze_offset(pull_request.get_message_queue(), pull_request.next_offset)
173                    .await;
174                offset_store.persist(pull_request.get_message_queue()).await;
175                push_consumer_impl
176                    .rebalance_impl
177                    .remove_process_queue(pull_request.get_message_queue())
178                    .await;
179                match push_consumer_impl
180                    .rebalance_impl
181                    .rebalance_impl_inner
182                    .get_mq_client_factory()
183                {
184                    Some(client_instance) => client_instance.re_balance_immediately(),
185                    None => {
186                        warn!("offset illegal repair skipped immediate rebalance: MQClientInstance is not initialized")
187                    }
188                }
189            }
190        };
191    }
192
193    fn on_exception(&mut self, err: ClientError) {
194        let Some(message_queue_inner) = self.message_queue_inner.take() else {
195            warn!(
196                "pull callback exception ignored: message queue is missing, error={}",
197                err
198            );
199            return;
200        };
201        let Some(pull_request) = self.pull_request.take() else {
202            warn!(
203                "pull callback exception ignored: pull request is missing, mq={}, error={}",
204                message_queue_inner, err
205            );
206            return;
207        };
208        let topic = message_queue_inner.topic_str();
209        let broker_code = broker_response_code(&err);
210        let consumer_group = self
211            .push_consumer_impl
212            .consumer_config_snapshot()
213            .consumer_group
214            .clone();
215        if !topic.starts_with(mix_all::RETRY_GROUP_TOPIC_PREFIX) {
216            if broker_code == Some(ResponseCode::SubscriptionNotLatest) {
217                warn!("the subscription is not latest, group={}", consumer_group,);
218            } else {
219                warn!(
220                    "execute the pull request exception, group={}, broker_code={:?}, error={}",
221                    consumer_group, broker_code, err
222                );
223            }
224        }
225        let time_delay = if broker_code == Some(ResponseCode::FlowControl) {
226            PULL_TIME_DELAY_MILLS_WHEN_BROKER_FLOW_CONTROL
227        } else {
228            self.push_consumer_impl.pull_time_delay_mills_when_exception
229        };
230
231        self.push_consumer_impl
232            .execute_pull_request_later(pull_request, time_delay);
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use bytes::Bytes;
239
240    use super::*;
241    use crate::base::client_config::ClientConfig;
242    use crate::consumer::default_mq_push_consumer::ConsumerConfig;
243    use crate::consumer::pull_result::PullResult;
244
245    fn new_callback() -> DefaultPullCallback {
246        let consumer_config = ConsumerConfig::default();
247        let push_consumer_impl = Arc::new(DefaultMQPushConsumerImpl::new(
248            crate::runtime::test_client_runtime("pull-callback-test"),
249            ClientConfig::default(),
250            consumer_config,
251            None,
252        ));
253        push_consumer_impl.initialize_self_reference();
254        DefaultPullCallback {
255            push_consumer_impl,
256            message_queue_inner: None,
257            subscription_data: None,
258            pull_request: None,
259        }
260    }
261
262    #[test]
263    fn broker_response_code_reads_broker_error_without_downcast() {
264        let error = crate::ClientError::broker_operation_failed(
265            "PULL_MESSAGE",
266            ResponseCode::FlowControl.to_i32(),
267            "flow control",
268        );
269
270        assert_eq!(broker_response_code(&error), Some(ResponseCode::FlowControl));
271    }
272
273    #[tokio::test]
274    async fn success_without_message_queue_is_ignored_without_panic() {
275        let mut callback = new_callback();
276        let pull_result_ext = PullResultExt {
277            pull_result: PullResult::new(PullStatus::NoNewMsg, 0, 0, 0, None),
278            suggest_which_broker_id: 0,
279            message_binary: Some(Bytes::new()),
280            offset_delta: None,
281        };
282
283        PullCallback::on_success(&mut callback, pull_result_ext).await;
284    }
285
286    #[test]
287    fn exception_without_pull_request_is_ignored_without_panic() {
288        let mut callback = new_callback();
289        callback.message_queue_inner = Some(MessageQueue::from_parts("topic", "broker-a", 0));
290
291        PullCallback::on_exception(&mut callback, ClientError::illegal_argument("test error"));
292    }
293}