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