[ISSUE #441]🔥Improving PutMessageHook logic🚀

Author: mxsm

Cargo.toml

@@ -66,4 +66,6 @@ parking_lot = "0.12"
 dirs = "5.0"
 trait-variant = "0.1.2"
 
-once_cell = "1.19.0"
+once_cell = "1.19.0"
+
+mockall = "0.12.1"

rocketmq-broker/Cargo.toml

@@ -56,6 +56,9 @@ dirs.workspace = true
 
 local-ip-address = "0.6.1"
 
+[dev-dependencies]
+mockall = "0.12.1"
+
 [[bin]]
 name = "rocketmq-broker-rust"
 path = "src/bin/broker_bootstrap_server.rs"

rocketmq-broker/src/broker_runtime.rs

@@ -47,6 +47,10 @@ use crate::{
     broker::broker_hook::BrokerShutdownHook,
     client::manager::producer_manager::ProducerManager,
     filter::manager::consumer_filter_manager::ConsumerFilterManager,
+    hook::{
+        batch_check_before_put_message::BatchCheckBeforePutMessageHook,
+        check_before_put_message::CheckBeforePutMessageHook,
+    },
     offset::manager::{
         consumer_offset_manager::ConsumerOffsetManager,
         consumer_order_info_manager::ConsumerOrderInfoManager,
@@ -260,6 +264,7 @@ impl BrokerRuntime {
             unimplemented!()
         }
         if self.message_store.is_some() {
+            self.register_message_store_hook();
             self.message_store.as_mut().unwrap().load().await;
         }
 
@@ -279,6 +284,18 @@ impl BrokerRuntime {
         result
     }
 
+    pub fn register_message_store_hook(&mut self) {
+        if let Some(ref mut message_store) = self.message_store {
+            message_store.set_put_message_hook(Box::new(CheckBeforePutMessageHook::new(
+                message_store.clone(),
+                self.message_store_config.clone(),
+            )));
+            message_store.set_put_message_hook(Box::new(BatchCheckBeforePutMessageHook::new(
+                self.topic_config_manager.topic_config_table(),
+            )));
+        }
+    }
+
     fn initialize_remoting_server(&mut self) {
 
         // fast broker server implementation in future versions

rocketmq-broker/src/hook.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 batch_check_before_put_message;
+pub(crate) mod check_before_put_message;

rocketmq-broker/src/hook/batch_check_before_put_message.rs

@@ -0,0 +1,44 @@
+/*
+ * 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::{collections::HashMap, sync::Arc};
+
+use rocketmq_common::common::{config::TopicConfig, message::message_single::MessageExt};
+use rocketmq_store::{
+    base::message_result::PutMessageResult, hook::put_message_hook::PutMessageHook,
+};
+
+use crate::util::hook_utils::HookUtils;
+
+pub struct BatchCheckBeforePutMessageHook {
+    topic_config_table: Arc<parking_lot::Mutex<HashMap<String, TopicConfig>>>,
+}
+
+impl BatchCheckBeforePutMessageHook {
+    pub fn new(topic_config_table: Arc<parking_lot::Mutex<HashMap<String, TopicConfig>>>) -> Self {
+        Self { topic_config_table }
+    }
+}
+
+impl PutMessageHook for BatchCheckBeforePutMessageHook {
+    fn hook_name(&self) -> String {
+        "batchCheckBeforePutMessage".to_string()
+    }
+
+    fn execute_before_put_message(&self, msg: &MessageExt) -> Option<PutMessageResult> {
+        HookUtils::check_inner_batch(&self.topic_config_table, msg)
+    }
+}

rocketmq-broker/src/hook/check_before_put_message.rs

@@ -0,0 +1,49 @@
+/*
+ * 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;
+
+use rocketmq_common::common::message::message_single::MessageExt;
+use rocketmq_store::{
+    base::message_result::PutMessageResult, config::message_store_config::MessageStoreConfig,
+    hook::put_message_hook::PutMessageHook, log_file::MessageStore,
+};
+
+use crate::util::hook_utils::HookUtils;
+
+pub struct CheckBeforePutMessageHook<MS> {
+    message_store: MS,
+    message_store_config: Arc<MessageStoreConfig>,
+}
+
+impl<MS: MessageStore> CheckBeforePutMessageHook<MS> {
+    pub fn new(message_store: MS, message_store_config: Arc<MessageStoreConfig>) -> Self {
+        Self {
+            message_store,
+            message_store_config,
+        }
+    }
+}
+
+impl<MS: MessageStore> PutMessageHook for CheckBeforePutMessageHook<MS> {
+    fn hook_name(&self) -> String {
+        "checkBeforePutMessage".to_string()
+    }
+
+    fn execute_before_put_message(&self, msg: &MessageExt) -> Option<PutMessageResult> {
+        HookUtils::check_before_put_message(&self.message_store, &self.message_store_config, msg)
+    }
+}

rocketmq-broker/src/lib.rs

@@ -29,6 +29,7 @@ mod client;
 mod coldctr;
 mod controller;
 mod filter;
+mod hook;
 mod longpolling;
 mod mqtrace;
 mod offset;

rocketmq-broker/src/util/hook_utils.rs

@@ -132,11 +132,11 @@ impl HookUtils {
                 PutMessageStatus::MessageIllegal,
             ));
         }
-        /*if message_store.is_os_page_cache_busy() {
-            return return Some(PutMessageResult::new_default(
+        if message_store.is_os_page_cache_busy() {
+            return Some(PutMessageResult::new_default(
                 PutMessageStatus::OsPageCacheBusy,
             ));
-        }*/
+        }
 
         None
     }
@@ -399,10 +399,12 @@ impl HookUtils {
 mod tests {
     use std::{collections::HashMap, error::Error, sync::Arc};
 
+    use parking_lot::RwLock;
     use rocketmq_common::common::{config::TopicConfig, message::message_single::MessageExt};
     use rocketmq_store::{
         base::{message_result::PutMessageResult, message_status_enum::PutMessageStatus},
         config::message_store_config::MessageStoreConfig,
+        hook::put_message_hook::BoxedPutMessageHook,
         log_file::MessageStore,
         store::running_flags::RunningFlags,
     };
@@ -464,6 +466,14 @@ mod tests {
         fn is_shutdown(&self) -> bool {
             false
         }
+
+        fn get_put_message_hook_list(&self) -> Arc<RwLock<Vec<BoxedPutMessageHook>>> {
+            todo!()
+        }
+
+        fn set_put_message_hook(&self, put_message_hook: BoxedPutMessageHook) {
+            todo!()
+        }
         // Implement required methods...
     }
 

rocketmq-store/Cargo.toml

@@ -60,4 +60,5 @@ log = "0.4.20"
 memmap2 = "0.9.4"
 trait-variant.workspace = true
 [dev-dependencies]
-tempfile = "3.10.0"
+tempfile = "3.10.0"
+mockall = { workspace = true }

rocketmq-store/src/config/message_store_config.rs

@@ -128,7 +128,7 @@ pub struct MessageStoreConfig {
     pub debug_lock_enable: bool,
     pub duplication_enable: bool,
     pub disk_fall_recorded: bool,
-    pub os_page_cache_busy_timeout_mills: usize,
+    pub os_page_cache_busy_timeout_mills: u64,
     pub default_query_max_num: usize,
     pub transient_store_pool_enable: bool,
     pub transient_store_pool_size: usize,
@@ -311,7 +311,7 @@ impl Default for MessageStoreConfig {
             debug_lock_enable: false,
             duplication_enable: false,
             disk_fall_recorded: false,
-            os_page_cache_busy_timeout_mills: 0,
+            os_page_cache_busy_timeout_mills: 1000,
             default_query_max_num: 0,
             transient_store_pool_enable: false,
             transient_store_pool_size: 0,

rocketmq-store/src/log_file.rs

@@ -15,17 +15,21 @@
  * limitations under the License.
  */
 
+use std::sync::Arc;
+
 use rocketmq_common::{
     common::message::message_single::MessageExtBrokerInner, TimeUtils::get_current_millis,
 };
 
-use crate::{base::message_result::PutMessageResult, store::running_flags::RunningFlags};
+use crate::{
+    base::message_result::PutMessageResult, hook::put_message_hook::BoxedPutMessageHook,
+    store::running_flags::RunningFlags,
+};
 
 pub mod commit_log;
 pub mod flush_manager_impl;
 pub mod mapped_file;
 
-#[allow(async_fn_in_trait)]
 #[trait_variant::make(MessageStore: Send)]
 pub trait RocketMQMessageStore: Clone + 'static {
     /// Load previously stored messages.
@@ -65,4 +69,8 @@ pub trait RocketMQMessageStore: Clone + 'static {
     fn get_running_flags(&self) -> &RunningFlags;
 
     fn is_shutdown(&self) -> bool;
+
+    fn get_put_message_hook_list(&self) -> Arc<parking_lot::RwLock<Vec<BoxedPutMessageHook>>>;
+
+    fn set_put_message_hook(&self, put_message_hook: BoxedPutMessageHook);
 }

rocketmq-store/src/log_file/commit_log.rs

@@ -15,7 +15,12 @@
  * limitations under the License.
  */
 
-use std::{cell::Cell, collections::HashMap, mem, sync::Arc};
+use std::{
+    cell::Cell,
+    collections::HashMap,
+    mem,
+    sync::{atomic::AtomicU64, Arc},
+};
 
 use bytes::{Buf, Bytes, BytesMut};
 use rocketmq_common::{
@@ -167,6 +172,7 @@ pub struct CommitLog {
     consume_queue_store: ConsumeQueueStore,
     flush_manager: Arc<tokio::sync::Mutex<DefaultFlushManager>>,
     //flush_manager: Arc<parking_lot::Mutex<DefaultFlushManager>>,
+    begin_time_in_lock: Arc<AtomicU64>,
 }
 
 impl CommitLog {
@@ -203,11 +209,7 @@ impl CommitLog {
                 mapped_file_queue,
                 store_checkpoint,
             ))),
-            /*flush_manager: Arc::new(parking_lot::Mutex::new(DefaultFlushManager::new(
-                message_store_config,
-                mapped_file_queue,
-                store_checkpoint,
-            ))),*/
+            begin_time_in_lock: Arc::new(AtomicU64::new(0)),
         }
     }
 }
@@ -319,6 +321,10 @@ impl CommitLog {
         msg.encoded_buff = Some(encoded_buff);
         let put_message_context = PutMessageContext::new(topic_queue_key);
         let lock = self.put_message_lock.lock().await;
+        self.begin_time_in_lock.store(
+            time_utils::get_current_millis(),
+            std::sync::atomic::Ordering::Release,
+        );
         let start_time = Instant::now();
         // Here settings are stored timestamp, in order to ensure an orderly global
         if !self.message_store_config.duplication_enable {
@@ -333,6 +339,13 @@ impl CommitLog {
 
         if mapped_file.is_none() {
             drop(lock);
+            error!(
+                "create mapped file error, topic: {}  clientAddr: {}",
+                msg.topic(),
+                msg.born_host()
+            );
+            self.begin_time_in_lock
+                .store(0, std::sync::atomic::Ordering::Release);
             return PutMessageResult::new_default(PutMessageStatus::CreateMappedFileFailed);
         }
 
@@ -353,6 +366,8 @@ impl CommitLog {
                     .mapped_file_queue
                     .get_last_mapped_file_mut_start_offset(0, true);
                 if mapped_file.is_none() {
+                    self.begin_time_in_lock
+                        .store(0, std::sync::atomic::Ordering::Release);
                     error!(
                         "create mapped file error, topic: {}  clientAddr: {}",
                         msg.topic(),
@@ -379,15 +394,21 @@ impl CommitLog {
             }
             AppendMessageStatus::MessageSizeExceeded
             | AppendMessageStatus::PropertiesSizeExceeded => {
+                self.begin_time_in_lock
+                    .store(0, std::sync::atomic::Ordering::Release);
                 PutMessageResult::new_append_result(PutMessageStatus::MessageIllegal, Some(result))
             }
             AppendMessageStatus::UnknownError => {
+                self.begin_time_in_lock
+                    .store(0, std::sync::atomic::Ordering::Release);
                 PutMessageResult::new_append_result(PutMessageStatus::UnknownError, Some(result))
             }
         };
         let elapsed_time_in_lock = start_time.elapsed().as_millis() as u64;
         drop(lock);
-        if elapsed_time_in_lock > 100 {
+        self.begin_time_in_lock
+            .store(0, std::sync::atomic::Ordering::Release);
+        if elapsed_time_in_lock > 500 {
             warn!(
                 "[NOTIFYME]putMessage in lock cost time(ms)={}, bodyLength={} \
                  AppendMessageResult={:?}",
@@ -873,6 +894,10 @@ impl CommitLog {
     pub fn check_self(&self) {
         self.mapped_file_queue.check_self();
     }
+
+    pub fn begin_time_in_lock(&self) -> &Arc<AtomicU64> {
+        &self.begin_time_in_lock
+    }
 }
 
 pub fn check_message_and_return_size(