[ISSUE #353]Add Broker shutdown logic

Author: mxsm

rocketmq-broker/src/broker.rs

@@ -0,0 +1,18 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+pub mod broker_hook;

rocketmq-broker/src/broker/broker_hook.rs

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+use std::sync::Arc;
+
+pub trait ShutdownHook {
+    fn before_shutdown(&self);
+}
+
+pub type BrokerShutdownHook = Arc<dyn ShutdownHook + Send + Sync + 'static>;

rocketmq-broker/src/broker_runtime.rs

@@ -44,6 +44,7 @@ use rocketmq_store::{
 use tracing::{info, warn};
 
 use crate::{
+    broker::broker_hook::BrokerShutdownHook,
     client::manager::producer_manager::ProducerManager,
     filter::manager::consumer_filter_manager::ConsumerFilterManager,
     offset::manager::{
@@ -84,6 +85,8 @@ pub(crate) struct BrokerRuntime {
     broker_runtime: Option<RocketMQRuntime>,
     producer_manager: Arc<ProducerManager>,
     drop: Arc<AtomicBool>,
+    shutdown: Arc<AtomicBool>,
+    shutdown_hook: Option<BrokerShutdownHook>,
 }
 
 impl Clone for BrokerRuntime {
@@ -105,6 +108,8 @@ impl Clone for BrokerRuntime {
             broker_runtime: None,
             producer_manager: self.producer_manager.clone(),
             drop: self.drop.clone(),
+            shutdown: self.shutdown.clone(),
+            shutdown_hook: self.shutdown_hook.clone(),
         }
     }
 }
@@ -151,6 +156,8 @@ impl BrokerRuntime {
             broker_runtime: Some(runtime),
             producer_manager: Arc::new(ProducerManager::new()),
             drop: Arc::new(AtomicBool::new(false)),
+            shutdown: Arc::new(AtomicBool::new(false)),
+            shutdown_hook: None,
         }
     }
 
@@ -161,6 +168,22 @@ impl BrokerRuntime {
     pub(crate) fn message_store_config(&self) -> &MessageStoreConfig {
         &self.message_store_config
     }
+
+    pub fn shutdown(&mut self) {
+        self.broker_out_api.shutdown();
+        self.message_store.as_mut().unwrap().shutdown();
+        if let Some(runtime) = self.broker_runtime.take() {
+            runtime.shutdown();
+        }
+    }
+
+    pub(crate) fn shutdown_basic_service(&mut self) {
+        self.shutdown.store(true, Ordering::SeqCst);
+
+        if let Some(hook) = self.shutdown_hook.as_ref() {
+            hook.before_shutdown();
+        }
+    }
 }
 
 impl Drop for BrokerRuntime {
@@ -170,14 +193,8 @@ impl Drop for BrokerRuntime {
                 .clone()
                 .compare_exchange(false, true, Ordering::SeqCst, Ordering::Relaxed);
         if result.is_ok() {
-            if let Some(runtime) = self.broker_runtime.take() {
-                runtime.shutdown();
-            }
+            self.shutdown();
         }
-
-        /*if let Some(runtime) = self.broker_runtime.take() {
-            runtime.shutdown();
-        }*/
     }
 }
 

rocketmq-broker/src/lib.rs

@@ -21,6 +21,7 @@ pub use broker_bootstrap::{BrokerBootstrap, Builder};
 
 pub mod command;
 
+mod broker;
 mod broker_bootstrap;
 mod broker_path_config_helper;
 mod broker_runtime;

rocketmq-broker/src/out_api/broker_outer_api.rs

@@ -220,4 +220,6 @@ impl BrokerOuterAPI {
             let _result = tokio::join!(handle);
         }
     }
+
+    pub fn shutdown(&self) {}
 }

rocketmq-store/src/consume_queue/mapped_file_queue.rs

