1use std::cmp::Ordering;
16use std::collections::HashMap;
17use std::collections::HashSet;
18use std::sync::atomic::AtomicI64;
19use std::sync::Arc;
20use std::time::Duration;
21
22use cheetah_string::CheetahString;
23use dashmap::DashMap;
24use futures::future;
25use rand::seq::IndexedRandom;
26use rocketmq_common::common::base::service_state::ServiceState;
27use rocketmq_common::common::boundary_type::BoundaryType;
28use rocketmq_common::common::config::TopicConfig;
29use rocketmq_common::common::constant::PermName;
30use rocketmq_common::common::filter::expression_type::ExpressionType;
31use rocketmq_common::common::message::message_decoder;
32use rocketmq_common::common::message::message_ext::MessageExt;
33use rocketmq_common::common::message::message_queue::MessageQueue;
34use rocketmq_common::common::message::message_queue_assignment::MessageQueueAssignment;
35use rocketmq_common::common::mix_all;
36use rocketmq_common::common::mq_version::CURRENT_VERSION;
37use rocketmq_common::TimeUtils::current_millis;
38use rocketmq_remoting::base::connection_net_event::ConnectionNetEvent;
39use rocketmq_remoting::protocol::body::acl_info::AclInfo;
40use rocketmq_remoting::protocol::body::broker_body::cluster_info::ClusterInfo;
41use rocketmq_remoting::protocol::body::consume_message_directly_result::ConsumeMessageDirectlyResult;
42use rocketmq_remoting::protocol::body::consumer_running_info::ConsumerRunningInfo;
43use rocketmq_remoting::protocol::body::user_info::UserInfo;
44use rocketmq_remoting::protocol::header::get_topic_config_request_header::GetTopicConfigRequestHeader;
45use rocketmq_remoting::protocol::header::pull_message_request_header::PullMessageRequestHeader;
46use rocketmq_remoting::protocol::header::query_consumer_offset_request_header::QueryConsumerOffsetRequestHeader;
47use rocketmq_remoting::protocol::header::update_consumer_offset_header::UpdateConsumerOffsetRequestHeader;
48use rocketmq_remoting::protocol::heartbeat::consumer_data::ConsumerData;
49use rocketmq_remoting::protocol::heartbeat::heartbeat_data::HeartbeatData;
50use rocketmq_remoting::protocol::heartbeat::message_model::MessageModel;
51use rocketmq_remoting::protocol::heartbeat::producer_data::ProducerData;
52use rocketmq_remoting::protocol::namespace_util::NamespaceUtil;
53use rocketmq_remoting::protocol::route::topic_route_data::TopicRouteData;
54use rocketmq_remoting::protocol::subscription::subscription_group_config::SubscriptionGroupConfig;
55use rocketmq_remoting::rpc::client_metadata::ClientMetadata;
56use rocketmq_remoting::runtime::config::client_config::TokioClientConfig;
57use rocketmq_remoting::runtime::RPCHook;
58use rocketmq_rust::schedule::simple_scheduler::ScheduledTaskManager;
59use rocketmq_rust::ArcMut;
60use rocketmq_rust::RocketMQTokioMutex;
61use rocketmq_rust::WeakArcMut;
62use tracing::error;
63use tracing::info;
64use tracing::warn;
65
66use crate::admin::mq_admin_ext_async_inner::MQAdminExtInnerImpl;
67use crate::base::client_config::ClientConfig;
68use crate::consumer::consumer_impl::pull_message_service::PullMessageService;
69use crate::consumer::consumer_impl::pull_request_ext::PullResultExt;
70use crate::consumer::consumer_impl::re_balance::rebalance_service::RebalanceService;
71use crate::consumer::mq_consumer_inner::MQConsumerInner;
72use crate::consumer::mq_consumer_inner::MQConsumerInnerImpl;
73use crate::consumer::pull_callback::PullCallback;
74use crate::consumer::pull_result::PullResult;
75use crate::consumer::pull_status::PullStatus;
76use crate::factory::client_tables::*;
77use crate::implementation::client_remoting_processor::ClientRemotingProcessor;
78use crate::implementation::communication_mode::CommunicationMode;
79use crate::implementation::find_broker_result::FindBrokerResult;
80use crate::implementation::mq_admin_impl::MQAdminImpl;
81use crate::implementation::mq_client_api_impl::MQClientAPIImpl;
82use crate::producer::default_mq_producer::DefaultMQProducer;
83use crate::producer::default_mq_producer::ProducerConfig;
84use crate::producer::producer_impl::mq_producer_inner::MQProducerInnerImpl;
85use crate::producer::producer_impl::topic_publish_info::TopicPublishInfo;
86use crate::stat::consumer_stats_manager::ConsumerStatsManager;
87
88const LOCK_TIMEOUT_MILLIS: u64 = 3000;
89
90fn get_topic_config_request_header(topic: CheetahString) -> GetTopicConfigRequestHeader {
91 let mut request_header = GetTopicConfigRequestHeader {
92 topic,
93 topic_request_header: None,
94 };
95 rocketmq_remoting::protocol::header::message_operation_header::TopicRequestHeaderTrait::set_lo(
96 &mut request_header,
97 Some(true),
98 );
99 request_header
100}
101
102fn update_name_server_address_list_sync(api_impl: ArcMut<MQClientAPIImpl>, addr: String) {
103 futures::executor::block_on(async move {
104 api_impl.update_name_server_address_list(&addr).await;
105 });
106}
107
108async fn run_connection_net_event_listener(
109 mut rx: tokio::sync::broadcast::Receiver<ConnectionNetEvent>,
110 weak_instance: WeakArcMut<MQClientInstance>,
111) {
112 while let Ok(value) = rx.recv().await {
113 if let Some(instance_) = weak_instance.upgrade() {
114 match value {
115 ConnectionNetEvent::CONNECTED(remote_address) => {
116 info!("ConnectionNetEvent CONNECTED");
117 let remote_address = remote_address.to_string();
118 instance_.on_channel_active(&remote_address).await;
119 }
120 ConnectionNetEvent::DISCONNECTED => {
121 instance_.on_channel_close("");
122 }
123 ConnectionNetEvent::EXCEPTION => {
124 instance_.on_channel_exception("");
125 }
126 ConnectionNetEvent::IDLE => {
127 instance_.on_channel_idle("");
128 }
129 }
130 }
131 }
132 warn!("ConnectionNetEvent recv error");
133}
134
135fn spawn_connection_net_event_listener(
136 rx: tokio::sync::broadcast::Receiver<ConnectionNetEvent>,
137 weak_instance: WeakArcMut<MQClientInstance>,
138) {
139 match tokio::runtime::Handle::try_current() {
140 Ok(handle) => {
141 let task = handle.spawn(run_connection_net_event_listener(rx, weak_instance));
142 drop(task);
143 }
144 Err(error) => {
145 warn!(
146 "No Tokio runtime available while creating MQClientInstance; starting dedicated connection event \
147 listener thread. error={}",
148 error
149 );
150 let thread = std::thread::Builder::new().name("rocketmq-client-connection-events".to_string());
151 if let Err(error) = thread.spawn(move || {
152 let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
153 Ok(runtime) => runtime,
154 Err(error) => {
155 error!("Failed to create MQClientInstance connection event runtime: {}", error);
156 return;
157 }
158 };
159 runtime.block_on(run_connection_net_event_listener(rx, weak_instance));
160 }) {
161 error!(
162 "Failed to spawn MQClientInstance connection event listener thread: {}",
163 error
164 );
165 }
166 }
167 }
168}
169
170fn schedule_rebalance_wakeup(service: RebalanceService, delay: Duration) {
171 if delay.is_zero() {
172 service.wakeup();
173 return;
174 }
175
176 match tokio::runtime::Handle::try_current() {
177 Ok(handle) => {
178 let task = handle.spawn(async move {
179 tokio::time::sleep(delay).await;
180 service.wakeup();
181 });
182 drop(task);
183 }
184 Err(error) => {
185 warn!(
186 "No Tokio runtime available for delayed rebalance wakeup; using a dedicated sleeper thread. error={}",
187 error
188 );
189 let thread = std::thread::Builder::new().name("rocketmq-client-rebalance-delay".to_string());
190 if let Err(error) = thread.spawn(move || {
191 std::thread::sleep(delay);
192 service.wakeup();
193 }) {
194 error!("Failed to spawn delayed rebalance wakeup thread: {}", error);
195 }
196 }
197 }
198}
199
200pub struct MQClientInstance {
201 pub(crate) client_config: ArcMut<ClientConfig>,
202 pub(crate) client_id: CheetahString,
203 boot_timestamp: u64,
204 producer_table: ProducerTable,
209 consumer_table: ConsumerTable,
214 admin_ext_table: AdminExtTable,
219 pub(crate) mq_client_api_impl: Option<ArcMut<MQClientAPIImpl>>,
220 pub(crate) mq_admin_impl: ArcMut<MQAdminImpl>,
221 pub(crate) topic_route_table: TopicRouteTable,
222 topic_end_points_table: TopicEndPointsTable,
223 lock_namesrv: Arc<RocketMQTokioMutex<()>>,
224 lock_heartbeat: Arc<RocketMQTokioMutex<()>>,
225
226 service_state: ServiceState,
227 pub(crate) pull_message_service: ArcMut<PullMessageService>,
228 rebalance_service: RebalanceService,
229 pub(crate) default_producer: ArcMut<DefaultMQProducer>,
230 broker_addr_table: BrokerAddrTable,
231 broker_version_table: BrokerVersionTable,
232 send_heartbeat_times_total: Arc<AtomicI64>,
233 scheduled_task_manager: ScheduledTaskManager,
234 broker_heartbeat_fingerprint_table: BrokerHeartbeatFingerprintTable,
236 broker_support_v2_heartbeat_set: BrokerSupportV2HeartbeatSet,
238 consumer_stats_manager: ConsumerStatsManager,
239}
240
241impl MQClientInstance {
242 pub fn new_arc(
243 client_config: ClientConfig,
244 _instance_index: i32,
245 client_id: impl Into<CheetahString>,
246 rpc_hook: Option<Arc<dyn RPCHook>>,
247 ) -> ArcMut<MQClientInstance> {
248 let client_id = client_id.into();
249 let shared_config = ArcMut::new(client_config.clone());
250
251 let broker_addr_table = Arc::new(DashMap::default());
252 let producer_table = Arc::new(DashMap::default());
253 let consumer_table = Arc::new(DashMap::default());
254 let admin_ext_table = Arc::new(DashMap::default());
255 let topic_route_table = Arc::new(DashMap::default());
256 let topic_end_points_table = Arc::new(DashMap::default());
257 let broker_version_table = Arc::new(DashMap::default());
258 let broker_heartbeat_fingerprint_table = Arc::new(DashMap::default());
259 let broker_support_v2_heartbeat_set = Arc::new(DashMap::default());
260
261 let default_producer = ArcMut::new(
262 DefaultMQProducer::builder()
263 .producer_group(mix_all::CLIENT_INNER_PRODUCER_GROUP)
264 .client_config(client_config.clone())
265 .build(),
266 );
267 let mut instance = ArcMut::new(MQClientInstance {
268 client_config: shared_config,
269 client_id,
270 boot_timestamp: current_millis(),
271 producer_table,
272 consumer_table,
273 admin_ext_table,
274 mq_client_api_impl: None,
275 mq_admin_impl: ArcMut::new(MQAdminImpl::new()),
276 topic_route_table,
277 topic_end_points_table,
278 lock_namesrv: Arc::default(),
279 lock_heartbeat: Arc::default(),
280 service_state: ServiceState::CreateJust,
281 pull_message_service: ArcMut::new(PullMessageService::new()),
282 rebalance_service: RebalanceService::new(),
283 default_producer,
284 broker_addr_table,
285 broker_version_table,
286 send_heartbeat_times_total: Arc::new(AtomicI64::new(0)),
287 scheduled_task_manager: ScheduledTaskManager::new(),
288 broker_heartbeat_fingerprint_table,
289 broker_support_v2_heartbeat_set,
290 consumer_stats_manager: ConsumerStatsManager::new(),
291 });
292
293 let instance_clone = instance.clone();
295 instance.mq_admin_impl.set_client(instance_clone);
296
297 let (tx, rx) = tokio::sync::broadcast::channel::<ConnectionNetEvent>(16);
298
299 let mq_client_api_impl = ArcMut::new(MQClientAPIImpl::new(
300 Arc::new(TokioClientConfig::default()),
301 ClientRemotingProcessor::new(instance.clone()),
302 rpc_hook,
303 Arc::new(client_config.clone()),
304 Some(tx),
305 ));
306
307 if let Some(namesrv_addr) = client_config.namesrv_addr.as_deref() {
308 let api_impl = mq_client_api_impl.clone();
309 let addr = namesrv_addr.to_string();
310 update_name_server_address_list_sync(api_impl, addr);
311 }
312
313 instance.mq_client_api_impl = Some(mq_client_api_impl);
314
315 let weak_instance = ArcMut::downgrade(&instance);
317 spawn_connection_net_event_listener(rx, weak_instance);
318 instance
319 }
320
321 pub fn consumer_stats_manager(&self) -> &ConsumerStatsManager {
323 &self.consumer_stats_manager
324 }
325
326 pub fn on_channel_connect(&self, _remote_addr: &str) {}
327
328 pub fn on_channel_close(&self, _remote_addr: &str) {}
329
330 pub fn on_channel_exception(&self, _remote_addr: &str) {}
331
332 pub fn on_channel_idle(&self, _remote_addr: &str) {}
333
334 pub async fn on_channel_active(&self, remote_addr: &str) {
335 let matched_brokers: Vec<_> = self
336 .broker_addr_table
337 .iter()
338 .flat_map(|entry| {
339 let (name, addrs) = entry.pair();
340 addrs
341 .iter()
342 .filter(|(_, addr)| addr.as_str() == remote_addr)
343 .map(|(id, addr)| (*id, name.clone(), addr.clone()))
344 .collect::<Vec<_>>()
345 })
346 .collect();
347
348 for (id, broker_name, addr) in matched_brokers {
349 if self.send_heartbeat_to_broker(id, &broker_name, &addr, false).await {
350 self.re_balance_immediately();
351 }
352 }
353 }
354
355 pub fn re_balance_immediately(&self) {
356 self.rebalance_service.wakeup();
357 }
358
359 pub fn rebalance_immediately(&self) {
360 self.re_balance_immediately();
361 }
362
363 pub fn re_balance_later(&self, delay_millis: Duration) {
364 schedule_rebalance_wakeup(self.rebalance_service.clone(), delay_millis);
365 }
366
367 pub async fn start(&mut self, this: ArcMut<Self>) -> rocketmq_error::RocketMQResult<()> {
368 match self.service_state {
369 ServiceState::CreateJust => {
370 self.service_state = ServiceState::StartFailed;
371 if self.client_config.namesrv_addr.is_none() {
373 let Some(mq_client_api_impl) = self.mq_client_api_impl.as_mut() else {
374 return Err(mq_client_err!("mq_client_api_impl is None"));
375 };
376 mq_client_api_impl.fetch_name_server_addr().await;
377 }
378 let Some(mq_client_api_impl) = self.mq_client_api_impl.as_mut() else {
380 return Err(mq_client_err!("mq_client_api_impl is None"));
381 };
382 mq_client_api_impl.start().await;
383 self.start_scheduled_task(this.clone());
385 let instance = this.clone();
387 if let Err(e) = self.pull_message_service.start(instance).await {
388 error!("Failed to start pull message service: {:?}", e);
389 }
390 if let Err(e) = self.rebalance_service.start(this).await {
392 error!("Failed to start rebalance service: {:?}", e);
393 }
394 let Some(default_producer_impl) = self.default_producer.default_mqproducer_impl.as_mut() else {
397 return Err(mq_client_err!("default producer impl is None"));
398 };
399 default_producer_impl.start_with_factory(false).await?;
400 self.consumer_stats_manager.start();
402 info!("the client factory[{}] start OK", self.client_id);
403 self.service_state = ServiceState::Running;
404 }
405 ServiceState::Running => {}
406 ServiceState::ShutdownAlready => {}
407 ServiceState::StartFailed => {
408 return Err(mq_client_err!(format!(
409 "The Factory object[{}] has been created before, and failed.",
410 self.client_id
411 )));
412 }
413 }
414 Ok(())
415 }
416
417 pub async fn shutdown(&mut self) {
418 match self.service_state {
419 ServiceState::CreateJust | ServiceState::ShutdownAlready | ServiceState::StartFailed => {
420 warn!(
421 "MQClientInstance shutdown called but state is {:?}, ignoring shutdown request",
422 self.service_state
423 );
424 return;
425 }
426 ServiceState::Running => {
427 info!(
428 "MQClientInstance[{}] shutdown starting, current state: Running",
429 self.client_id
430 );
431 }
432 }
433
434 info!("MQClientInstance[{}] shutting down rebalance service", self.client_id);
435 if let Err(e) = self.rebalance_service.shutdown(3000).await {
436 warn!("Failed to shutdown rebalance service: {:?}", e);
437 }
438
439 info!(
440 "MQClientInstance[{}] shutting down pull message service",
441 self.client_id
442 );
443 if let Err(e) = self.pull_message_service.shutdown_default().await {
444 warn!("Failed to shutdown pull message service: {:?}", e);
445 }
446
447 info!("MQClientInstance[{}] persisting all consumer offsets", self.client_id);
448 self.persist_all_consumer_offset().await;
449
450 info!("MQClientInstance[{}] shutting down default producer", self.client_id);
451 if let Some(producer_impl) = self.default_producer.default_mqproducer_impl.as_mut() {
452 let shutdown_future = producer_impl.shutdown_with_factory(false);
453 if let Err(e) = Box::pin(shutdown_future).await {
454 warn!("Failed to shutdown default producer: {:?}", e);
455 }
456 }
457
458 info!("MQClientInstance[{}] unregistering all consumers", self.client_id);
459 self.unregister_all_consumers().await;
460
461 info!("MQClientInstance[{}] unregistering all producers", self.client_id);
462 self.unregister_all_producers().await;
463
464 if self.mq_client_api_impl.is_some() {
465 info!(
466 "MQClientInstance[{}] network client shutdown (deferred to Drop)",
467 self.client_id
468 );
469 }
470
471 info!("MQClientInstance[{}] canceling scheduled tasks", self.client_id);
472 self.scheduled_task_manager.cancel_all();
473 info!(
474 "MQClientInstance[{}] scheduled tasks canceled, total canceled: {}",
475 self.client_id,
476 self.scheduled_task_manager.task_count()
477 );
478
479 info!("MQClientInstance[{}] clearing all registration tables", self.client_id);
480 self.producer_table.clear();
481 self.consumer_table.clear();
482 self.admin_ext_table.clear();
483
484 self.topic_route_table.clear();
485 self.topic_end_points_table.clear();
486 self.broker_addr_table.clear();
487 self.broker_version_table.clear();
488
489 self.service_state = ServiceState::ShutdownAlready;
490
491 info!("MQClientInstance[{}] shutdown completed successfully", self.client_id);
492 }
493
494 async fn unregister_all_producers(&mut self) {
496 let producer_groups: Vec<CheetahString> = self.producer_table.iter().map(|entry| entry.key().clone()).collect();
498
499 for group in producer_groups {
501 info!("Unregistering producer group: {}", group);
502 self.unregister_producer(group).await;
503 }
504 }
505
506 async fn unregister_all_consumers(&mut self) {
508 let consumer_groups: Vec<CheetahString> = self.consumer_table.iter().map(|entry| entry.key().clone()).collect();
510
511 for group in consumer_groups {
513 info!("Unregistering consumer group: {}", group);
514 self.unregister_consumer(group).await;
515 }
516 }
517
518 pub async fn register_producer(&mut self, group: &str, producer: MQProducerInnerImpl) -> bool {
519 if group.is_empty() {
520 return false;
521 }
522 if self.producer_table.contains_key(group) {
523 warn!("the producer group[{}] exist already.", group);
524 return false;
525 }
526 self.producer_table.insert(group.into(), producer);
527 true
528 }
529
530 pub async fn register_admin_ext(&mut self, group: &str, admin: MQAdminExtInnerImpl) -> bool {
531 if group.is_empty() {
532 return false;
533 }
534 if self.admin_ext_table.contains_key(group) {
535 warn!("the admin group[{}] exist already.", group);
536 return false;
537 }
538 self.admin_ext_table.insert(group.into(), admin);
539 true
540 }
541
542 fn start_scheduled_task(&mut self, this: ArcMut<Self>) {
543 info!("Starting scheduled tasks with ScheduledTaskManager");
544
545 if self.client_config.namesrv_addr.is_none() {
546 if let Some(mq_client_api_impl) = self.mq_client_api_impl.as_ref().cloned() {
547 self.scheduled_task_manager.add_fixed_rate_task_async(
548 Duration::from_secs(10),
549 Duration::from_secs(120),
550 async move |_token| {
551 let mut api = mq_client_api_impl.clone();
552 info!("ScheduledTask: fetchNameServerAddr");
553 api.fetch_name_server_addr().await;
554 Ok(())
555 },
556 );
557 } else {
558 warn!(
559 "ScheduledTask: fetchNameServerAddr skipped because mq_client_api_impl is None. [{}]",
560 self.client_id
561 );
562 }
563 }
564
565 let client_instance = this.clone();
566 let poll_name_server_interval = self.client_config.poll_name_server_interval;
567 self.scheduled_task_manager.add_fixed_rate_task_async(
568 Duration::from_millis(10),
569 Duration::from_millis(poll_name_server_interval as u64),
570 async move |_token| {
571 let mut instance = client_instance.clone();
572 info!("ScheduledTask: update_topic_route_info_from_name_server");
573 instance.update_topic_route_info_from_name_server().await;
574 Ok(())
575 },
576 );
577
578 let client_instance = this.clone();
579 let heartbeat_broker_interval = self.client_config.heartbeat_broker_interval;
580 self.scheduled_task_manager.add_fixed_rate_task_async(
581 Duration::from_secs(1),
582 Duration::from_millis(heartbeat_broker_interval as u64),
583 async move |_token| {
584 let mut instance = client_instance.clone();
585 info!("ScheduledTask: clean_offline_broker and send_heartbeat");
586 instance.clean_offline_broker().await;
587 instance.send_heartbeat_to_all_broker_with_lock().await;
588 Ok(())
589 },
590 );
591
592 let client_instance = this;
593 let persist_consumer_offset_interval = self.client_config.persist_consumer_offset_interval as u64;
594 self.scheduled_task_manager.add_fixed_rate_task_async(
595 Duration::from_secs(10),
596 Duration::from_millis(persist_consumer_offset_interval),
597 async move |_token| {
598 let mut instance = client_instance.clone();
599 info!("ScheduledTask: persistAllConsumerOffset");
600 instance.persist_all_consumer_offset().await;
601 Ok(())
602 },
603 );
604
605 info!(
606 "All scheduled tasks started, total tasks: {}",
607 self.scheduled_task_manager.task_count()
608 );
609 }
610
611 pub async fn update_topic_route_info_from_name_server(&mut self) {
612 let mut topic_list = HashSet::new();
613
614 for entry in self.producer_table.iter() {
615 topic_list.extend(entry.value().get_publish_topic_list());
616 }
617
618 for entry in self.consumer_table.iter() {
619 entry.value().subscriptions().iter().for_each(|sub| {
620 topic_list.insert(sub.topic.clone());
621 });
622 }
623
624 for topic in topic_list.iter() {
625 self.update_topic_route_info_from_name_server_topic(topic).await;
626 }
627 }
628
629 #[inline]
630 pub async fn update_topic_route_info_from_name_server_topic(&mut self, topic: &CheetahString) -> bool {
631 self.update_topic_route_info_from_name_server_default(topic, false, None)
632 .await
633 }
634
635 pub async fn find_consumer_id_list(
636 &mut self,
637 topic: &CheetahString,
638 group: &CheetahString,
639 ) -> Option<Vec<CheetahString>> {
640 let mut broker_addr = self.find_broker_addr_by_topic(topic).await;
641 if broker_addr.is_none() {
642 self.update_topic_route_info_from_name_server_topic(topic).await;
643 broker_addr = self.find_broker_addr_by_topic(topic).await;
644 }
645 if let Some(broker_addr) = broker_addr {
646 if let Some(mq_client_api_impl) = self.mq_client_api_impl.as_mut() {
647 match mq_client_api_impl
648 .get_consumer_id_list_by_group(
649 broker_addr.as_str(),
650 group,
651 self.client_config.mq_client_api_timeout,
652 )
653 .await
654 {
655 Ok(value) => return Some(value),
656 Err(e) => {
657 warn!(
658 "getConsumerIdListByGroup exception,{} {}, err:{}",
659 broker_addr,
660 group,
661 e.to_string()
662 );
663 }
664 }
665 } else {
666 warn!(
667 "getConsumerIdListByGroup skipped because mq_client_api_impl is None, topic={}, group={}",
668 topic, group
669 );
670 }
671 }
672 None
673 }
674
675 pub async fn find_broker_addr_by_topic(&self, topic: &str) -> Option<CheetahString> {
676 if let Some(topic_route_data) = self.topic_route_table.get(topic) {
677 let brokers = &topic_route_data.value().broker_datas;
678 if !brokers.is_empty() {
679 let bd = brokers.choose(&mut rand::rng());
680 if let Some(bd) = bd {
681 return bd.select_broker_addr();
682 }
683 }
684 }
685 None
686 }
687
688 pub async fn query_topic_route_data(&mut self, topic: &CheetahString) -> Option<TopicRouteData> {
689 if let Some(topic_route_data) = self.topic_route_table.get(topic) {
690 return Some(TopicRouteData::from_existing(topic_route_data.value()));
691 }
692
693 self.update_topic_route_info_from_name_server_topic(topic).await;
694 self.topic_route_table
695 .get(topic)
696 .map(|entry| TopicRouteData::from_existing(entry.value()))
697 }
698
699 pub async fn update_topic_route_info_from_name_server_default(
700 &mut self,
701 topic: &CheetahString,
702 is_default: bool,
703 producer_config: Option<&Arc<ProducerConfig>>,
704 ) -> bool {
705 let lock = self.lock_namesrv.lock().await;
706 let Some(mq_client_api_impl) = self.mq_client_api_impl.as_mut() else {
707 error!(
708 "updateTopicRouteInfoFromNameServer skipped because mq_client_api_impl is None, Topic: {}. [{}]",
709 topic, self.client_id
710 );
711 drop(lock);
712 return false;
713 };
714 let topic_route_data = if let (true, Some(producer_config)) = (is_default, producer_config) {
715 let mut result = match mq_client_api_impl
716 .get_default_topic_route_info_from_name_server(self.client_config.mq_client_api_timeout)
717 .await
718 {
719 Ok(value) => value,
720 Err(e) => {
721 error!(
722 "get_default_topic_route_info_from_name_server failed, topic: {}, error: {}",
723 topic, e
724 );
725 None
726 }
727 };
728 if let Some(topic_route_data) = result.as_mut() {
729 for data in topic_route_data.queue_datas.iter_mut() {
730 let queue_nums = producer_config.default_topic_queue_nums().max(data.read_queue_nums);
731 data.read_queue_nums = queue_nums;
732 data.write_queue_nums = queue_nums;
733 }
734 }
735 result
736 } else {
737 mq_client_api_impl
738 .get_topic_route_info_from_name_server(topic, self.client_config.mq_client_api_timeout)
739 .await
740 .unwrap_or(None)
741 };
742 if let Some(mut topic_route_data) = topic_route_data {
743 let old = self.topic_route_table.get(topic).map(|entry| entry.value().clone());
744 let mut changed = topic_route_data.topic_route_data_changed(old.as_ref());
745 if !changed {
746 changed = self.is_need_update_topic_route_info(topic).await;
747 } else {
748 info!(
749 "the topic[{}] route info changed, old[{:?}] ,new[{:?}]",
750 topic, old, topic_route_data
751 )
752 }
753 if changed {
754 for bd in topic_route_data.broker_datas.iter() {
755 self.broker_addr_table
756 .insert(bd.broker_name().clone(), bd.broker_addrs().clone());
757 }
758
759 {
761 let mq_end_points =
762 ClientMetadata::topic_route_data2endpoints_for_static_topic(topic, &topic_route_data);
763 if let Some(mq_end_points) = mq_end_points {
764 if !mq_end_points.is_empty() {
765 self.topic_end_points_table.insert(topic.into(), mq_end_points);
766 }
767 }
768 }
769
770 {
772 let mut publish_info = topic_route_data2topic_publish_info(topic, &mut topic_route_data);
773 publish_info.have_topic_router_info = true;
774 for mut entry in self.producer_table.iter_mut() {
775 entry
776 .value_mut()
777 .update_topic_publish_info(topic.to_string(), Some(publish_info.clone()));
778 }
779 }
780
781 if !self.consumer_table.is_empty() {
783 let subscribe_info = topic_route_data2topic_subscribe_info(topic, &topic_route_data);
784
785 let consumers: Vec<_> = self.consumer_table.iter().map(|entry| entry.value().clone()).collect();
786
787 for consumer in consumers {
788 consumer
789 .update_topic_subscribe_info(topic.clone(), &subscribe_info)
790 .await;
791 }
792 }
793 let clone_topic_route_data = TopicRouteData::from_existing(&topic_route_data);
794 self.topic_route_table.insert(topic.clone(), clone_topic_route_data);
795 return true;
796 }
797 } else {
798 warn!(
799 "updateTopicRouteInfoFromNameServer, getTopicRouteInfoFromNameServer return null, Topic: {}. [{}]",
800 topic, self.client_id
801 );
802 }
803
804 drop(lock);
805 false
806 }
807
808 async fn is_need_update_topic_route_info(&self, topic: &CheetahString) -> bool {
809 for entry in self.producer_table.iter() {
810 if entry.value().is_publish_topic_need_update(topic) {
811 return true;
812 }
813 }
814
815 for entry in self.consumer_table.iter() {
816 if entry.value().is_subscribe_topic_need_update(topic).await {
817 return true;
818 }
819 }
820 false
821 }
822
823 pub async fn persist_all_consumer_offset(&mut self) {
824 for entry in self.consumer_table.iter() {
825 entry.value().persist_consumer_offset().await;
826 }
827 }
828
829 pub async fn clean_offline_broker(&mut self) {
830 let lock = self
831 .lock_namesrv
832 .try_lock_timeout(Duration::from_millis(LOCK_TIMEOUT_MILLIS))
833 .await;
834 if let Some(lock) = lock {
835 let mut updated_table = HashMap::new();
836 let mut broker_name_set = HashSet::new();
837
838 for broker_entry in self.broker_addr_table.iter() {
839 let (broker_name, one_table) = broker_entry.pair();
840 let mut clone_addr_table = one_table.clone();
841 let mut remove_id_set = HashSet::new();
842 for (id, addr) in one_table.iter() {
843 if !self.is_broker_addr_exist_in_topic_route_table(addr).await {
844 remove_id_set.insert(*id);
845 }
846 }
847 clone_addr_table.retain(|k, _| !remove_id_set.contains(k));
848 if clone_addr_table.is_empty() {
849 info!("the broker[{}] name's host is offline, remove it", broker_name);
850 broker_name_set.insert(broker_name.clone());
851 } else {
852 updated_table.insert(broker_name.clone(), clone_addr_table);
853 }
854 }
855
856 for broker_name in broker_name_set {
858 self.broker_addr_table.remove(&broker_name);
859 }
860 for (broker_name, addrs) in updated_table {
861 self.broker_addr_table.insert(broker_name, addrs);
862 }
863 }
864 }
865 pub async fn send_heartbeat_to_all_broker_with_lock(&self) -> bool {
866 let _guard = match self.lock_heartbeat.try_lock().await {
867 Some(g) => g,
868 None => {
869 warn!("lock heartBeat, but failed. [{}]", self.client_id);
870 return false;
871 }
872 };
873
874 if self.client_config.enable_concurrent_heartbeat {
876 return self.send_heartbeat_to_all_broker_concurrently().await;
877 }
878
879 if self.client_config.use_heartbeat_v2 {
881 self.send_heartbeat_to_all_broker_v2(false).await
882 } else {
883 self.send_heartbeat_to_all_broker().await
884 }
885 }
886
887 pub async fn send_heartbeat_to_all_broker_with_lock_v2(&self, is_rebalance: bool) -> bool {
888 let _guard = match self.lock_heartbeat.try_lock_timeout(Duration::from_secs(2)).await {
889 Some(g) => g,
890 None => {
891 warn!("lock heartBeat, but failed. [{}]", self.client_id);
892 return false;
893 }
894 };
895
896 if self.client_config.use_heartbeat_v2 {
897 self.send_heartbeat_to_all_broker_v2(is_rebalance).await
898 } else {
899 self.send_heartbeat_to_all_broker().await
900 }
901 }
902
903 pub fn get_mq_client_api_impl(&self) -> rocketmq_error::RocketMQResult<ArcMut<MQClientAPIImpl>> {
904 self.mq_client_api_impl
905 .as_ref()
906 .cloned()
907 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)
908 }
909
910 pub async fn pull_message_from_broker(
911 &self,
912 broker_addr: &str,
913 request_header: PullMessageRequestHeader,
914 timeout_millis: u64,
915 ) -> rocketmq_error::RocketMQResult<PullResult> {
916 struct NoopPullCallback;
917
918 impl PullCallback for NoopPullCallback {
919 async fn on_success(&mut self, _pull_result: PullResultExt) {}
920
921 fn on_exception(&mut self, _e: Box<dyn std::error::Error + Send>) {}
922 }
923
924 let api_impl = self
925 .mq_client_api_impl
926 .as_ref()
927 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
928
929 let mut result = MQClientAPIImpl::pull_message(
930 api_impl.clone(),
931 CheetahString::from(broker_addr),
932 request_header,
933 timeout_millis,
934 CommunicationMode::Sync,
935 NoopPullCallback,
936 )
937 .await?
938 .ok_or_else(|| rocketmq_error::RocketMQError::Internal("pull_message returned None in sync mode".into()))?;
939
940 if result.pull_result.pull_status == PullStatus::Found {
941 if let Some(mut message_binary) = result.message_binary.take() {
942 let messages = message_decoder::decodes_batch(&mut message_binary, true, true);
943 result
944 .pull_result
945 .set_msg_found_list(Some(messages.into_iter().map(ArcMut::new).collect()));
946 }
947 }
948
949 Ok(result.pull_result)
950 }
951
952 pub async fn query_consumer_offset(
953 &self,
954 broker_addr: &str,
955 request_header: QueryConsumerOffsetRequestHeader,
956 timeout_millis: u64,
957 ) -> rocketmq_error::RocketMQResult<i64> {
958 let api_impl = self
959 .mq_client_api_impl
960 .as_ref()
961 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
962 api_impl
963 .clone()
964 .query_consumer_offset(broker_addr, request_header, timeout_millis)
965 .await
966 }
967
968 pub async fn update_consumer_offset(
969 &self,
970 broker_addr: &CheetahString,
971 request_header: UpdateConsumerOffsetRequestHeader,
972 timeout_millis: u64,
973 ) -> rocketmq_error::RocketMQResult<()> {
974 let api_impl = self
975 .mq_client_api_impl
976 .as_ref()
977 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
978 api_impl
979 .clone()
980 .update_consumer_offset(broker_addr, request_header, timeout_millis)
981 .await
982 }
983
984 pub async fn consumer_send_message_back(
985 &mut self,
986 broker_addr: &str,
987 broker_name: Option<&str>,
988 message: &MessageExt,
989 consumer_group: &str,
990 delay_level: i32,
991 timeout_millis: u64,
992 max_consume_retry_times: i32,
993 ) -> rocketmq_error::RocketMQResult<()> {
994 let api_impl = self
995 .mq_client_api_impl
996 .as_ref()
997 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
998 api_impl
999 .clone()
1000 .consumer_send_message_back(
1001 broker_addr,
1002 broker_name,
1003 message,
1004 consumer_group,
1005 delay_level,
1006 timeout_millis,
1007 max_consume_retry_times,
1008 )
1009 .await
1010 }
1011
1012 pub async fn get_max_offset(
1013 &self,
1014 broker_addr: &str,
1015 message_queue: &MessageQueue,
1016 timeout_millis: u64,
1017 ) -> rocketmq_error::RocketMQResult<i64> {
1018 let api_impl = self
1019 .mq_client_api_impl
1020 .as_ref()
1021 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1022 api_impl
1023 .clone()
1024 .get_max_offset(broker_addr, message_queue, timeout_millis)
1025 .await
1026 }
1027
1028 pub async fn get_min_offset(
1029 &self,
1030 broker_addr: &str,
1031 message_queue: &MessageQueue,
1032 timeout_millis: u64,
1033 ) -> rocketmq_error::RocketMQResult<i64> {
1034 let api_impl = self
1035 .mq_client_api_impl
1036 .as_ref()
1037 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1038 api_impl
1039 .clone()
1040 .get_min_offset(broker_addr, message_queue, timeout_millis)
1041 .await
1042 }
1043
1044 pub async fn search_offset_by_timestamp(
1045 &self,
1046 broker_addr: &str,
1047 message_queue: &MessageQueue,
1048 timestamp: i64,
1049 boundary_type: BoundaryType,
1050 timeout_millis: u64,
1051 ) -> rocketmq_error::RocketMQResult<i64> {
1052 let api_impl = self
1053 .mq_client_api_impl
1054 .as_ref()
1055 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1056 api_impl
1057 .clone()
1058 .search_offset_by_timestamp(broker_addr, message_queue, timestamp, boundary_type, timeout_millis)
1059 .await
1060 }
1061
1062 pub async fn get_topic_config(
1063 &self,
1064 broker_addr: &CheetahString,
1065 topic: CheetahString,
1066 timeout_millis: u64,
1067 ) -> rocketmq_error::RocketMQResult<TopicConfig> {
1068 let request_header = get_topic_config_request_header(topic);
1069 let api_impl = self
1070 .mq_client_api_impl
1071 .as_ref()
1072 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1073 let topic_mapping = api_impl
1074 .get_topic_config(broker_addr, request_header, timeout_millis)
1075 .await?;
1076 Ok(topic_mapping.topic_config)
1077 }
1078
1079 pub async fn get_subscription_group_config(
1080 &self,
1081 broker_addr: &CheetahString,
1082 group: CheetahString,
1083 timeout_millis: u64,
1084 ) -> rocketmq_error::RocketMQResult<SubscriptionGroupConfig> {
1085 let api_impl = self
1086 .mq_client_api_impl
1087 .as_ref()
1088 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1089 api_impl
1090 .get_subscription_group_config(broker_addr, group, timeout_millis)
1091 .await
1092 }
1093
1094 pub async fn get_broker_cluster_info(&self, timeout_millis: u64) -> rocketmq_error::RocketMQResult<ClusterInfo> {
1095 let api_impl = self
1096 .mq_client_api_impl
1097 .as_ref()
1098 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1099 api_impl.get_broker_cluster_info(timeout_millis).await
1100 }
1101
1102 pub async fn get_user(
1103 &self,
1104 broker_addr: CheetahString,
1105 username: CheetahString,
1106 timeout_millis: u64,
1107 ) -> rocketmq_error::RocketMQResult<Option<UserInfo>> {
1108 let api_impl = self
1109 .mq_client_api_impl
1110 .as_ref()
1111 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1112 api_impl.get_user(broker_addr, username, timeout_millis).await
1113 }
1114
1115 pub async fn get_acl(
1116 &self,
1117 broker_addr: CheetahString,
1118 subject: CheetahString,
1119 timeout_millis: u64,
1120 ) -> rocketmq_error::RocketMQResult<Option<AclInfo>> {
1121 let api_impl = self
1122 .mq_client_api_impl
1123 .as_ref()
1124 .ok_or(rocketmq_error::RocketMQError::ClientNotStarted)?;
1125 let acl_infos = api_impl
1126 .list_acl(broker_addr, subject.clone(), CheetahString::default(), timeout_millis)
1127 .await?;
1128 Ok(acl_infos
1129 .into_iter()
1130 .find(|acl_info| acl_info.subject.as_ref() == Some(&subject)))
1131 }
1132
1133 pub async fn get_broker_name_from_message_queue(&self, message_queue: &MessageQueue) -> CheetahString {
1134 if let Some(broker_name) = self.topic_end_points_table.get(message_queue.topic_str()) {
1135 if let Some(addr) = broker_name.value().get(message_queue) {
1136 return addr.clone();
1137 }
1138 }
1139 message_queue.broker_name().clone()
1140 }
1141
1142 pub fn find_broker_address_in_publish(&self, broker_name: &CheetahString) -> Option<CheetahString> {
1143 if broker_name.is_empty() {
1144 return None;
1145 }
1146 let map = self.broker_addr_table.get(broker_name);
1147 if let Some(map) = map {
1148 return map.get(&(mix_all::MASTER_ID)).cloned();
1149 }
1150 None
1151 }
1152
1153 pub async fn find_broker_address_in_admin(&self, broker_name: &CheetahString) -> Option<FindBrokerResult> {
1154 if broker_name.is_empty() {
1155 return None;
1156 }
1157
1158 let map_ref = self.broker_addr_table.get(broker_name)?;
1159 for (id, broker_addr) in map_ref.value() {
1160 if broker_addr.is_empty() {
1161 continue;
1162 }
1163 let broker_addr = broker_addr.clone();
1164 let broker_version = self.find_broker_version(broker_name, broker_addr.as_str()).await;
1165 return Some(FindBrokerResult {
1166 broker_addr,
1167 slave: *id != mix_all::MASTER_ID,
1168 broker_version,
1169 });
1170 }
1171 None
1172 }
1173
1174 pub fn parse_offset_table_from_broker(
1175 &self,
1176 offset_table: HashMap<MessageQueue, i64>,
1177 namespace: &str,
1178 ) -> HashMap<MessageQueue, i64> {
1179 if namespace.is_empty() {
1180 return offset_table;
1181 }
1182
1183 offset_table
1184 .into_iter()
1185 .map(|(mut queue, offset)| {
1186 let topic = NamespaceUtil::without_namespace_with_namespace(queue.topic_str(), namespace);
1187 queue.set_topic(CheetahString::from_string(topic));
1188 (queue, offset)
1189 })
1190 .collect()
1191 }
1192
1193 async fn send_heartbeat_to_all_broker_v2(&self, is_rebalance: bool) -> bool {
1194 let heartbeat_data_with_sub = self.prepare_heartbeat_data(false).await;
1195 let producer_empty = heartbeat_data_with_sub.producer_data_set.is_empty();
1196 let consumer_empty = heartbeat_data_with_sub.consumer_data_set.is_empty();
1197
1198 if producer_empty && consumer_empty {
1199 warn!(
1200 "sendHeartbeatToAllBrokerV2 sending heartbeat, but no consumer and no producer. [{}]",
1201 self.client_id
1202 );
1203 return false;
1204 }
1205
1206 if self.broker_addr_table.is_empty() {
1207 return false;
1208 }
1209
1210 if is_rebalance {
1212 self.reset_broker_addr_heartbeat_fingerprint_map().await;
1213 }
1214
1215 let current_fingerprint = heartbeat_data_with_sub.compute_heartbeat_fingerprint();
1217
1218 for broker_entry in self.broker_addr_table.iter() {
1220 let (broker_name, broker_addrs) = broker_entry.pair();
1221 if broker_addrs.is_empty() {
1222 continue;
1223 }
1224
1225 for (id, addr) in broker_addrs.iter() {
1226 if addr.is_empty() {
1227 continue;
1228 }
1229 if consumer_empty && *id != mix_all::MASTER_ID {
1231 continue;
1232 }
1233
1234 let mut heartbeat_data = heartbeat_data_with_sub.clone();
1236 heartbeat_data.heartbeat_fingerprint = current_fingerprint;
1237
1238 self.send_heartbeat_to_broker_v2(*id, broker_name, addr, heartbeat_data)
1239 .await;
1240 }
1241 }
1242
1243 true
1244 }
1245
1246 async fn reset_broker_addr_heartbeat_fingerprint_map(&self) {
1247 self.broker_heartbeat_fingerprint_table.clear();
1248 }
1249
1250 async fn send_heartbeat_to_all_broker_concurrently(&self) -> bool {
1252 let heartbeat_data = self.prepare_heartbeat_data(false).await;
1253 let producer_empty = heartbeat_data.producer_data_set.is_empty();
1254 let consumer_empty = heartbeat_data.consumer_data_set.is_empty();
1255
1256 if producer_empty && consumer_empty {
1257 warn!(
1258 "sending heartbeat concurrently, but no consumer and no producer. [{}]",
1259 self.client_id
1260 );
1261 return false;
1262 }
1263
1264 let broker_list: Vec<(u64, CheetahString, CheetahString)> = {
1266 if self.broker_addr_table.is_empty() {
1267 return false;
1268 }
1269
1270 let mut list = Vec::new();
1271 for broker_entry in self.broker_addr_table.iter() {
1272 let (broker_name, broker_addrs) = broker_entry.pair();
1273 if broker_addrs.is_empty() {
1274 continue;
1275 }
1276
1277 for (id, addr) in broker_addrs.iter() {
1278 if addr.is_empty() {
1279 continue;
1280 }
1281 if consumer_empty && *id != mix_all::MASTER_ID {
1283 continue;
1284 }
1285 list.push((*id, broker_name.clone(), addr.clone()));
1286 }
1287 }
1288 list
1289 };
1290
1291 if broker_list.is_empty() {
1292 return false;
1293 }
1294
1295 let task_count = broker_list.len();
1296
1297 let mut tasks = Vec::new();
1299 let Some(mq_client_api) = self.mq_client_api_impl.as_ref().cloned() else {
1300 warn!("send heartbeat concurrently skipped because mq_client_api_impl is None");
1301 return false;
1302 };
1303 for (broker_id, broker_name, addr) in broker_list {
1304 let heartbeat_data = heartbeat_data.clone();
1306 let client_id = self.client_id.clone();
1307 let mq_client_api = mq_client_api.clone();
1308 let timeout = self.client_config.mq_client_api_timeout;
1309 let send_heartbeat_times_total = self.send_heartbeat_times_total.clone();
1310 let topic_route_table_clone = self.topic_route_table.clone();
1311
1312 let task = async move {
1314 match mq_client_api
1315 .mut_from_ref()
1316 .send_heartbeat(&addr, &heartbeat_data, timeout)
1317 .await
1318 {
1319 Ok((version, _response)) => {
1320 let times = send_heartbeat_times_total.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1321 if times % 20 == 0 {
1322 info!(
1323 "send heart beat to broker[{} {} {}] success (concurrent)",
1324 broker_name, broker_id, addr
1325 );
1326 }
1327 (broker_name, addr, Some(version), true)
1328 }
1329 Err(_) => {
1330 let is_in_ns = topic_route_table_clone.iter().any(|route_entry| {
1332 route_entry.value().broker_datas.iter().any(|bd| {
1333 bd.broker_addrs()
1334 .iter()
1335 .any(|(_, broker_addr)| broker_addr.as_str() == addr.as_str())
1336 })
1337 });
1338
1339 if is_in_ns {
1340 warn!(
1341 "send heart beat to broker[{} {} {}] failed (concurrent)",
1342 broker_name, broker_id, addr
1343 );
1344 } else {
1345 warn!(
1346 "send heart beat to broker[{} {} {}] exception, because the broker not up, forget it \
1347 (concurrent)",
1348 broker_name, broker_id, addr
1349 );
1350 }
1351 (broker_name, addr, None, false)
1352 }
1353 }
1354 };
1355
1356 tasks.push(task);
1357 }
1358
1359 let results = match tokio::time::timeout(Duration::from_millis(3000), future::join_all(tasks)).await {
1361 Ok(results) => results,
1362 Err(_) => {
1363 warn!(
1364 "Concurrent heartbeat timeout after 3000ms for client [{}]",
1365 self.client_id
1366 );
1367 return false;
1368 }
1369 };
1370
1371 let mut success_count = 0;
1374
1375 for (broker_name, addr, version_opt, success) in results {
1376 if success {
1377 success_count += 1;
1378 if let Some(version) = version_opt {
1379 self.broker_version_table
1380 .entry(broker_name.clone())
1381 .or_default()
1382 .insert(addr, version);
1383 }
1384 }
1385 }
1386
1387 info!(
1388 "Concurrent heartbeat completed for client [{}]: {}/{} succeeded",
1389 self.client_id, success_count, task_count
1390 );
1391
1392 success_count > 0
1394 }
1395
1396 async fn send_heartbeat_to_all_broker(&self) -> bool {
1397 let heartbeat_data = self.prepare_heartbeat_data(false).await;
1398 let producer_empty = heartbeat_data.producer_data_set.is_empty();
1399 let consumer_empty = heartbeat_data.consumer_data_set.is_empty();
1400 if producer_empty && consumer_empty {
1401 warn!(
1402 "sending heartbeat, but no consumer and no producer. [{}]",
1403 self.client_id
1404 );
1405 return false;
1406 }
1407 if self.broker_addr_table.is_empty() {
1408 return false;
1409 }
1410 for broker_entry in self.broker_addr_table.iter() {
1411 let (broker_name, broker_addrs) = broker_entry.pair();
1412 if broker_addrs.is_empty() {
1413 continue;
1414 }
1415 for (id, addr) in broker_addrs.iter() {
1416 if addr.is_empty() {
1417 continue;
1418 }
1419 if consumer_empty && *id != mix_all::MASTER_ID {
1420 continue;
1421 }
1422 self.send_heartbeat_to_broker_inner(*id, broker_name, addr, &heartbeat_data)
1423 .await;
1424 }
1425 }
1426
1427 true
1428 }
1429
1430 pub async fn send_heartbeat_to_broker(
1431 &self,
1432 id: u64,
1433 broker_name: &CheetahString,
1434 addr: &CheetahString,
1435 strict_lock_mode: bool,
1436 ) -> bool {
1437 let _guard = match self.lock_heartbeat.try_lock().await {
1438 Some(g) => g,
1439 None => {
1440 if strict_lock_mode {
1441 warn!("lock heartBeat, but failed. [{}]", self.client_id);
1442 }
1443 return false;
1444 }
1445 };
1446
1447 let heartbeat_data = self.prepare_heartbeat_data(false).await;
1448 let producer_empty = heartbeat_data.producer_data_set.is_empty();
1449 let consumer_empty = heartbeat_data.consumer_data_set.is_empty();
1450 if producer_empty && consumer_empty {
1451 warn!(
1452 "sending heartbeat, but no consumer and no producer. [{}]",
1453 self.client_id
1454 );
1455 return false;
1456 }
1457
1458 if self.client_config.use_heartbeat_v2 {
1459 self.send_heartbeat_to_broker_v2(id, broker_name, addr, heartbeat_data)
1460 .await
1461 } else {
1462 let (result, _) = self
1463 .send_heartbeat_to_broker_inner(id, broker_name, addr, &heartbeat_data)
1464 .await;
1465 result
1466 }
1467 }
1468
1469 async fn send_heartbeat_to_broker_v2(
1471 &self,
1472 id: u64,
1473 broker_name: &CheetahString,
1474 addr: &CheetahString,
1475 mut heartbeat_data_with_sub: HeartbeatData,
1476 ) -> bool {
1477 let current_fingerprint = heartbeat_data_with_sub.compute_heartbeat_fingerprint();
1479 heartbeat_data_with_sub.heartbeat_fingerprint = current_fingerprint;
1480
1481 let broker_support_v2 = self.broker_support_v2_heartbeat_set.contains_key(addr);
1483
1484 let last_fingerprint = self
1486 .broker_heartbeat_fingerprint_table
1487 .get(addr)
1488 .map(|entry| *entry.value());
1489
1490 let should_send_minimal = broker_support_v2 && last_fingerprint == Some(current_fingerprint);
1492
1493 let heartbeat_data = if should_send_minimal {
1494 let mut minimal = self.prepare_heartbeat_data(true).await;
1496 minimal.heartbeat_fingerprint = current_fingerprint;
1497 minimal
1498 } else {
1499 heartbeat_data_with_sub.clone()
1501 };
1502
1503 let (result, support_v2) = self
1505 .send_heartbeat_to_broker_inner(id, broker_name, addr, &heartbeat_data)
1506 .await;
1507
1508 if result {
1509 if let Some(support_v2) = support_v2 {
1510 if support_v2 {
1511 self.broker_support_v2_heartbeat_set.insert(addr.clone(), ());
1513
1514 if !should_send_minimal {
1516 self.broker_heartbeat_fingerprint_table
1517 .insert(addr.clone(), current_fingerprint);
1518 }
1519 } else {
1520 self.broker_support_v2_heartbeat_set.remove(addr);
1522 self.broker_heartbeat_fingerprint_table.remove(addr);
1523
1524 warn!("Broker {} does not support HeartbeatV2, downgrading to V1", addr);
1525 }
1526 }
1527 }
1528
1529 result
1530 }
1531
1532 async fn send_heartbeat_to_broker_inner(
1533 &self,
1534 id: u64,
1535 broker_name: &CheetahString,
1536 addr: &CheetahString,
1537 heartbeat_data: &HeartbeatData,
1538 ) -> (bool, Option<bool>) {
1539 let Some(mq_client_api_impl) = self.mq_client_api_impl.as_ref() else {
1540 warn!(
1541 "send heart beat to broker[{} {} {}] skipped because mq_client_api_impl is None",
1542 broker_name, id, addr
1543 );
1544 return (false, None);
1545 };
1546
1547 match mq_client_api_impl
1548 .mut_from_ref()
1549 .send_heartbeat(addr, heartbeat_data, self.client_config.mq_client_api_timeout)
1550 .await
1551 {
1552 Ok((version, response)) => {
1553 self.broker_version_table
1554 .entry(broker_name.clone())
1555 .or_default()
1556 .insert(addr.clone(), version);
1557
1558 let times = self
1559 .send_heartbeat_times_total
1560 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1561 if times % 20 == 0 {
1562 info!("send heart beat to broker[{} {} {}] success", broker_name, id, addr);
1563 }
1564
1565 let support_v2 = response.and_then(|resp| {
1567 resp.ext_fields().and_then(|fields| {
1568 fields
1569 .get(rocketmq_common::common::mix_all::IS_SUPPORT_HEART_BEAT_V2)
1570 .and_then(|v| v.parse::<bool>().ok())
1571 })
1572 });
1573
1574 (true, support_v2)
1575 }
1576 Err(_) => {
1577 if self.is_broker_in_name_server(addr).await {
1578 warn!("send heart beat to broker[{} {} {}] failed", broker_name, id, addr);
1579 } else {
1580 warn!(
1581 "send heart beat to broker[{} {} {}] exception, because the broker not up, forget it",
1582 broker_name, id, addr
1583 );
1584 }
1585 (false, None)
1586 }
1587 }
1588 }
1589
1590 async fn is_broker_in_name_server(&self, broker_name: &str) -> bool {
1591 for entry in self.topic_route_table.iter() {
1592 for bd in entry.value().broker_datas.iter() {
1593 for value in bd.broker_addrs().values() {
1594 if value.as_str() == broker_name {
1595 return true;
1596 }
1597 }
1598 }
1599 }
1600 false
1601 }
1602
1603 async fn prepare_heartbeat_data(&self, is_without_sub: bool) -> HeartbeatData {
1604 let mut heartbeat_data = HeartbeatData {
1605 client_id: self.client_id.clone(),
1606 ..Default::default()
1607 };
1608
1609 for entry in self.consumer_table.iter() {
1610 let mut consumer_data = ConsumerData {
1611 group_name: entry.value().group_name(),
1612 consume_type: entry.value().consume_type(),
1613 message_model: entry.value().message_model(),
1614 consume_from_where: entry.value().consume_from_where(),
1615 subscription_data_set: entry.value().subscriptions(),
1616 unit_mode: entry.value().is_unit_mode(),
1617 };
1618 if !is_without_sub {
1619 entry.value().subscriptions().iter().for_each(|sub| {
1620 consumer_data.subscription_data_set.insert(sub.clone());
1621 });
1622 }
1623 heartbeat_data.consumer_data_set.insert(consumer_data);
1624 }
1625 for entry in self.producer_table.iter() {
1626 let producer_data = ProducerData {
1627 group_name: entry.key().clone(),
1628 };
1629 heartbeat_data.producer_data_set.insert(producer_data);
1630 }
1631 heartbeat_data.is_without_sub = is_without_sub;
1632 heartbeat_data
1633 }
1634
1635 pub async fn register_consumer(&mut self, group: &CheetahString, consumer: MQConsumerInnerImpl) -> bool {
1636 if self.consumer_table.contains_key(group) {
1637 warn!("the consumer group[{}] exist already.", group);
1638 return false;
1639 }
1640 self.consumer_table.insert(group.clone(), consumer);
1641 true
1642 }
1643
1644 pub async fn check_client_in_broker(&mut self) -> rocketmq_error::RocketMQResult<()> {
1645 for entry in self.consumer_table.iter() {
1646 let subscription_inner = entry.value().subscriptions();
1647 if subscription_inner.is_empty() {
1648 return Ok(());
1649 }
1650 for subscription_data in subscription_inner.iter() {
1651 if ExpressionType::is_tag_type(Some(subscription_data.expression_type.as_str())) {
1652 continue;
1653 }
1654 let addr = self.find_broker_addr_by_topic(subscription_data.topic.as_str()).await;
1655 if let Some(addr) = addr {
1656 let Some(mq_client_api_impl) = self.mq_client_api_impl.as_mut() else {
1657 return Err(rocketmq_error::RocketMQError::ClientNotStarted);
1658 };
1659 match mq_client_api_impl
1660 .check_client_in_broker(
1661 addr.as_str(),
1662 entry.key().as_str(),
1663 self.client_id.as_str(),
1664 subscription_data,
1665 self.client_config.mq_client_api_timeout,
1666 )
1667 .await
1668 {
1669 Ok(_) => {}
1670 Err(e) => match e {
1671 rocketmq_error::RocketMQError::IllegalArgument(_) => {
1672 return Err(e);
1673 }
1674 _ => {
1675 let desc = format!(
1676 "Check client in broker error, maybe because you use {} to filter message, but \
1677 server has not been upgraded to support!This error would not affect the launch \
1678 of consumer, but may has impact on message receiving if you have use the new \
1679 features which are not supported by server, please check the log!",
1680 subscription_data.expression_type
1681 );
1682 }
1683 },
1684 }
1685 }
1686 }
1687 }
1688
1689 Ok(())
1690 }
1691
1692 pub async fn do_rebalance(&mut self) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
1693 let mut balanced = true;
1694 for entry in self.consumer_table.iter() {
1695 match entry.value().try_rebalance().await {
1696 Ok(result) => {
1697 if !result {
1698 balanced = false;
1699 }
1700 }
1701 Err(e) => {
1702 error!(
1703 "doRebalance for consumer group [{}] exception:{}",
1704 entry.key(),
1705 e.to_string()
1706 );
1707 balanced = false;
1708 }
1709 }
1710 }
1711 Ok(balanced)
1712 }
1713
1714 pub fn rebalance_later(&mut self, delay_millis: u64) {
1715 schedule_rebalance_wakeup(self.rebalance_service.clone(), Duration::from_millis(delay_millis));
1716 }
1717
1718 pub async fn find_broker_address_in_subscribe(
1719 &mut self,
1720 broker_name: &CheetahString,
1721 broker_id: u64,
1722 only_this_broker: bool,
1723 ) -> Option<FindBrokerResult> {
1724 if broker_name.is_empty() {
1725 return None;
1726 }
1727
1728 let mut broker_addr: Option<CheetahString> = None;
1729 let mut slave = false;
1730 let mut found = false;
1731
1732 if let Some(map_ref) = self.broker_addr_table.get(broker_name) {
1733 let map = map_ref.value();
1734 if let Some(addr) = map.get(&broker_id) {
1735 broker_addr = Some(addr.clone());
1736 slave = broker_id != mix_all::MASTER_ID;
1737 found = true;
1738 } else if !found && broker_id != mix_all::MASTER_ID {
1739 if let Some(addr) = map.get(&(broker_id + 1)) {
1740 broker_addr = Some(addr.clone());
1741 found = true;
1742 }
1743 }
1744 if !found && !only_this_broker {
1745 if let Some((key, value)) = map.iter().next() {
1746 broker_addr = Some(value.clone());
1747 slave = *key != mix_all::MASTER_ID;
1748 found = !value.is_empty();
1749 }
1750 }
1751 }
1752
1753 if found {
1754 let broker_addr = broker_addr?;
1755 let broker_version = self.find_broker_version(broker_name, broker_addr.as_str()).await;
1756 Some(FindBrokerResult {
1757 broker_addr,
1758 slave,
1759 broker_version,
1760 })
1761 } else {
1762 None
1763 }
1764 }
1765 async fn find_broker_version(&self, broker_name: &str, broker_addr: &str) -> i32 {
1766 if let Some(map) = self.broker_version_table.get(broker_name) {
1767 if let Some(version) = map.value().get(broker_addr) {
1768 return *version;
1769 }
1770 }
1771 0
1772 }
1773
1774 pub async fn select_consumer(&self, group: &str) -> Option<MQConsumerInnerImpl> {
1775 self.consumer_table.get(group).map(|entry| entry.value().clone())
1776 }
1777
1778 pub async fn reset_offset(
1779 &self,
1780 topic: &CheetahString,
1781 group: &CheetahString,
1782 offset_table: HashMap<MessageQueue, i64>,
1783 ) -> bool {
1784 let Some(consumer) = self.select_consumer(group).await else {
1785 warn!("reset offset failed because consumer group does not exist. group={group}, topic={topic}");
1786 return false;
1787 };
1788
1789 consumer.reset_offsets(topic, offset_table).await;
1790 true
1791 }
1792
1793 pub async fn get_consumer_status(
1794 &self,
1795 topic: &CheetahString,
1796 group: &CheetahString,
1797 ) -> Option<HashMap<MessageQueue, i64>> {
1798 let Some(consumer) = self.select_consumer(group).await else {
1799 warn!("get consumer status failed because consumer group does not exist. group={group}, topic={topic}");
1800 return None;
1801 };
1802
1803 Some(consumer.consumer_status(topic).await)
1804 }
1805
1806 pub async fn consumer_running_info(&self, consumer_group: &CheetahString) -> Option<ConsumerRunningInfo> {
1807 let Some(consumer) = self.select_consumer(consumer_group).await else {
1808 warn!("get consumer running info failed because consumer group does not exist. group={consumer_group}");
1809 return None;
1810 };
1811
1812 let mut info = consumer.consumer_running_info().await;
1813 let namesrv_addr = self
1814 .client_config
1815 .get_namesrv_addr()
1816 .map(|addr| {
1817 if addr.ends_with(';') {
1818 addr.to_string()
1819 } else {
1820 format!("{addr};")
1821 }
1822 })
1823 .unwrap_or_default();
1824
1825 info.set_property(ConsumerRunningInfo::PROP_NAMESERVER_ADDR, namesrv_addr);
1826 info.set_property(
1827 ConsumerRunningInfo::PROP_CONSUME_TYPE,
1828 match consumer.consume_type() {
1829 rocketmq_remoting::protocol::heartbeat::consume_type::ConsumeType::ConsumeActively => {
1830 "CONSUME_ACTIVELY"
1831 }
1832 rocketmq_remoting::protocol::heartbeat::consume_type::ConsumeType::ConsumePassively => {
1833 "CONSUME_PASSIVELY"
1834 }
1835 rocketmq_remoting::protocol::heartbeat::consume_type::ConsumeType::ConsumePop => "CONSUME_POP",
1836 },
1837 );
1838 info.set_property(ConsumerRunningInfo::PROP_CLIENT_VERSION, CURRENT_VERSION.name());
1839 info.sync_derived_fields_from_properties();
1840
1841 Some(info)
1842 }
1843
1844 pub async fn select_producer(&self, group: &str) -> Option<MQProducerInnerImpl> {
1845 self.producer_table.get(group).map(|entry| entry.value().clone())
1846 }
1847
1848 pub async fn unregister_consumer(&mut self, group: impl Into<CheetahString>) -> bool {
1849 let group = group.into();
1850 if group.is_empty() {
1851 warn!("unregister_consumer: group name is empty");
1852 return false;
1853 }
1854
1855 let removed = self.consumer_table.remove(&group).is_some();
1856
1857 if removed {
1858 info!("unregister consumer [{}] OK", group);
1859 self.unregister_client(None, Some(group)).await;
1860 true
1861 } else {
1862 warn!("unregister consumer [{}] failed: not found in consumer table", group);
1863 false
1864 }
1865 }
1866
1867 pub async fn unregister_producer(&mut self, group: impl Into<CheetahString>) -> bool {
1868 let group = group.into();
1869 if group.is_empty() {
1870 warn!("unregister_producer: group name is empty");
1871 return false;
1872 }
1873
1874 let removed = self.producer_table.remove(&group).is_some();
1875
1876 if removed {
1877 info!("unregister producer [{}] OK", group);
1878 self.unregister_client(Some(group), None).await;
1879 true
1880 } else {
1881 warn!("unregister producer [{}] failed: not found in producer table", group);
1882 false
1883 }
1884 }
1885
1886 pub async fn unregister_admin_ext(&mut self, group: impl Into<CheetahString>) {
1887 let _ = self.admin_ext_table.remove(&group.into());
1888 }
1889 async fn unregister_client(
1890 &mut self,
1891 producer_group: Option<CheetahString>,
1892 consumer_group: Option<CheetahString>,
1893 ) {
1894 let Some(mq_client_api_impl) = self.mq_client_api_impl.as_mut() else {
1895 warn!(
1896 "unregister client[Producer: {:?} Consumer: {:?}] skipped because mq_client_api_impl is None",
1897 producer_group, consumer_group
1898 );
1899 return;
1900 };
1901 for broker_entry in self.broker_addr_table.iter() {
1902 let (broker_name, broker_addrs) = broker_entry.pair();
1903 for (id, addr) in broker_addrs.iter() {
1904 if let Err(err) = mq_client_api_impl
1905 .unregister_client(
1906 addr,
1907 self.client_id.clone(),
1908 producer_group.clone(),
1909 consumer_group.clone(),
1910 self.client_config.mq_client_api_timeout,
1911 )
1912 .await
1913 {
1914 warn!(
1915 "unregister client[Producer: {:?} Consumer: {:?}] from broker[{} {} {}] failed: {}",
1916 producer_group, consumer_group, broker_name, id, addr, err
1917 );
1918 } else {
1919 info!(
1920 "unregister client[Producer: {:?} Consumer: {:?}] from broker[{} {} {}] success",
1921 producer_group, consumer_group, broker_name, id, addr,
1922 );
1923 }
1924 }
1925 }
1926 }
1927
1928 async fn is_broker_addr_exist_in_topic_route_table(&self, addr: &str) -> bool {
1929 for entry in self.topic_route_table.iter() {
1930 for bd in entry.value().broker_datas.iter() {
1931 for value in bd.broker_addrs().values() {
1932 if value.as_str() == addr {
1933 return true;
1934 }
1935 }
1936 }
1937 }
1938 false
1939 }
1940
1941 pub async fn query_assignment(
1961 &mut self,
1962 topic: &CheetahString,
1963 consumer_group: &CheetahString,
1964 strategy_name: &CheetahString,
1965 message_model: MessageModel,
1966 timeout: u64,
1967 ) -> rocketmq_error::RocketMQResult<Option<HashSet<MessageQueueAssignment>>> {
1968 let mut broker_addr = self.find_broker_addr_by_topic(topic).await;
1970
1971 if broker_addr.is_none() {
1973 self.update_topic_route_info_from_name_server_topic(topic).await;
1974 broker_addr = self.find_broker_addr_by_topic(topic).await;
1975 }
1976 if let Some(broker_addr) = broker_addr {
1977 let client_id = self.client_id.clone();
1978 match self.mq_client_api_impl.as_mut() {
1979 Some(api_impl) => {
1980 api_impl
1981 .query_assignment(
1982 &broker_addr,
1983 topic.clone(),
1984 consumer_group.clone(),
1985 client_id,
1986 strategy_name.clone(),
1987 message_model,
1988 timeout,
1989 )
1990 .await
1991 }
1992 None => Err(mq_client_err!("mq_client_api_impl is None")),
1993 }
1994 } else {
1995 Ok(None)
1996 }
1997 }
1998
1999 pub async fn consume_message_directly(
2000 &self,
2001 message: MessageExt,
2002 consumer_group: &CheetahString,
2003 broker_name: Option<CheetahString>,
2004 ) -> Option<ConsumeMessageDirectlyResult> {
2005 let consumer_inner = self.consumer_table.get(consumer_group);
2006 if let Some(entry) = consumer_inner {
2007 return entry.value().consume_message_directly(message, broker_name).await;
2008 }
2009
2010 None
2011 }
2012}
2013
2014#[allow(clippy::unnecessary_unwrap)]
2015pub fn topic_route_data2topic_publish_info(topic: &str, route: &mut TopicRouteData) -> TopicPublishInfo {
2016 let mut info = TopicPublishInfo {
2017 topic_route_data: Some(route.clone()),
2018 ..Default::default()
2019 };
2020 if let Some(order_topic_conf) = route.order_topic_conf.as_ref().filter(|conf| !conf.is_empty()) {
2021 for broker in order_topic_conf.split(';') {
2022 let item = broker.split(":").collect::<Vec<&str>>();
2023 if item.len() == 2 {
2024 let queue_num = match item[1].parse::<i32>() {
2025 Ok(queue_num) => queue_num,
2026 Err(error) => {
2027 warn!(
2028 "ignore invalid order topic conf entry for topic={}, broker={}, queue_nums={}, error={}",
2029 topic, item[0], item[1], error
2030 );
2031 continue;
2032 }
2033 };
2034 for i in 0..queue_num {
2035 let mq = MessageQueue::from_parts(topic, item[0], i);
2036 info.message_queue_list.push(mq);
2037 }
2038 }
2039 }
2040 info.order_topic = true;
2041 } else if route.order_topic_conf.is_none()
2042 && route
2043 .topic_queue_mapping_by_broker
2044 .as_ref()
2045 .is_some_and(|topic_queue_mapping_by_broker| !topic_queue_mapping_by_broker.is_empty())
2046 {
2047 info.order_topic = false;
2048 let mq_end_points = ClientMetadata::topic_route_data2endpoints_for_static_topic(topic, route);
2049 if let Some(mq_end_points) = mq_end_points {
2050 for (mq, broker_name) in mq_end_points {
2051 info.message_queue_list.push(mq);
2052 }
2053 }
2054 info.message_queue_list
2055 .sort_by(|a, b| match a.queue_id().cmp(&b.queue_id()) {
2056 Ordering::Less => std::cmp::Ordering::Less,
2057 Ordering::Equal => std::cmp::Ordering::Equal,
2058 Ordering::Greater => std::cmp::Ordering::Greater,
2059 });
2060 } else {
2061 route.queue_datas.sort();
2062 for queue_data in route.queue_datas.iter() {
2063 if PermName::is_writeable(queue_data.perm) {
2064 let mut broker_data = None;
2065 for bd in route.broker_datas.iter() {
2066 if bd.broker_name() == queue_data.broker_name.as_str() {
2067 broker_data = Some(bd.clone());
2068 break;
2069 }
2070 }
2071 let Some(broker_data) = broker_data else {
2072 continue;
2073 };
2074 if !broker_data.broker_addrs().contains_key(&(mix_all::MASTER_ID)) {
2075 continue;
2076 }
2077 for i in 0..queue_data.write_queue_nums {
2078 let mq = MessageQueue::from_parts(topic, queue_data.broker_name.as_str(), i as i32);
2079 info.message_queue_list.push(mq);
2080 }
2081 }
2082 }
2083 }
2084 info
2085}
2086
2087pub fn topic_route_data2_topic_publish_info(topic: &str, route: &mut TopicRouteData) -> TopicPublishInfo {
2088 topic_route_data2topic_publish_info(topic, route)
2089}
2090
2091pub fn topic_route_data2topic_subscribe_info(topic: &str, route: &TopicRouteData) -> HashSet<MessageQueue> {
2092 if let Some(ref topic_queue_mapping_by_broker) = route.topic_queue_mapping_by_broker {
2093 if !topic_queue_mapping_by_broker.is_empty() {
2094 let mq_endpoints = ClientMetadata::topic_route_data2endpoints_for_static_topic(topic, route);
2095 return mq_endpoints
2096 .unwrap_or_default()
2097 .keys()
2098 .cloned()
2099 .collect::<HashSet<MessageQueue>>();
2100 }
2101 }
2102 let mut mq_list = HashSet::new();
2103 for qd in &route.queue_datas {
2104 if PermName::is_readable(qd.perm) {
2105 for i in 0..qd.read_queue_nums {
2106 let mq = MessageQueue::from_parts(topic, qd.broker_name.as_str(), i as i32);
2107 mq_list.insert(mq);
2108 }
2109 }
2110 }
2111 mq_list
2112}
2113
2114pub fn topic_route_data2_topic_subscribe_info(topic: &str, route: &TopicRouteData) -> HashSet<MessageQueue> {
2115 topic_route_data2topic_subscribe_info(topic, route)
2116}
2117
2118#[cfg(test)]
2119mod tests {
2120 use rocketmq_error::RocketMQError;
2121
2122 use super::*;
2123
2124 #[tokio::test]
2125 async fn get_mq_client_api_impl_returns_client_not_started_instead_of_panicking() {
2126 let mut instance = MQClientInstance::new_arc(ClientConfig::default(), 0, "test-client", None);
2127 instance.mq_client_api_impl = None;
2128
2129 let result = instance.get_mq_client_api_impl();
2130
2131 assert!(matches!(result, Err(RocketMQError::ClientNotStarted)));
2132 }
2133
2134 #[tokio::test]
2135 async fn mq_client_api_impl_inherits_tls_flag_from_client_config() {
2136 let mut client_config = ClientConfig::default();
2137 client_config.set_use_tls(true);
2138 client_config.set_tls_client_trust_cert_path("/certs/ca.pem");
2139
2140 let instance = MQClientInstance::new_arc(client_config, 0, "test-client-tls", None);
2141 let api_impl = instance
2142 .get_mq_client_api_impl()
2143 .expect("MQClientInstance should initialize MQClientAPIImpl");
2144
2145 assert!(api_impl.is_use_tls());
2146 assert_eq!(
2147 api_impl
2148 .get_remoting_client()
2149 .tls_config()
2150 .client
2151 .trust_cert_path
2152 .as_deref(),
2153 Some("/certs/ca.pem")
2154 );
2155 }
2156
2157 #[test]
2158 fn new_arc_initializes_namesrv_inside_current_thread_runtime_without_block_in_place_panic() {
2159 let runtime = tokio::runtime::Builder::new_current_thread()
2160 .enable_all()
2161 .build()
2162 .unwrap();
2163
2164 runtime.block_on(async {
2165 let client_config = ClientConfig {
2166 namesrv_addr: Some(CheetahString::from_static_str("127.0.0.1:9876;127.0.0.2:9876")),
2167 ..Default::default()
2168 };
2169
2170 let instance = MQClientInstance::new_arc(client_config, 0, "current-thread-namesrv-test", None);
2171 let api_impl = instance
2172 .get_mq_client_api_impl()
2173 .expect("MQClientInstance should initialize MQClientAPIImpl");
2174 let addrs = api_impl.get_name_server_address_list();
2175
2176 assert_eq!(addrs.len(), 2);
2177 assert!(addrs.contains(&CheetahString::from_static_str("127.0.0.1:9876")));
2178 assert!(addrs.contains(&CheetahString::from_static_str("127.0.0.2:9876")));
2179 });
2180 }
2181
2182 #[test]
2183 fn new_arc_without_tokio_runtime_initializes_client_without_spawn_panic() {
2184 let client_config = ClientConfig {
2185 namesrv_addr: None,
2186 ..Default::default()
2187 };
2188
2189 let instance = MQClientInstance::new_arc(client_config, 0, "no-runtime-client-test", None);
2190
2191 assert!(instance.get_mq_client_api_impl().is_ok());
2192 }
2193
2194 #[test]
2195 fn delayed_rebalance_without_tokio_runtime_does_not_spawn_panic() {
2196 let client_config = ClientConfig {
2197 namesrv_addr: None,
2198 ..Default::default()
2199 };
2200 let mut instance = MQClientInstance::new_arc(client_config, 0, "no-runtime-rebalance-test", None);
2201
2202 instance.re_balance_later(Duration::from_millis(1));
2203 instance.rebalance_later(1);
2204
2205 std::thread::sleep(Duration::from_millis(20));
2206 }
2207
2208 #[tokio::test]
2209 async fn channel_event_methods_match_java_listener_contract_without_match() {
2210 let client_config = ClientConfig {
2211 namesrv_addr: None,
2212 ..Default::default()
2213 };
2214 let instance = MQClientInstance::new_arc(client_config, 0, "channel-event-test", None);
2215
2216 instance.on_channel_connect("127.0.0.1:10911");
2217 instance.on_channel_close("127.0.0.1:10911");
2218 instance.on_channel_exception("127.0.0.1:10911");
2219 instance.on_channel_idle("127.0.0.1:10911");
2220 instance.on_channel_active("127.0.0.1:10911").await;
2221 }
2222
2223 #[test]
2224 fn topic_publish_info_skips_invalid_order_topic_conf_without_panicking() {
2225 let mut route = TopicRouteData {
2226 order_topic_conf: Some(CheetahString::from("broker-a:not-a-number;broker-b:2")),
2227 ..Default::default()
2228 };
2229
2230 let info = topic_route_data2topic_publish_info("topic-a", &mut route);
2231
2232 assert!(info.order_topic);
2233 assert_eq!(info.message_queue_list.len(), 2);
2234 assert_eq!(info.message_queue_list[0].broker_name(), "broker-b");
2235 assert_eq!(info.message_queue_list[0].queue_id(), 0);
2236 assert_eq!(info.message_queue_list[1].broker_name(), "broker-b");
2237 assert_eq!(info.message_queue_list[1].queue_id(), 1);
2238 }
2239
2240 #[test]
2241 fn java_named_topic_publish_info_alias_delegates_to_route_converter() {
2242 let mut route = TopicRouteData {
2243 order_topic_conf: Some(CheetahString::from("broker-a:1")),
2244 ..Default::default()
2245 };
2246
2247 let info = topic_route_data2_topic_publish_info("topic-a", &mut route);
2248
2249 assert!(info.order_topic);
2250 assert_eq!(info.message_queue_list.len(), 1);
2251 assert_eq!(info.message_queue_list[0].topic(), "topic-a");
2252 assert_eq!(info.message_queue_list[0].broker_name(), "broker-a");
2253 assert_eq!(info.message_queue_list[0].queue_id(), 0);
2254 }
2255
2256 #[test]
2257 fn java_named_topic_subscribe_info_alias_delegates_to_route_converter() {
2258 let route = TopicRouteData {
2259 queue_datas: vec![rocketmq_remoting::protocol::route::route_data_view::QueueData {
2260 broker_name: CheetahString::from("broker-a"),
2261 read_queue_nums: 2,
2262 perm: PermName::PERM_READ,
2263 ..Default::default()
2264 }],
2265 ..Default::default()
2266 };
2267
2268 let info = topic_route_data2_topic_subscribe_info("topic-a", &route);
2269
2270 assert_eq!(info.len(), 2);
2271 assert!(info.contains(&MessageQueue::from_parts("topic-a", "broker-a", 0)));
2272 assert!(info.contains(&MessageQueue::from_parts("topic-a", "broker-a", 1)));
2273 }
2274
2275 #[tokio::test]
2276 async fn rebalance_immediately_alias_matches_existing_wakeup_path() {
2277 let client_config = ClientConfig {
2278 namesrv_addr: None,
2279 ..Default::default()
2280 };
2281 let instance = MQClientInstance::new_arc(client_config, 0, "rebalance-alias-test", None);
2282
2283 instance.rebalance_immediately();
2284 }
2285
2286 #[tokio::test]
2287 async fn parse_offset_table_from_broker_removes_namespace_like_java() {
2288 let client_config = ClientConfig {
2289 namesrv_addr: None,
2290 ..Default::default()
2291 };
2292 let instance = MQClientInstance::new_arc(client_config, 0, "parse-offset-test", None);
2293 let mut offsets = HashMap::new();
2294 offsets.insert(MessageQueue::from_parts("ns%topic-a", "broker-a", 0), 42);
2295
2296 let parsed = instance.parse_offset_table_from_broker(offsets, "ns");
2297
2298 assert_eq!(parsed.len(), 1);
2299 assert_eq!(
2300 parsed.get(&MessageQueue::from_parts("topic-a", "broker-a", 0)),
2301 Some(&42)
2302 );
2303 }
2304
2305 #[tokio::test]
2306 async fn find_broker_address_in_admin_returns_any_known_broker_like_java() {
2307 let client_config = ClientConfig {
2308 namesrv_addr: None,
2309 ..Default::default()
2310 };
2311 let instance = MQClientInstance::new_arc(client_config, 0, "admin-broker-test", None);
2312 let broker_name = CheetahString::from_static_str("broker-a");
2313 let broker_addr = CheetahString::from_static_str("127.0.0.1:10912");
2314 let mut addrs = HashMap::new();
2315 addrs.insert(1, broker_addr.clone());
2316 instance.broker_addr_table.insert(broker_name.clone(), addrs);
2317 let mut versions = HashMap::new();
2318 versions.insert(broker_addr.clone(), 321);
2319 instance.broker_version_table.insert(broker_name.clone(), versions);
2320
2321 let result = instance
2322 .find_broker_address_in_admin(&broker_name)
2323 .await
2324 .expect("known broker should be returned");
2325
2326 assert_eq!(result.broker_addr, broker_addr);
2327 assert!(result.slave);
2328 assert_eq!(result.broker_version, 321);
2329 }
2330
2331 #[tokio::test]
2332 async fn query_topic_route_data_returns_cached_route_before_nameserver_lookup() {
2333 let client_config = ClientConfig {
2334 namesrv_addr: None,
2335 ..Default::default()
2336 };
2337 let mut instance = MQClientInstance::new_arc(client_config, 0, "query-route-test", None);
2338 let topic = CheetahString::from_static_str("topic-a");
2339 let route = TopicRouteData {
2340 order_topic_conf: Some(CheetahString::from_static_str("broker-a:1")),
2341 ..Default::default()
2342 };
2343 instance.topic_route_table.insert(topic.clone(), route.clone());
2344
2345 let result = instance
2346 .query_topic_route_data(&topic)
2347 .await
2348 .expect("cached topic route should be returned");
2349
2350 assert_eq!(result, route);
2351 }
2352
2353 #[test]
2354 fn get_topic_config_request_header_sets_lo_like_java_client() {
2355 let header = get_topic_config_request_header(CheetahString::from("TopicTest"));
2356
2357 assert_eq!(
2358 rocketmq_remoting::protocol::header::message_operation_header::TopicRequestHeaderTrait::topic(&header),
2359 "TopicTest"
2360 );
2361 assert_eq!(
2362 rocketmq_remoting::protocol::header::message_operation_header::TopicRequestHeaderTrait::lo(&header),
2363 Some(true)
2364 );
2365 }
2366}