[ISSUE #1076]PullMessageService and RebalanceService add shutdown method

mxsm · mxsm · commit a464f860475e · 2024-10-22T10:00:58.000Z
diff --git a/rocketmq-client/src/consumer/consumer_impl/pull_message_service.rs b/rocketmq-client/src/consumer/consumer_impl/pull_message_service.rs
@@ -16,6 +16,7 @@
  */
 use rocketmq_common::common::message::message_enum::MessageRequestMode;
 use rocketmq_common::ArcRefCellWrapper;
+use rocketmq_rust::Shutdown;
 use tracing::info;
 use tracing::warn;
 
@@ -27,19 +28,25 @@ use crate::factory::mq_client_instance::MQClientInstance;
 #[derive(Clone)]
 pub struct PullMessageService {
     tx: Option<tokio::sync::mpsc::Sender<Box<dyn MessageRequest + Send + 'static>>>,
+    tx_shutdown: Option<tokio::sync::broadcast::Sender<()>>,
 }
 
 impl PullMessageService {
     pub fn new() -> Self {
-        PullMessageService { tx: None }
+        PullMessageService {
+            tx: None,
+            tx_shutdown: None,
+        }
     }
     pub async fn start(&mut self, mut instance: ArcRefCellWrapper<MQClientInstance>) {
         let (tx, mut rx) =
             tokio::sync::mpsc::channel::<Box<dyn MessageRequest + Send + 'static>>(1024 * 4);
+        let (mut shutdown, tx_shutdown) = Shutdown::new(1);
         self.tx = Some(tx);
+        self.tx_shutdown = Some(tx_shutdown);
         tokio::spawn(async move {
             info!(">>>>>>>>>>>>>>>>>>>>>>>PullMessageService  started<<<<<<<<<<<<<<<<<<<<<<<<<<<<");
-            while let Some(request) = rx.recv().await {
+            /*while let Some(request) = rx.recv().await {
                 if request.get_message_request_mode() == MessageRequestMode::Pull {
                     let pull_request =
                         unsafe { *Box::from_raw(Box::into_raw(request) as *mut PullRequest) };
@@ -49,6 +56,32 @@ impl PullMessageService {
                         unsafe { *Box::from_raw(Box::into_raw(request) as *mut PopRequest) };
                     PullMessageService::pop_message(pop_request, instance.as_mut()).await;
                 }
+            }*/
+            if shutdown.is_shutdown() {
+                info!("PullMessageService shutdown");
+                return;
+            }
+            loop {
+                tokio::select! {
+                    _ = shutdown.recv() => {
+                        info!("PullMessageService shutdown");
+                    }
+                    Some(request) = rx.recv() => {
+                        if request.get_message_request_mode() == MessageRequestMode::Pull {
+                            let pull_request =
+                                unsafe { *Box::from_raw(Box::into_raw(request) as *mut PullRequest) };
+                            PullMessageService::pull_message(pull_request, instance.as_mut()).await;
+                        } else {
+                            let pop_request =
+                                unsafe { *Box::from_raw(Box::into_raw(request) as *mut PopRequest) };
+                            PullMessageService::pop_message(pop_request, instance.as_mut()).await;
+                        }
+                    }
+                }
+                if shutdown.is_shutdown() {
+                    info!("PullMessageService shutdown");
+                    break;
+                }
             }
         });
     }
@@ -106,4 +139,10 @@ impl PullMessageService {
             warn!("Failed to send pull request to pull_tx, error: {:?}", e);
         }
     }
+
+    pub fn shutdown(&self) {
+        if let Err(e) = self.tx_shutdown.as_ref().unwrap().send(()) {
+            warn!("Failed to send shutdown signal to pull_tx, error: {:?}", e);
+        }
+    }
 }
diff --git a/rocketmq-client/src/consumer/consumer_impl/re_balance/rebalance_service.rs b/rocketmq-client/src/consumer/consumer_impl/re_balance/rebalance_service.rs
@@ -19,10 +19,12 @@ use std::time::Duration;
 
 use once_cell::sync::Lazy;
 use rocketmq_common::ArcRefCellWrapper;
+use rocketmq_rust::Shutdown;
 use tokio::select;
 use tokio::sync::Notify;
 use tokio::time::Instant;
 use tracing::info;
+use tracing::warn;
 
 use crate::factory::mq_client_instance::MQClientInstance;
 
@@ -47,17 +49,21 @@ static MIN_INTERVAL: Lazy<Duration> = Lazy::new(|| {
 #[derive(Clone)]
 pub struct RebalanceService {
     notify: Arc<Notify>,
+    tx_shutdown: Option<tokio::sync::broadcast::Sender<()>>,
 }
 
 impl RebalanceService {
     pub fn new() -> Self {
         RebalanceService {
             notify: Arc::new(Notify::new()),
+            tx_shutdown: None,
         }
     }
 
     pub async fn start(&mut self, mut instance: ArcRefCellWrapper<MQClientInstance>) {
         let notify = self.notify.clone();
+        let (mut shutdown, tx_shutdown) = Shutdown::new(1);
+        self.tx_shutdown = Some(tx_shutdown);
         tokio::spawn(async move {
             let mut last_rebalance_timestamp = Instant::now();
             let min_interval = *MIN_INTERVAL;
@@ -66,8 +72,12 @@ impl RebalanceService {
             loop {
                 select! {
                     _ = notify.notified() => {}
+                    _ = shutdown.recv() => {info!("RebalanceService shutdown");}
                     _ = tokio::time::sleep(real_wait_interval) => {}
                 }
+                if shutdown.is_shutdown() {
+                    return;
+                }
                 let interval = Instant::now() - last_rebalance_timestamp;
                 if interval < min_interval {
                     real_wait_interval = min_interval - interval;
@@ -87,4 +97,10 @@ impl RebalanceService {
     pub fn wakeup(&self) {
         self.notify.notify_waiters();
     }
+
+    pub fn shutdown(&self) {
+        if let Err(e) = self.tx_shutdown.as_ref().unwrap().send(()) {
+            warn!("Failed to send shutdown signal to pull_tx, error: {:?}", e);
+        }
+    }
 }