@@ -265,8 +265,6 @@ impl MappedFileQueue {
 
 #[cfg(test)]
 mod tests {
-    use std::fs;
-
     use super::*;
 
     #[test]

rocketmq-store/src/lib.rs

@@ -30,5 +30,6 @@ pub mod message_store;
 mod queue;
 pub(crate) mod services;
 pub mod status;
+mod store;
 pub mod store_path_config_helper;
 pub mod timer;

rocketmq-store/src/log_file.rs

@@ -39,6 +39,8 @@ pub trait RocketMQMessageStore: Clone {
     /// Throws an `Exception` if there is any error.
     fn start(&mut self) -> Result<(), Box<dyn std::error::Error>>;
 
+    fn shutdown(&mut self);
+
     fn set_confirm_offset(&mut self, phy_offset: i64);
 
     fn get_max_phy_offset(&self) -> i64;

rocketmq-store/src/log_file/commit_log.rs

@@ -127,6 +127,9 @@ impl CommitLog {
         result
     }
 
+    pub fn shutdown(&mut self) {}
+
+    pub fn destroy(&mut self) {}
     /*    pub fn set_local_file_message_store(
         &mut self,
         local_file_message_store: Weak<Mutex<LocalFileMessageStore>>,
@@ -280,7 +283,6 @@ impl CommitLog {
                     true,
                     &message_store_config,
                 );
-                println!("{}", dispatch_request);
                 current_pos += size;
                 if dispatch_request.success && dispatch_request.msg_size > 0 {
                     last_valid_msg_phy_offset = process_offset + mapped_file_offset;
@@ -402,7 +404,6 @@ fn check_message_and_return_size(
     read_body: bool,
     message_store_config: &Arc<MessageStoreConfig>,
 ) -> DispatchRequest {
-    info!("check_message_and_return_size");
     let total_size = bytes.get_i32();
     let magic_code = bytes.get_i32();
     if magic_code == MESSAGE_MAGIC_CODE || magic_code == MESSAGE_MAGIC_CODE_V2 {

rocketmq-store/src/message_store/default_message_store.rs

@@ -20,8 +20,9 @@ use std::{
     collections::HashMap,
     error::Error,
     fs,
+    path::PathBuf,
     sync::{
-        atomic::{AtomicI64, Ordering},
+        atomic::{AtomicBool, AtomicI64, Ordering},
         Arc,
     },
     time::Instant,
@@ -57,6 +58,7 @@ use crate::{
         build_consume_queue::CommitLogDispatcherBuildConsumeQueue,
         local_file_consume_queue_store::ConsumeQueueStore, ConsumeQueueStoreTrait,
     },
+    store::running_flags::RunningFlags,
     store_path_config_helper::{get_abort_file, get_store_checkpoint},
 };
 
@@ -77,6 +79,8 @@ pub struct DefaultMessageStore {
     dispatcher: CommitLogDispatcherDefault,
     broker_init_max_offset: Arc<AtomicI64>,
     state_machine_version: Arc<AtomicI64>,
+    shutdown: Arc<AtomicBool>,
+    running_flags: Arc<RunningFlags>,
 }
 
 impl Clone for DefaultMessageStore {
@@ -96,6 +100,8 @@ impl Clone for DefaultMessageStore {
             dispatcher: self.dispatcher.clone(),
             broker_init_max_offset: self.broker_init_max_offset.clone(),
             state_machine_version: self.state_machine_version.clone(),
+            shutdown: self.shutdown.clone(),
+            running_flags: self.running_flags.clone(),
         }
     }
 }
@@ -141,6 +147,8 @@ impl DefaultMessageStore {
             dispatcher,
             broker_init_max_offset: Arc::new(AtomicI64::new(-1)),
             state_machine_version: Arc::new(AtomicI64::new(0)),
+            shutdown: Arc::new(AtomicBool::new(false)),
+            running_flags: Arc::new(RunningFlags::new()),
         }
     }
 }
@@ -268,6 +276,17 @@ impl DefaultMessageStore {
     pub fn consume_queue_store_mut(&mut self) -> &mut ConsumeQueueStore {
         &mut self.consume_queue_store
     }
+
+    fn delete_file(&mut self, file_name: String) {
+        match fs::remove_file(PathBuf::from(file_name.as_str())) {
+            Ok(_) => {
+                info!("delete OK, file:{}", file_name);
+            }
+            Err(err) => {
+                error!("delete error, file:{}, {:?}", file_name, err);
+            }
+        }
+    }
 }
 
 impl MessageStore for DefaultMessageStore {
@@ -328,6 +347,20 @@ impl MessageStore for DefaultMessageStore {
         Ok(())
     }
 
+    fn shutdown(&mut self) {
+        if !self.shutdown.load(Ordering::Relaxed) {
+            self.shutdown.store(true, Ordering::SeqCst);
+            self.commit_log.shutdown();
+
+            if self.running_flags.is_writeable() {
+                //delete abort file
+                self.delete_file(get_abort_file(
+                    self.message_store_config.store_path_root_dir.as_str(),
+                ))
+            }
+        }
+    }
+
     fn set_confirm_offset(&mut self, phy_offset: i64) {
 
         // self.commit_log.set_confirm_offset(phy_offset);

rocketmq-store/src/store.rs

@@ -0,0 +1,18 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+pub(crate) mod running_flags;