rocketmq_client_rust/implementation/mq_client_api_impl/
consumer.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 super::admin::parse_lite_order_count_info_like_java;
16use super::admin::pop_msg_queue_offset_for_index;
17use super::admin::split_lite_dispatch_value;
18use super::callback_executor::ClientCallbackExecutor;
19use super::request_builder::notification_request;
20use super::response_decoder::notify_result_from_response;
21use super::*;
22use rocketmq_error::fields;
23use rocketmq_error::Error;
24use rocketmq_error::ErrorContext;
25use rocketmq_transport::api::OutboundRequestContractReason;
26use rocketmq_transport::api::OutboundRequestRejectionReason;
27use rocketmq_transport::api::OutboundRequestStage;
28
29fn pop_background_task_cancelled(actual: impl Into<String>) -> ClientError {
30    ClientError::invalid_state("active client POP request", actual.into())
31}
32
33fn pop_background_task_spawn_failed(source: impl std::error::Error + Send + Sync + 'static) -> ClientError {
34    let context = ErrorContext::new()
35        .with_text(fields::EXPECTED_STATE, "active client POP request")
36        .with_text(fields::ACTUAL_STATE, "client POP task spawn failed");
37    ClientError::from_error(
38        Error::caused_by(&rocketmq_error::CLIENT_LIFECYCLE_INVALID_STATE, source).with_context(context),
39    )
40}
41
42fn consumer_request_error(operation: &'static str, input: RetryInput) -> ClientError {
43    match input {
44        RetryInput::Transport(error) => ClientError::from_shared(error.into_shared_error()),
45        RetryInput::Rejected(rejection) => match rejection.reason() {
46            OutboundRequestRejectionReason::DeadlineExpired => {
47                let mut context = ErrorContext::new().with_text(fields::OPERATION_DIAGNOSTIC, operation);
48                if let Some(timeout_millis) = rejection.timeout_millis() {
49                    context = context.with_u64(fields::TIMEOUT_MS, timeout_millis);
50                }
51                ClientError::from_shared(Arc::new(
52                    Error::new(&rocketmq_error::CORE_OPERATION_TIMED_OUT).with_context(context),
53                ))
54            }
55            OutboundRequestRejectionReason::ClientStopping => ClientError::not_started(),
56            OutboundRequestRejectionReason::QueueSaturated => ClientError::from_shared(Arc::new(Error::new(
57                &rocketmq_error::TRANSPORT_ADMISSION_QUEUE_SATURATED,
58            ))),
59            OutboundRequestRejectionReason::Cancelled
60            | OutboundRequestRejectionReason::SessionClosed
61            | OutboundRequestRejectionReason::EndpointUnavailable => {
62                let phase = match rejection.reason() {
63                    OutboundRequestRejectionReason::SessionClosed => "closed",
64                    _ => consumer_request_stage_phase(rejection.stage()),
65                };
66                let mut context = ErrorContext::new().with_text(fields::PHASE, phase);
67                if rejection.remote_addr_present() {
68                    context = context.with_secret_presence(fields::REMOTE_ADDR_PRESENT);
69                }
70                ClientError::from_shared(Arc::new(
71                    Error::new(&rocketmq_error::TRANSPORT_CONNECTION_FAILED).with_context(context),
72                ))
73            }
74        },
75        RetryInput::Contract(contract) => match contract.reason() {
76            OutboundRequestContractReason::NameServerEndpointMissing => {
77                ClientError::from_shared(Arc::new(Error::new(&rocketmq_error::TRANSPORT_CONNECTION_FAILED)))
78            }
79        },
80        RetryInput::Response { terminal_error, .. } | RetryInput::BusinessError(terminal_error) => terminal_error,
81        // Consumer request methods never construct producer status or route-policy inputs.
82        RetryInput::SendStatus(_) | RetryInput::RouteUnavailable => {
83            unreachable!("consumer request returned an unrelated retry input")
84        }
85    }
86}
87
88const fn consumer_request_stage_phase(stage: OutboundRequestStage) -> &'static str {
89    match stage {
90        OutboundRequestStage::BeforeWrite => "before_write",
91        OutboundRequestStage::Writing => "writing",
92        OutboundRequestStage::AwaitingResponse => "awaiting_response",
93        OutboundRequestStage::ResponseReceived => "response_received",
94    }
95}
96
97pub struct ConsumerClient<'a> {
98    api: &'a MQClientAPIImpl,
99}
100
101impl ConsumerClient<'_> {
102    pub async fn consumer_offset(
103        &self,
104        addr: &str,
105        request_header: QueryConsumerOffsetRequestHeader,
106        timeout_millis: u64,
107    ) -> crate::ClientResult<i64> {
108        self.api
109            .query_consumer_offset(addr, request_header, timeout_millis)
110            .await
111    }
112}
113
114impl MQClientAPIImpl {
115    pub(super) fn spawn_pop_callback_task<PC, F>(
116        service_context: &ChildServiceContext,
117        tracker: &TaskTracker,
118        shutdown_token: &CancellationToken,
119        callback_executor: ClientCallbackExecutor,
120        request: F,
121        pop_callback: PC,
122    ) where
123        PC: PopCallback + 'static,
124        F: Future<Output = crate::ClientResult<PopResult>> + Send + 'static,
125    {
126        let callback = Arc::new(std::sync::Mutex::new(Some(pop_callback)));
127        if shutdown_token.is_cancelled() {
128            if let Some(mut callback) = callback
129                .lock()
130                .unwrap_or_else(std::sync::PoisonError::into_inner)
131                .take()
132            {
133                callback.on_error(pop_background_task_cancelled("client API is shutting down"));
134            }
135            return;
136        }
137
138        let task_callback = callback.clone();
139        let shutdown_token = shutdown_token.clone();
140        let tracked_task = tracker.track_future(async move {
141            let outcome = tokio::select! {
142                biased;
143                _ = shutdown_token.cancelled() => {
144                    Err(pop_background_task_cancelled("client API shutdown cancelled the POP request"))
145                }
146                outcome = request => outcome,
147            };
148            let callback = task_callback
149                .lock()
150                .unwrap_or_else(std::sync::PoisonError::into_inner)
151                .take();
152            let Some(mut callback) = callback else {
153                return;
154            };
155            let _ = callback_executor
156                .execute(async move {
157                    match outcome {
158                        Ok(pop_result) => callback.on_success(pop_result).await,
159                        Err(error) => callback.on_error(error),
160                    }
161                })
162                .await;
163        });
164
165        if let Err(error) = spawn_client_task_with_context(
166            service_context,
167            "rocketmq-client-pop-message-async",
168            Box::pin(tracked_task),
169        ) {
170            if let Some(mut callback) = callback
171                .lock()
172                .unwrap_or_else(std::sync::PoisonError::into_inner)
173                .take()
174            {
175                callback.on_error(pop_background_task_spawn_failed(error));
176            }
177        }
178    }
179
180    #[must_use]
181    pub fn consumer_client(&self) -> ConsumerClient<'_> {
182        ConsumerClient { api: self }
183    }
184}
185
186impl MQClientAPIImpl {
187    pub async fn notification(
188        &self,
189        broker_addr: &CheetahString,
190        request_header: NotificationRequestHeader,
191        timeout_millis: u64,
192    ) -> crate::ClientResult<bool> {
193        self.notification_with_polling_stats(broker_addr, request_header, timeout_millis)
194            .await
195            .map(|result| result.is_has_msg())
196    }
197
198    pub async fn notification_with_polling_stats(
199        &self,
200        broker_addr: &CheetahString,
201        request_header: NotificationRequestHeader,
202        timeout_millis: u64,
203    ) -> crate::ClientResult<NotifyResult> {
204        let request = notification_request(&self.command_factory, request_header);
205        let outcome = self
206            .remoting_client
207            .invoke_request(Some(broker_addr), request, timeout_millis)
208            .await;
209        let response = match outcome {
210            Ok(OutboundRequestOutcome::Response(response)) => response,
211            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
212                return Err(consumer_request_error(
213                    "notification_with_polling_stats",
214                    RetryInput::Rejected(rejection),
215                ));
216            }
217            Ok(OutboundRequestOutcome::Contract(contract)) => {
218                return Err(consumer_request_error(
219                    "notification_with_polling_stats",
220                    RetryInput::Contract(contract),
221                ));
222            }
223            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
224        };
225
226        match ResponseCode::from(response.code()) {
227            ResponseCode::Success => notify_result_from_response(&response),
228            _ => Err(client_broker_err!(
229                response.code(),
230                response.remark().map_or_else(String::new, |remark| remark.to_string()),
231                broker_addr
232            )),
233        }
234    }
235
236    pub async fn get_consumer_id_list_by_group(
237        &self,
238        addr: &str,
239        consumer_group: &str,
240        timeout_millis: u64,
241    ) -> crate::ClientResult<Vec<CheetahString>> {
242        let request_header = GetConsumerListByGroupRequestHeader {
243            consumer_group: CheetahString::from_slice(consumer_group),
244            rpc: None,
245        };
246        let request = self.create_request_command(RequestCode::GetConsumerListByGroup, request_header);
247        let outcome = self
248            .remoting_client
249            .invoke_request(
250                Some(mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref()),
251                request,
252                timeout_millis,
253            )
254            .await;
255        let response = match outcome {
256            Ok(OutboundRequestOutcome::Response(response)) => response,
257            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
258                return Err(consumer_request_error(
259                    "get_consumer_id_list_by_group",
260                    RetryInput::Rejected(rejection),
261                ));
262            }
263            Ok(OutboundRequestOutcome::Contract(contract)) => {
264                return Err(consumer_request_error(
265                    "get_consumer_id_list_by_group",
266                    RetryInput::Contract(contract),
267                ));
268            }
269            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
270        };
271        match ResponseCode::from(response.code()) {
272            ResponseCode::Success => {
273                let _body = response.body();
274                if let Some(body) = response.body() {
275                    return match GetConsumerListByGroupResponseBody::decode(body) {
276                        Ok(value) => Ok(value.consumer_id_list),
277                        Err(_e) => Err(mq_client_err!(response
278                            .remark()
279                            .map_or("".to_string(), |s| s.to_string()))),
280                    };
281                }
282            }
283            _ => {
284                return Err(client_broker_err!(
285                    response.code(),
286                    response.remark().map_or("".to_string(), |s| s.to_string()),
287                    addr
288                ))
289            }
290        }
291        Err(client_broker_err!(
292            response.code(),
293            response.remark().map_or("".to_string(), |s| s.to_string()),
294            addr
295        ))
296    }
297
298    pub async fn get_consumer_connection_list(
299        &self,
300        addr: &str,
301        consumer_group: CheetahString,
302        timeout_millis: u64,
303    ) -> crate::ClientResult<rocketmq_protocol::protocol::body::consumer_connection::ConsumerConnection> {
304        let request_header = GetConsumerConnectionListRequestHeader {
305            consumer_group,
306            rpc_request_header: None,
307        };
308        let request = self.create_request_command(RequestCode::GetConsumerConnectionList, request_header);
309        let outcome = self
310            .remoting_client
311            .invoke_request(
312                Some(mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref()),
313                request,
314                timeout_millis,
315            )
316            .await;
317        let response = match outcome {
318            Ok(OutboundRequestOutcome::Response(response)) => response,
319            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
320                return Err(consumer_request_error(
321                    "get_consumer_connection_list",
322                    RetryInput::Rejected(rejection),
323                ));
324            }
325            Ok(OutboundRequestOutcome::Contract(contract)) => {
326                return Err(consumer_request_error(
327                    "get_consumer_connection_list",
328                    RetryInput::Contract(contract),
329                ));
330            }
331            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
332        };
333        match ResponseCode::from(response.code()) {
334            ResponseCode::Success => {
335                if let Some(body) = response.body() {
336                    return ConsumerConnection::decode(body).map_err(ClientError::from);
337                }
338            }
339            _ => {
340                return Err(client_broker_err!(
341                    response.code(),
342                    response.remark().map_or("".to_string(), |s| s.to_string()),
343                    addr
344                ))
345            }
346        }
347        Err(client_broker_err!(
348            response.code(),
349            response.remark().map_or("".to_string(), |s| s.to_string()),
350            addr
351        ))
352    }
353
354    pub async fn get_producer_connection_list(
355        &self,
356        addr: &str,
357        producer_group: CheetahString,
358        timeout_millis: u64,
359    ) -> crate::ClientResult<ProducerConnection> {
360        let request_header = GetProducerConnectionListRequestHeader {
361            producer_group,
362            rpc_request_header: None,
363        };
364        let request = self.create_request_command(RequestCode::GetProducerConnectionList, request_header);
365        let outcome = self
366            .remoting_client
367            .invoke_request(
368                Some(mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref()),
369                request,
370                timeout_millis,
371            )
372            .await;
373        let response = match outcome {
374            Ok(OutboundRequestOutcome::Response(response)) => response,
375            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
376                return Err(consumer_request_error(
377                    "get_producer_connection_list",
378                    RetryInput::Rejected(rejection),
379                ));
380            }
381            Ok(OutboundRequestOutcome::Contract(contract)) => {
382                return Err(consumer_request_error(
383                    "get_producer_connection_list",
384                    RetryInput::Contract(contract),
385                ));
386            }
387            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
388        };
389        match ResponseCode::from(response.code()) {
390            ResponseCode::Success => {
391                if let Some(body) = response.body() {
392                    return ProducerConnection::decode(body).map_err(ClientError::from);
393                }
394            }
395            _ => {
396                return Err(client_broker_err!(
397                    response.code(),
398                    response.remark().map_or("".to_string(), |s| s.to_string()),
399                    addr
400                ));
401            }
402        }
403        Err(client_broker_err!(
404            response.code(),
405            response.remark().map_or("".to_string(), |s| s.to_string()),
406            addr
407        ))
408    }
409
410    pub async fn invoke_broker_to_get_consumer_status(
411        &self,
412        addr: &str,
413        topic: CheetahString,
414        group: CheetahString,
415        client_addr: CheetahString,
416        timeout_millis: u64,
417    ) -> crate::ClientResult<HashMap<CheetahString, HashMap<MessageQueue, i64>>> {
418        let request_header = GetConsumerStatusRequestHeader {
419            topic,
420            group,
421            client_addr: if client_addr.is_empty() {
422                None
423            } else {
424                Some(client_addr)
425            },
426            topic_request_header: None,
427        };
428        let request = self.create_request_command(RequestCode::InvokeBrokerToGetConsumerStatus, request_header);
429        let outcome = self
430            .remoting_client
431            .invoke_request(
432                Some(mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref()),
433                request,
434                timeout_millis,
435            )
436            .await;
437        let response = match outcome {
438            Ok(OutboundRequestOutcome::Response(response)) => response,
439            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
440                return Err(consumer_request_error(
441                    "invoke_broker_to_get_consumer_status",
442                    RetryInput::Rejected(rejection),
443                ));
444            }
445            Ok(OutboundRequestOutcome::Contract(contract)) => {
446                return Err(consumer_request_error(
447                    "invoke_broker_to_get_consumer_status",
448                    RetryInput::Contract(contract),
449                ));
450            }
451            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
452        };
453        match ResponseCode::from(response.code()) {
454            ResponseCode::Success => {
455                if let Some(body) = response.body() {
456                    if let Some(status_body) = GetConsumerStatusBody::decode(body) {
457                        return Ok(status_body.consumer_table);
458                    }
459                }
460                Ok(HashMap::new())
461            }
462            _ => Err(mq_client_err!(
463                response.code(),
464                format!(
465                    "invoke broker to get consumer status failed, remark={}",
466                    response.remark().map_or("".to_string(), |s| s.to_string()),
467                )
468            )),
469        }
470    }
471
472    pub async fn get_all_producer_info(
473        &self,
474        addr: &str,
475        timeout_millis: u64,
476    ) -> crate::ClientResult<ProducerTableInfo> {
477        let request = self.create_request_command(RequestCode::GetAllProducerInfo, EmptyHeader {});
478        let outcome = self
479            .remoting_client
480            .invoke_request(
481                Some(mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref()),
482                request,
483                timeout_millis,
484            )
485            .await;
486        let response = match outcome {
487            Ok(OutboundRequestOutcome::Response(response)) => response,
488            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
489                return Err(consumer_request_error(
490                    "get_all_producer_info",
491                    RetryInput::Rejected(rejection),
492                ));
493            }
494            Ok(OutboundRequestOutcome::Contract(contract)) => {
495                return Err(consumer_request_error(
496                    "get_all_producer_info",
497                    RetryInput::Contract(contract),
498                ));
499            }
500            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
501        };
502        match ResponseCode::from(response.code()) {
503            ResponseCode::Success => {
504                if let Some(body) = response.body() {
505                    return ProducerTableInfo::decode(body).map_err(ClientError::from);
506                }
507            }
508            _ => {
509                return Err(client_broker_err!(
510                    response.code(),
511                    response.remark().map_or("".to_string(), |s| s.to_string()),
512                    addr
513                ));
514            }
515        }
516        Err(client_broker_err!(
517            response.code(),
518            response.remark().map_or("".to_string(), |s| s.to_string()),
519            addr
520        ))
521    }
522
523    pub async fn update_consumer_offset_oneway(
524        &self,
525        addr: &str,
526        request_header: UpdateConsumerOffsetRequestHeader,
527        timeout_millis: u64,
528    ) -> crate::ClientResult<()> {
529        let request = self.create_request_command(RequestCode::UpdateConsumerOffset, request_header);
530        self.remoting_client
531            .invoke_request_oneway(
532                mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref(),
533                request,
534                timeout_millis,
535            )
536            .await
537            .map_err(ClientError::from)
538    }
539
540    pub async fn update_consumer_offset_one_way(
541        &self,
542        addr: &str,
543        request_header: UpdateConsumerOffsetRequestHeader,
544        timeout_millis: u64,
545    ) -> crate::ClientResult<()> {
546        self.update_consumer_offset_oneway(addr, request_header, timeout_millis)
547            .await
548    }
549
550    pub async fn update_consumer_offset(
551        &self,
552        addr: &CheetahString,
553        request_header: UpdateConsumerOffsetRequestHeader,
554        timeout_millis: u64,
555    ) -> crate::ClientResult<()> {
556        let request = self.create_request_command(RequestCode::UpdateConsumerOffset, request_header);
557        let outcome = self
558            .remoting_client
559            .invoke_request(Some(addr), request, timeout_millis)
560            .await;
561        let response = match outcome {
562            Ok(OutboundRequestOutcome::Response(response)) => response,
563            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
564                return Err(consumer_request_error(
565                    "update_consumer_offset",
566                    RetryInput::Rejected(rejection),
567                ));
568            }
569            Ok(OutboundRequestOutcome::Contract(contract)) => {
570                return Err(consumer_request_error(
571                    "update_consumer_offset",
572                    RetryInput::Contract(contract),
573                ));
574            }
575            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
576        };
577        if ResponseCode::from(response.code()) != ResponseCode::Success {
578            Err(client_broker_err!(
579                response.code(),
580                response.remark().map_or("".to_string(), |s| s.to_string()),
581                addr
582            ))
583        } else {
584            Ok(())
585        }
586    }
587
588    pub async fn update_consumer_offset_async(
589        &self,
590        addr: &CheetahString,
591        request_header: UpdateConsumerOffsetRequestHeader,
592        timeout_millis: u64,
593    ) -> crate::ClientResult<()> {
594        self.update_consumer_offset(addr, request_header, timeout_millis).await
595    }
596
597    pub async fn query_consumer_offset(
598        &self,
599        addr: &str,
600        request_header: QueryConsumerOffsetRequestHeader,
601        timeout_millis: u64,
602    ) -> crate::ClientResult<i64> {
603        let request = self.create_request_command(RequestCode::QueryConsumerOffset, request_header);
604        let outcome = self
605            .remoting_client
606            .invoke_request(
607                Some(mix_all::broker_vip_channel(self.client_config.vip_channel_enabled, addr).as_ref()),
608                request,
609                timeout_millis,
610            )
611            .await;
612        let response = match outcome {
613            Ok(OutboundRequestOutcome::Response(response)) => response,
614            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
615                return Err(consumer_request_error(
616                    "query_consumer_offset",
617                    RetryInput::Rejected(rejection),
618                ));
619            }
620            Ok(OutboundRequestOutcome::Contract(contract)) => {
621                return Err(consumer_request_error(
622                    "query_consumer_offset",
623                    RetryInput::Contract(contract),
624                ));
625            }
626            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
627        };
628        match ResponseCode::from(response.code()) {
629            ResponseCode::Success => {
630                let response_header = response.decode_command_custom_header::<QueryConsumerOffsetResponseHeader>()?;
631                return response_header.offset.ok_or_else(|| {
632                    client_broker_err!(
633                        response.code(),
634                        "QueryConsumerOffset response header missing offset",
635                        addr
636                    )
637                });
638            }
639            ResponseCode::QueryNotFound => {
640                return Err(client_broker_err!(
641                    response.code(),
642                    response.remark().map_or("".to_string(), |s| s.to_string()),
643                    addr
644                ));
645            }
646            _ => {}
647        }
648        Err(client_broker_err!(
649            response.code(),
650            response.remark().map_or("".to_string(), |s| s.to_string()),
651            addr
652        ))
653    }
654
655    pub async fn query_consumer_offset_with_future(
656        &self,
657        addr: &str,
658        request_header: QueryConsumerOffsetRequestHeader,
659        timeout_millis: u64,
660    ) -> crate::ClientResult<i64> {
661        self.query_consumer_offset(addr, request_header, timeout_millis).await
662    }
663
664    pub async fn query_message(
665        this: &Arc<Self>,
666        addr: &CheetahString,
667        request_header: QueryMessageRequestHeader,
668        unique_key_flag: bool,
669        timeout_millis: u64,
670    ) -> crate::ClientResult<Option<(QueryMessageResponseHeader, Option<bytes::Bytes>)>> {
671        let mut request = this.create_request_command(RequestCode::QueryMessage, request_header);
672        if unique_key_flag {
673            request.ensure_ext_fields_initialized();
674            request.add_ext_field(mix_all::UNIQUE_MSG_QUERY_FLAG, "true");
675        }
676        let outcome = this
677            .remoting_client
678            .invoke_request(Some(addr), request, timeout_millis)
679            .await;
680        let response = match outcome {
681            Ok(OutboundRequestOutcome::Response(response)) => response,
682            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
683                return Err(consumer_request_error("query_message", RetryInput::Rejected(rejection)));
684            }
685            Ok(OutboundRequestOutcome::Contract(contract)) => {
686                return Err(consumer_request_error("query_message", RetryInput::Contract(contract)));
687            }
688            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
689        };
690        match ResponseCode::from(response.code()) {
691            ResponseCode::Success => {
692                let response_header = response
693                    .decode_command_custom_header::<QueryMessageResponseHeader>()
694                    .map_err(|e| ClientError::response_process_source("decode QueryMessageResponseHeader", e))?;
695                let body = response.body().cloned();
696                Ok(Some((response_header, body)))
697            }
698            ResponseCode::QueryNotFound => Ok(None),
699            _ => Err(client_broker_err!(
700                response.code(),
701                response.remark().map_or("".to_string(), |s| s.to_string()),
702                addr
703            )),
704        }
705    }
706
707    pub async fn pull_message<PCB>(
708        this: Arc<Self>,
709        addr: CheetahString,
710        request_header: PullMessageRequestHeader,
711        timeout_millis: u64,
712        communication_mode: CommunicationMode,
713        pull_callback: PCB,
714    ) -> crate::ClientResult<Option<PullResultExt>>
715    where
716        PCB: PullCallback + 'static,
717    {
718        let request = if PullSysFlag::has_lite_pull_flag(request_header.sys_flag as u32) {
719            this.create_request_command(RequestCode::LitePullMessage, request_header)
720        } else {
721            this.create_request_command(RequestCode::PullMessage, request_header)
722        };
723        match communication_mode {
724            CommunicationMode::Sync => {
725                let result_ext = this.pull_message_sync(&addr, request, timeout_millis).await?;
726                Ok(Some(result_ext))
727            }
728            CommunicationMode::Async => {
729                let tracker = this.background_tasks.clone();
730                let shutdown_token = this.background_shutdown.clone();
731                let service_context = this.service_context.clone();
732                Self::spawn_api_background_task(
733                    &service_context,
734                    "rocketmq-client-pull-message-async",
735                    &tracker,
736                    &shutdown_token,
737                    async move {
738                        let _ = this
739                            .pull_message_async(&addr, request, timeout_millis, pull_callback)
740                            .await;
741                    },
742                );
743                Ok(None)
744            }
745            CommunicationMode::Oneway => Ok(None),
746        }
747    }
748
749    pub(super) async fn pull_message_sync(
750        &self,
751        addr: &CheetahString,
752        request: RemotingCommand,
753        timeout_millis: u64,
754    ) -> crate::ClientResult<PullResultExt> {
755        let outcome = self
756            .remoting_client
757            .invoke_request(Some(addr), request, timeout_millis)
758            .await;
759        let response = match outcome {
760            Ok(OutboundRequestOutcome::Response(response)) => response,
761            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
762                return Err(consumer_request_error(
763                    "pull_message_sync",
764                    RetryInput::Rejected(rejection),
765                ));
766            }
767            Ok(OutboundRequestOutcome::Contract(contract)) => {
768                return Err(consumer_request_error(
769                    "pull_message_sync",
770                    RetryInput::Contract(contract),
771                ));
772            }
773            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
774        };
775        self.process_pull_response(response, addr).await
776    }
777
778    pub(super) async fn pull_message_async<PCB>(
779        &self,
780        addr: &CheetahString,
781        request: RemotingCommand,
782        timeout_millis: u64,
783        mut pull_callback: PCB,
784    ) -> crate::ClientResult<()>
785    where
786        PCB: PullCallback,
787    {
788        match self
789            .remoting_client
790            .invoke_request(Some(addr), request, timeout_millis)
791            .await
792        {
793            Ok(OutboundRequestOutcome::Response(response)) => {
794                let result = self.process_pull_response(response, addr).await;
795                let _ = self
796                    .callback_executor
797                    .execute(async {
798                        match result {
799                            Ok(pull_result) => pull_callback.on_success(pull_result).await,
800                            Err(error) => pull_callback.on_exception(error),
801                        }
802                    })
803                    .await;
804            }
805            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
806                let error = consumer_request_error("pull_message_async", RetryInput::Rejected(rejection));
807                let _ = self
808                    .callback_executor
809                    .execute(async { pull_callback.on_exception(error) })
810                    .await;
811            }
812            Ok(OutboundRequestOutcome::Contract(contract)) => {
813                let error = consumer_request_error("pull_message_async", RetryInput::Contract(contract));
814                let _ = self
815                    .callback_executor
816                    .execute(async { pull_callback.on_exception(error) })
817                    .await;
818            }
819            Err(error) => {
820                let error = ClientError::from_shared(error.into_shared_error());
821                let _ = self
822                    .callback_executor
823                    .execute(async { pull_callback.on_exception(error) })
824                    .await;
825            }
826        }
827        Ok(())
828    }
829
830    pub(super) async fn process_pull_response(
831        &self,
832        mut response: RemotingCommand,
833        addr: &CheetahString,
834    ) -> crate::ClientResult<PullResultExt> {
835        let pull_status = match ResponseCode::from(response.code()) {
836            ResponseCode::Success => PullStatus::Found,
837            ResponseCode::PullNotFound => PullStatus::NoNewMsg,
838            ResponseCode::PullRetryImmediately => PullStatus::NoMatchedMsg,
839            ResponseCode::PullOffsetMoved => PullStatus::OffsetIllegal,
840            _ => {
841                return Err(client_broker_err!(
842                    response.code(),
843                    response.remark().map_or("".to_string(), |s| s.to_string()),
844                    addr
845                ));
846            }
847        };
848        let response_header = response.decode_command_custom_header::<PullMessageResponseHeader>()?;
849        let next_begin_offset =
850            java_long_to_u64_field("pullMessage", "nextBeginOffset", response_header.next_begin_offset)?;
851        let min_offset = java_long_to_u64_field("pullMessage", "minOffset", response_header.min_offset)?;
852        let max_offset = java_long_to_u64_field("pullMessage", "maxOffset", response_header.max_offset)?;
853        let pull_result = PullResultExt {
854            pull_result: PullResult {
855                pull_status,
856                next_begin_offset,
857                min_offset,
858                max_offset,
859                msg_found_list: Some(vec![]),
860            },
861            suggest_which_broker_id: response_header.suggest_which_broker_id,
862            message_binary: response.take_body(),
863            offset_delta: response_header.offset_delta,
864        };
865        Ok(pull_result)
866    }
867
868    pub async fn consumer_send_message_back(
869        &self,
870        addr: &str,
871        broker_name: Option<&str>,
872        msg: &MessageExt,
873        consumer_group: &str,
874        delay_level: i32,
875        timeout_millis: u64,
876        max_consume_retry_times: i32,
877    ) -> crate::ClientResult<()> {
878        let header = Self::consumer_send_message_back_request_header(
879            msg,
880            broker_name,
881            consumer_group,
882            delay_level,
883            max_consume_retry_times,
884        );
885
886        let request_command = self.create_request_command(RequestCode::ConsumerSendMsgBack, header);
887        let outcome = self
888            .remoting_client
889            .invoke_request(
890                Some(&mix_all::broker_vip_channel(
891                    self.client_config.vip_channel_enabled,
892                    addr,
893                )),
894                request_command,
895                timeout_millis,
896            )
897            .await;
898        let response = match outcome {
899            Ok(OutboundRequestOutcome::Response(response)) => response,
900            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
901                return Err(consumer_request_error(
902                    "consumer_send_message_back",
903                    RetryInput::Rejected(rejection),
904                ));
905            }
906            Ok(OutboundRequestOutcome::Contract(contract)) => {
907                return Err(consumer_request_error(
908                    "consumer_send_message_back",
909                    RetryInput::Contract(contract),
910                ));
911            }
912            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
913        };
914        if ResponseCode::from(response.code()) == ResponseCode::Success {
915            Ok(())
916        } else {
917            Err(client_broker_err!(
918                response.code(),
919                response.remark().map_or("".to_string(), |s| s.to_string()),
920                addr
921            ))
922        }
923    }
924
925    pub(super) fn consumer_send_message_back_request_header(
926        msg: &MessageExt,
927        broker_name: Option<&str>,
928        consumer_group: &str,
929        delay_level: i32,
930        max_consume_retry_times: i32,
931    ) -> ConsumerSendMsgBackRequestHeader {
932        ConsumerSendMsgBackRequestHeader {
933            offset: msg.commit_log_offset,
934            group: CheetahString::from_slice(consumer_group),
935            delay_level,
936            origin_msg_id: Some(CheetahString::from_slice(msg.msg_id.as_str())),
937            origin_topic: Some(CheetahString::from_slice(msg.topic())),
938            unit_mode: false,
939            max_reconsume_times: Some(max_consume_retry_times),
940            rpc_request_header: Some(RpcRequestHeader {
941                namespace: None,
942                namespaced: None,
943                broker_name: broker_name.map(CheetahString::from_slice),
944                oneway: None,
945            }),
946        }
947    }
948
949    pub async fn send_message_back_async(
950        &self,
951        addr: &CheetahString,
952        request_header: ConsumerSendMsgBackRequestHeader,
953        timeout_millis: u64,
954    ) -> crate::ClientResult<RemotingCommand> {
955        let request = self.create_request_command(RequestCode::ConsumerSendMsgBack, request_header);
956        let outcome = self
957            .remoting_client
958            .invoke_request(Some(addr), request, timeout_millis)
959            .await;
960        match outcome {
961            Ok(OutboundRequestOutcome::Response(response)) => Ok(response),
962            Ok(OutboundRequestOutcome::Rejected(rejection)) => Err(consumer_request_error(
963                "send_message_back_async",
964                RetryInput::Rejected(rejection),
965            )),
966            Ok(OutboundRequestOutcome::Contract(contract)) => Err(consumer_request_error(
967                "send_message_back_async",
968                RetryInput::Contract(contract),
969            )),
970            Err(error) => Err(ClientError::from_shared(error.into_shared_error())),
971        }
972    }
973
974    pub async fn unregister_client(
975        &self,
976        addr: &CheetahString,
977        client_id: CheetahString,
978        producer_group: Option<CheetahString>,
979        consumer_group: Option<CheetahString>,
980        timeout_millis: u64,
981    ) -> crate::ClientResult<()> {
982        let request_header = UnregisterClientRequestHeader {
983            client_id,
984            producer_group,
985            consumer_group,
986            rpc_request_header: None,
987        };
988        let request = self.create_request_command(RequestCode::UnregisterClient, request_header);
989        let outcome = self
990            .remoting_client
991            .invoke_request(Some(addr), request, timeout_millis)
992            .await;
993        let response = match outcome {
994            Ok(OutboundRequestOutcome::Response(response)) => response,
995            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
996                return Err(consumer_request_error(
997                    "unregister_client",
998                    RetryInput::Rejected(rejection),
999                ));
1000            }
1001            Ok(OutboundRequestOutcome::Contract(contract)) => {
1002                return Err(consumer_request_error(
1003                    "unregister_client",
1004                    RetryInput::Contract(contract),
1005                ));
1006            }
1007            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
1008        };
1009        if ResponseCode::from(response.code()) == ResponseCode::Success {
1010            Ok(())
1011        } else {
1012            Err(client_broker_err!(
1013                response.code(),
1014                response.remark().map_or("".to_string(), |s| s.to_string()),
1015                addr
1016            ))
1017        }
1018    }
1019
1020    pub async fn unlock_batch_mq(
1021        &self,
1022        addr: &CheetahString,
1023        request_body: UnlockBatchRequestBody,
1024        timeout_millis: u64,
1025        oneway: bool,
1026    ) -> crate::ClientResult<()> {
1027        let mut request =
1028            self.create_request_command(RequestCode::UnlockBatchMq, UnlockBatchMqRequestHeader::default());
1029        request.set_body_mut_ref(request_body.encode()?);
1030        if oneway {
1031            self.remoting_client
1032                .invoke_request_oneway(addr, request, timeout_millis)
1033                .await
1034                .map_err(ClientError::from)
1035        } else {
1036            let outcome = self
1037                .remoting_client
1038                .invoke_request(
1039                    Some(&mix_all::broker_vip_channel(
1040                        self.client_config.vip_channel_enabled,
1041                        addr,
1042                    )),
1043                    request,
1044                    timeout_millis,
1045                )
1046                .await;
1047            let response = match outcome {
1048                Ok(OutboundRequestOutcome::Response(response)) => response,
1049                Ok(OutboundRequestOutcome::Rejected(rejection)) => {
1050                    return Err(consumer_request_error(
1051                        "unlock_batch_mq",
1052                        RetryInput::Rejected(rejection),
1053                    ));
1054                }
1055                Ok(OutboundRequestOutcome::Contract(contract)) => {
1056                    return Err(consumer_request_error(
1057                        "unlock_batch_mq",
1058                        RetryInput::Contract(contract),
1059                    ));
1060                }
1061                Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
1062            };
1063            if ResponseCode::from(response.code()) == ResponseCode::Success {
1064                Ok(())
1065            } else {
1066                Err(client_broker_err!(
1067                    response.code(),
1068                    response.remark().map_or("".to_string(), |s| s.to_string()),
1069                    addr
1070                ))
1071            }
1072        }
1073    }
1074
1075    pub async fn unlock_batch_mq_oneway(
1076        &self,
1077        addr: &CheetahString,
1078        request_body: UnlockBatchRequestBody,
1079        timeout_millis: u64,
1080    ) -> crate::ClientResult<()> {
1081        self.unlock_batch_mq(addr, request_body, timeout_millis, true).await
1082    }
1083
1084    pub async fn lock_batch_mq(
1085        &self,
1086        addr: &str,
1087        request_body: LockBatchRequestBody,
1088        timeout_millis: u64,
1089    ) -> crate::ClientResult<HashSet<MessageQueue>> {
1090        let mut request = self.create_request_command(RequestCode::LockBatchMq, LockBatchMqRequestHeader::default());
1091        request.set_body_mut_ref(request_body.encode()?);
1092        let outcome = self
1093            .remoting_client
1094            .invoke_request(
1095                Some(&mix_all::broker_vip_channel(
1096                    self.client_config.vip_channel_enabled,
1097                    addr,
1098                )),
1099                request,
1100                timeout_millis,
1101            )
1102            .await;
1103        let response = match outcome {
1104            Ok(OutboundRequestOutcome::Response(response)) => response,
1105            Ok(OutboundRequestOutcome::Rejected(rejection)) => {
1106                return Err(consumer_request_error("lock_batch_mq", RetryInput::Rejected(rejection)));
1107            }
1108            Ok(OutboundRequestOutcome::Contract(contract)) => {
1109                return Err(consumer_request_error("lock_batch_mq", RetryInput::Contract(contract)));
1110            }
1111            Err(error) => return Err(ClientError::from_shared(error.into_shared_error())),
1112        };
1113        if ResponseCode::from(response.code()) == ResponseCode::Success {
1114            if let Some(body) = response.body() {
1115                LockBatchResponseBody::decode(body.as_ref())
1116                    .map(|body| body.lock_ok_mq_set)
1117                    .map_err(|source| {
1118                        ClientError::broker_operation_source("lock_batch_mq", response.code(), Some(addr), source)
1119                    })
1120            } else {
1121                Err(client_broker_err!(response.code(), "Response body is empty", addr))
1122            }
1123        } else {
1124            Err(client_broker_err!(
1125                response.code(),
1126                response.remark().map_or("".to_string(), |s| s.to_string()),
1127                addr
1128            ))
1129        }
1130    }
1131
1132    pub async fn lock_batch_mq_with_future(
1133        &self,
1134        addr: &str,
1135        request_body: LockBatchRequestBody,
1136        timeout_millis: u64,
1137    ) -> crate::ClientResult<HashSet<MessageQueue>> {
1138        self.lock_batch_mq(addr, request_body, timeout_millis).await
1139    }
1140}
1141
1142mod message_operations;
1143
1144#[cfg(test)]
1145mod tests {
1146    use super::*;
1147
1148    #[test]
1149    fn pop_background_task_spawn_failure_preserves_cause() {
1150        let error = pop_background_task_spawn_failed(std::io::Error::other("task group closed"));
1151
1152        assert_eq!(
1153            error.descriptor().code(),
1154            rocketmq_error::CLIENT_LIFECYCLE_INVALID_STATE.code()
1155        );
1156        assert!(error.source_ref::<std::io::Error>().is_some());
1157    }
1158}
rocketmq_client_rust/implementation/mq_client_api_impl/consumer.rs

rocketmq_client_rust/implementation/mq_client_api_impl/
consumer.rs
