[ISSUE #2227]💫Implement broker graceful shutdown🧑‍💻

rocketmq-broker/src/broker_bootstrap.rs

@@ -14,9 +14,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
+use log::info;
 use rocketmq_common::common::broker::broker_config::BrokerConfig;
 use rocketmq_common::common::server::config::ServerConfig;
+use rocketmq_rust::wait_for_signal;
 use rocketmq_store::config::message_store_config::MessageStoreConfig;
 use tracing::error;
 
@@ -32,7 +33,10 @@
             error!("initialize fail");
             return;
         }
-        let (_start_result, _ctrl_c) = tokio::join!(self.start(), tokio::signal::ctrl_c());
+        let (shutdown_tx, shutdown_rx) = tokio::sync::broadcast::channel(1);
+        self.broker_runtime.shutdown_rx = Some(shutdown_rx);
+
+        tokio::join!(self.start(), wait_for_signal_inner(shutdown_tx));
     }
 
     async fn initialize(&mut self) -> bool {
@@ -44,6 +48,16 @@
     }
 }
 
+async fn wait_for_signal_inner(shutdown_tx: tokio::sync::broadcast::Sender<()>) {
+    tokio::select! {
+        _ = wait_for_signal() => {
+            info!("Broker Received signal, initiating shutdown...");
+        }
+    }
+    // Send shutdown signal to all tasks
+    let _ = shutdown_tx.send(());
+}
+
 pub struct Builder {
     broker_config: BrokerConfig,
     message_store_config: MessageStoreConfig,

rocketmq-broker/src/broker_runtime.rs

@@ -151,6 +151,8 @@
     broker_runtime: Option<RocketMQRuntime>,
     shutdown: Arc<AtomicBool>,
     shutdown_hook: Option<BrokerShutdownHook>,
+    // receiver for shutdown signal
+    pub(crate) shutdown_rx: Option<tokio::sync::broadcast::Receiver<()>>,
 }
 
 impl BrokerRuntime {
@@ -240,6 +242,7 @@
             broker_runtime: Some(runtime),
             shutdown: Arc::new(AtomicBool::new(false)),
             shutdown_hook: None,
+            shutdown_rx: None,
         }
     }
 
@@ -265,9 +268,9 @@
             pull_request_hold_service.shutdown();
         }
 
-        /*        if let Some(runtime) = self.broker_runtime.take() {
+        if let Some(runtime) = self.broker_runtime.take() {
             runtime.shutdown();
-        }*/
+        }
     }
 
     pub(crate) fn shutdown_basic_service(&mut self) {
@@ -913,6 +916,13 @@
             "Rocketmq Broker({} ----Rust) start success",
             self.inner.broker_config.broker_identity.broker_name
         );
+        tokio::select! {
+            _ = self.shutdown_rx.as_mut().unwrap().recv() => {
+                info!("Broker Shutdown received, initiating graceful shutdown...");
+                self.shutdown();
+                info!("Broker Shutdown complete");
+            }
+        }
     }
 
     pub(crate) fn schedule_send_heartbeat(&mut self) {}