[Refactor♻️]Refactor MessageQueueOpContext

 #3386

Author: Lori-Shu

rocketmq-broker/src/transaction/queue/default_transactional_message_service.rs

@@ -131,9 +131,9 @@ where
         let mut send_map = HashMap::<i32, Message>::new();
         let delete_context_mutex_guard = self.delete_context.lock().await;
         for (queue_id, mq_context) in delete_context_mutex_guard.iter() {
-            if mq_context.get_total_size() <= 0
-                || mq_context.context_queue().is_empty().await
-                || (mq_context.get_total_size() < max_size as i32
+            if mq_context.get_total_size().await <= 0
+                || mq_context.is_empty().await
+                || (mq_context.get_total_size().await < max_size as u32
                     && (start_time as i64 - mq_context.get_last_write_timestamp().await as i64)
                         < interval as i64)
             {
@@ -145,7 +145,7 @@ where
             }
             send_map.insert(*queue_id, op_message.unwrap());
             first_timestamp = first_timestamp.min(mq_context.get_last_write_timestamp().await);
-            if mq_context.get_total_size() >= max_size as i32 {
+            if mq_context.get_total_size().await >= max_size as u32 {
                 over_size = true;
             }
         }
@@ -172,8 +172,8 @@ where
         more_data: Option<String>,
     ) -> Option<Message> {
         let topic = TransactionalMessageUtil::build_op_topic();
-        let delete_context = self.delete_context.lock().await;
-        let mq_context = delete_context.get(&queue_id)?;
+        let mut delete_context = self.delete_context.lock().await;
+        let mq_context = delete_context.get_mut(&queue_id)?;
 
         let more_data_length = if let Some(ref data) = more_data {
             data.len()
@@ -187,7 +187,7 @@ where
             .broker_config()
             .transaction_op_msg_max_size as usize;
         if length < max_size {
-            let sz = mq_context.get_total_size() as usize;
+            let sz = mq_context.get_total_size().await as usize;
             if sz > max_size || length + sz > max_size {
                 length = max_size + 100;
             } else {
@@ -201,12 +201,12 @@ where
             sb.push_str(&data);
         }
 
-        while !mq_context.context_queue().is_empty().await {
+        while !mq_context.is_empty().await {
             if sb.len() >= max_size {
                 break;
             }
             {
-                if let Some(data) = mq_context.context_queue().try_poll().await {
+                if let Ok(data) = mq_context.pull().await {
                     sb.push_str(&data);
                 }
             }
@@ -1069,17 +1069,16 @@ where
         );
         let len = data.len();
         let res = mq_context
-            .context_queue()
             .offer(data.clone(), Duration::from_millis(100))
             .await;
-        if res {
-            let total_size = mq_context.total_size_add_and_get(len as i32);
+        if res.is_ok() {
+            let total_size = mq_context.total_size_add_and_get(len as u32).await;
             if total_size
                 > self
                     .transactional_message_bridge
                     .broker_runtime_inner
                     .broker_config()
-                    .transaction_op_msg_max_size
+                    .transaction_op_msg_max_size as u32
             {
                 if let Some(batch_service) = &self.transactional_op_batch_service {
                     batch_service.wakeup();

rocketmq-broker/src/transaction/queue/message_queue_op_context.rs

@@ -1,59 +1,71 @@
-/*
- * 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::atomic::AtomicI32;
+use std::sync::atomic::AtomicU32;
+use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
-use rocketmq_rust::RocketMQBlockingQueue;
-use tokio::sync::Mutex;
+use anyhow::Error;
+use rocketmq_error::Result;
+use tokio::sync::mpsc;
+use tokio::time;
 
 pub struct MessageQueueOpContext {
-    total_size: AtomicI32,
-    last_write_timestamp: Arc<Mutex<u64>>,
-    context_queue: RocketMQBlockingQueue<String>,
+    total_size: AtomicU32,
+    last_write_timestamp: AtomicU64,
+    context_queue: Arc<mpsc::UnboundedSender<String>>,
+    context_receiver: mpsc::UnboundedReceiver<String>,
+    queue_capacity: usize,
 }
 
 impl MessageQueueOpContext {
     pub fn new(timestamp: u64, queue_length: usize) -> Self {
+        let unbounded_channel = mpsc::unbounded_channel::<String>();
         MessageQueueOpContext {
-            total_size: AtomicI32::new(0),
-            last_write_timestamp: Arc::new(Mutex::new(timestamp)),
-            context_queue: RocketMQBlockingQueue::new(queue_length),
+            total_size: AtomicU32::new(0),
+            last_write_timestamp: AtomicU64::new(timestamp),
+            context_queue: Arc::new(unbounded_channel.0),
+            context_receiver: unbounded_channel.1,
+            queue_capacity: queue_length,
         }
     }
 
-    pub fn get_total_size(&self) -> i32 {
+    pub async fn get_total_size(&self) -> u32 {
         self.total_size.load(Ordering::Relaxed)
     }
 
-    pub fn total_size_add_and_get(&self, delta: i32) -> i32 {
+    pub async fn total_size_add_and_get(&self, delta: u32) -> u32 {
         self.total_size.fetch_add(delta, Ordering::AcqRel) + delta
     }
 
     pub async fn get_last_write_timestamp(&self) -> u64 {
-        *self.last_write_timestamp.lock().await
+        self.last_write_timestamp.load(Ordering::Relaxed)
     }
 
     pub async fn set_last_write_timestamp(&self, timestamp: u64) {
-        let mut last_timestamp = self.last_write_timestamp.lock().await;
-        *last_timestamp = timestamp;
+        self.last_write_timestamp
+            .store(timestamp, Ordering::Release);
     }
 
-    pub fn context_queue(&self) -> &RocketMQBlockingQueue<String> {
-        &self.context_queue
+    pub async fn push(&self, msg: String) -> Result<()> {
+        if self.context_receiver.len() > self.queue_capacity {
+            return Err(anyhow::Error::msg("queue is full".to_string()));
+        }
+        self.context_queue
+            .send(msg)
+            .map_err(|e| anyhow::Error::new(e))
+    }
+    pub async fn offer(&self, item: String, timeout: std::time::Duration) -> Result<()> {
+        time::timeout(timeout, self.push(item)).await.unwrap()
+    }
+    pub async fn pull(&mut self) -> Result<String> {
+        if let Some(item) = self.context_receiver.recv().await {
+            return Ok(item);
+        }
+        Err(Error::msg("pull failed, queue is empty".to_string()))
+    }
+    pub async fn is_empty(&self) -> bool {
+        if self.context_receiver.len() == 0 {
+            return true;
+        }
+        false
     }
 }