[ISSUE #1078]🚀Support send transaction message for client🍻 (#1088)

Author: mxsm

rocketmq-client/Cargo.toml

@@ -81,4 +81,8 @@ path = "examples/ordermessage/ordermessage_producer.rs"
 
 [[example]]
 name = "ordermessage-consumer"
-path = "examples/ordermessage/ordermessage_consumer.rs"
+path = "examples/ordermessage/ordermessage_consumer.rs"
+
+[[example]]
+name = "transaction-producer"
+path = "examples/transaction/transaction_producer.rs"

rocketmq-client/examples/transaction/transaction_producer.rs

@@ -0,0 +1,108 @@
+/*
+ * 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::any::Any;
+use std::collections::HashMap;
+use std::sync::atomic::AtomicI32;
+use std::sync::Arc;
+
+use parking_lot::Mutex;
+use rocketmq_client::producer::default_mq_producer::DefaultMQProducer;
+use rocketmq_client::producer::local_transaction_state::LocalTransactionState;
+use rocketmq_client::producer::mq_producer::MQProducer;
+use rocketmq_client::producer::transaction_listener::TransactionListener;
+use rocketmq_client::producer::transaction_mq_producer::TransactionMQProducer;
+use rocketmq_client::Result;
+use rocketmq_common::common::message::message_ext::MessageExt;
+use rocketmq_common::common::message::message_single::Message;
+use rocketmq_common::common::message::MessageTrait;
+use rocketmq_rust::rocketmq;
+
+pub const MESSAGE_COUNT: usize = 1;
+pub const PRODUCER_GROUP: &str = "please_rename_unique_group_name";
+pub const DEFAULT_NAMESRVADDR: &str = "127.0.0.1:9876";
+pub const TOPIC: &str = "TopicTest";
+pub const TAG: &str = "TagA";
+
+#[rocketmq::main]
+pub async fn main() -> Result<()> {
+    //init logger
+    rocketmq_common::log::init_logger();
+
+    // create a producer builder with default configuration
+    let builder = TransactionMQProducer::builder();
+
+    let mut producer = builder
+        .producer_group(PRODUCER_GROUP.to_string())
+        .name_server_addr(DEFAULT_NAMESRVADDR.to_string())
+        .topics(vec![TOPIC.to_string()])
+        .transaction_listener(TransactionListenerImpl::default())
+        .build();
+
+    producer.start().await?;
+
+    for _ in 0..10 {
+        let message = Message::with_tags(TOPIC, TAG, "Hello RocketMQ".as_bytes());
+        let send_result = producer
+            .send_message_in_transaction::<()>(message, None)
+            .await?;
+        println!("send result: {}", send_result);
+    }
+    let _ = tokio::signal::ctrl_c().await;
+    producer.shutdown().await;
+
+    Ok(())
+}
+
+struct TransactionListenerImpl {
+    local_trans: Arc<Mutex<HashMap<String, i32>>>,
+    transaction_index: AtomicI32,
+}
+
+impl Default for TransactionListenerImpl {
+    fn default() -> Self {
+        Self {
+            local_trans: Arc::new(Default::default()),
+            transaction_index: Default::default(),
+        }
+    }
+}
+
+impl TransactionListener for TransactionListenerImpl {
+    fn execute_local_transaction(
+        &self,
+        msg: &Message,
+        arg: Option<&(dyn Any + Send + Sync)>,
+    ) -> LocalTransactionState {
+        let value = self
+            .transaction_index
+            .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
+        let status = value % 3;
+        let mut guard = self.local_trans.lock();
+        guard.insert(msg.get_transaction_id().to_string(), status);
+        LocalTransactionState::Unknown
+    }
+
+    fn check_local_transaction(&self, msg: &MessageExt) -> LocalTransactionState {
+        let mut guard = self.local_trans.lock();
+        let status = guard.get(msg.get_transaction_id()).unwrap_or(&-1);
+        match status {
+            1 => LocalTransactionState::CommitMessage,
+            2 => LocalTransactionState::RollbackMessage,
+            _ => LocalTransactionState::Unknown,
+        }
+    }
+}

rocketmq-client/src/factory/mq_client_instance.rs

@@ -59,21 +59,21 @@ use crate::implementation::mq_admin_impl::MQAdminImpl;
 use crate::implementation::mq_client_api_impl::MQClientAPIImpl;
 use crate::producer::default_mq_producer::DefaultMQProducer;
 use crate::producer::default_mq_producer::ProducerConfig;
-use crate::producer::producer_impl::mq_producer_inner::MQProducerInner;
+use crate::producer::producer_impl::mq_producer_inner::MQProducerInnerImpl;
 use crate::producer::producer_impl::topic_publish_info::TopicPublishInfo;
 use crate::Result;
 
 const LOCK_TIMEOUT_MILLIS: u64 = 3000;
 
 pub struct MQClientInstance {
-    pub(crate) client_config: Arc<ClientConfig>,
+    pub(crate) client_config: ArcRefCellWrapper<ClientConfig>,
     pub(crate) client_id: String,
     boot_timestamp: u64,
     /**
      * The container of the producer in the current client. The key is the name of
      * producerGroup.
      */
-    producer_table: Arc<RwLock<HashMap<String, Box<dyn MQProducerInner>>>>,
+    producer_table: Arc<RwLock<HashMap<String, MQProducerInnerImpl>>>,
     /**
      * The container of the consumer in the current client. The key is the name of
      * consumer_group.
@@ -177,7 +177,7 @@ impl MQClientInstance {
     ) -> ArcRefCellWrapper<MQClientInstance> {
         let broker_addr_table = Arc::new(Default::default());
         let mut instance = ArcRefCellWrapper::new(MQClientInstance {
-            client_config: Arc::new(client_config.clone()),
+            client_config: ArcRefCellWrapper::new(client_config.clone()),
             client_id,
             boot_timestamp: get_current_millis(),
             producer_table: Arc::new(RwLock::new(HashMap::new())),
@@ -326,7 +326,7 @@ impl MQClientInstance {
 
     pub async fn shutdown(&mut self) {}
 
-    pub async fn register_producer(&mut self, group: &str, producer: impl MQProducerInner) -> bool {
+    pub async fn register_producer(&mut self, group: &str, producer: MQProducerInnerImpl) -> bool {
         if group.is_empty() {
             return false;
         }
@@ -335,7 +335,7 @@ impl MQClientInstance {
             warn!("the producer group[{}] exist already.", group);
             return false;
         }
-        producer_table.insert(group.to_string(), Box::new(producer));
+        producer_table.insert(group.to_string(), producer);
         true
     }
 
@@ -1055,6 +1055,11 @@ impl MQClientInstance {
         consumer_table.get(group).cloned()
     }
 
+    pub async fn select_producer(&self, group: &str) -> Option<MQProducerInnerImpl> {
+        let producer_table = self.producer_table.read().await;
+        producer_table.get(group).cloned()
+    }
+
     pub async fn unregister_consumer(&mut self, group: impl Into<String>) {
         self.unregister_client(None, Some(group.into())).await;
     }

rocketmq-client/src/hook/end_transaction_context.rs

@@ -14,4 +14,16 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-pub struct EndTransactionContext {}
+use rocketmq_common::common::message::message_single::Message;
+
+use crate::producer::local_transaction_state::LocalTransactionState;
+
+pub struct EndTransactionContext<'a> {
+    pub producer_group: String,
+    pub broker_addr: String,
+    pub message: &'a Message,
+    pub msg_id: String,
+    pub transaction_id: String,
+    pub transaction_state: LocalTransactionState,
+    pub from_transaction_check: bool,
+}

rocketmq-client/src/hook/send_message_context.rs

@@ -21,6 +21,7 @@ use std::sync::Arc;
 use rocketmq_common::common::message::message_enum::MessageType;
 use rocketmq_common::common::message::message_queue::MessageQueue;
 use rocketmq_common::common::message::MessageTrait;
+use rocketmq_common::ArcRefCellWrapper;
 
 use crate::implementation::communication_mode::CommunicationMode;
 use crate::producer::producer_impl::default_mq_producer_impl::DefaultMQProducerImpl;
@@ -38,7 +39,7 @@ pub struct SendMessageContext<'a> {
     pub exception: Option<Arc<Box<dyn Error + Send + Sync>>>,
     pub mq_trace_context: Option<Arc<Box<dyn std::any::Any + Send + Sync>>>,
     pub props: HashMap<String, String>,
-    pub producer: Option<DefaultMQProducerImpl>,
+    pub producer: Option<ArcRefCellWrapper<DefaultMQProducerImpl>>,
     pub msg_type: Option<MessageType>,
     pub namespace: Option<String>,
 }

rocketmq-client/src/implementation/client_remoting_processor.rs

@@ -20,6 +20,7 @@ use bytes::Bytes;
 use rocketmq_common::common::compression::compressor_factory::CompressorFactory;
 use rocketmq_common::common::message::message_ext::MessageExt;
 use rocketmq_common::common::message::MessageConst;
+use rocketmq_common::common::message::MessageTrait;
 use rocketmq_common::common::sys_flag::message_sys_flag::MessageSysFlag;
 use rocketmq_common::MessageAccessor::MessageAccessor;
 use rocketmq_common::MessageDecoder;
@@ -28,8 +29,10 @@ use rocketmq_common::WeakCellWrapper;
 use rocketmq_remoting::code::request_code::RequestCode;
 use rocketmq_remoting::code::response_code::ResponseCode;
 use rocketmq_remoting::net::channel::Channel;
+use rocketmq_remoting::protocol::header::check_transaction_state_request_header::CheckTransactionStateRequestHeader;
 use rocketmq_remoting::protocol::header::notify_consumer_ids_changed_request_header::NotifyConsumerIdsChangedRequestHeader;
 use rocketmq_remoting::protocol::header::reply_message_request_header::ReplyMessageRequestHeader;
+use rocketmq_remoting::protocol::namespace_util::NamespaceUtil;
 use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
 use rocketmq_remoting::runtime::connection_handler_context::ConnectionHandlerContext;
 use rocketmq_remoting::runtime::processor::RequestProcessor;
@@ -62,10 +65,26 @@ impl RequestProcessor for ClientRemotingProcessor {
         let request_code = RequestCode::from(request.code());
         info!("process_request: {:?}", request_code);
         match request_code {
+            RequestCode::CheckTransactionState => {
+                self.check_transaction_state(channel, ctx, request).await
+            }
+            RequestCode::ResetConsumerClientOffset => {
+                unimplemented!("ResetConsumerClientOffset")
+            }
+            RequestCode::GetConsumerStatusFromClient => {
+                unimplemented!("GetConsumerStatusFromClient")
+            }
+            RequestCode::GetConsumerRunningInfo => {
+                unimplemented!("GetConsumerRunningInfo")
+            }
+            RequestCode::ConsumeMessageDirectly => {
+                unimplemented!("ConsumeMessageDirectly")
+            }
             RequestCode::PushReplyMessageToClient => self.receive_reply_message(ctx, request).await,
             RequestCode::NotifyConsumerIdsChanged => {
                 self.notify_consumer_ids_changed(channel, ctx, request)
             }
+
             _ => {
                 info!("Unknown request code: {:?}", request_code);
                 Ok(None)
@@ -190,4 +209,64 @@ impl ClientRemotingProcessor {
         }
         Ok(None)
     }
+
+    async fn check_transaction_state(
+        &mut self,
+        channel: Channel,
+        ctx: ConnectionHandlerContext,
+        mut request: RemotingCommand,
+    ) -> Result<Option<RemotingCommand>> {
+        let request_header = request
+            .decode_command_custom_header::<CheckTransactionStateRequestHeader>()
+            .unwrap();
+        let message_ext = MessageDecoder::decode(
+            request.get_body_mut().unwrap(),
+            true,
+            true,
+            false,
+            false,
+            false,
+        );
+        if let Some(mut message_ext) = message_ext {
+            if let Some(mut client_instance) = self.client_instance.upgrade() {
+                if let Some(ref namespace) = client_instance.client_config.get_namespace() {
+                    let topic = NamespaceUtil::without_namespace_with_namespace(
+                        message_ext.get_topic(),
+                        client_instance
+                            .client_config
+                            .get_namespace()
+                            .unwrap_or_default()
+                            .as_str(),
+                    );
+                    message_ext.set_topic(topic.as_str());
+                }
+                let transaction_id =
+                    message_ext.get_property(MessageConst::PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX);
+                if let Some(transaction_id) = transaction_id {
+                    if !transaction_id.is_empty() {
+                        message_ext.set_transaction_id(transaction_id.as_str());
+                    }
+                }
+                let group = message_ext.get_property(MessageConst::PROPERTY_PRODUCER_GROUP);
+                if let Some(group) = group {
+                    let producer = client_instance.select_producer(&group).await;
+                    if let Some(producer) = producer {
+                        let addr = channel.remote_address().to_string();
+                        producer.check_transaction_state(
+                            addr.as_str(),
+                            message_ext,
+                            request_header,
+                        );
+                    } else {
+                        warn!("checkTransactionState, pick producer group failed");
+                    }
+                } else {
+                    warn!("checkTransactionState, pick producer group failed");
+                }
+            }
+        } else {
+            warn!("checkTransactionState, decode message failed");
+        };
+        Ok(None)
+    }
 }

rocketmq-client/src/implementation/mq_client_api_impl.rs

@@ -46,6 +46,7 @@ use rocketmq_remoting::protocol::body::response::lock_batch_response_body::LockB
 use rocketmq_remoting::protocol::body::unlock_batch_request_body::UnlockBatchRequestBody;
 use rocketmq_remoting::protocol::header::client_request_header::GetRouteInfoRequestHeader;
 use rocketmq_remoting::protocol::header::consumer_send_msg_back_request_header::ConsumerSendMsgBackRequestHeader;
+use rocketmq_remoting::protocol::header::end_transaction_request_header::EndTransactionRequestHeader;
 use rocketmq_remoting::protocol::header::get_consumer_listby_group_request_header::GetConsumerListByGroupRequestHeader;
 use rocketmq_remoting::protocol::header::heartbeat_request_header::HeartbeatRequestHeader;
 use rocketmq_remoting::protocol::header::lock_batch_mq_request_header::LockBatchMqRequestHeader;
@@ -1130,4 +1131,21 @@ impl MQClientAPIImpl {
             ))
         }
     }
+
+    pub async fn end_transaction_oneway(
+        &mut self,
+        addr: &str,
+        request_header: EndTransactionRequestHeader,
+        remark: String,
+        timeout_millis: u64,
+    ) -> Result<()> {
+        let request =
+            RemotingCommand::create_request_command(RequestCode::EndTransaction, request_header)
+                .set_remark(Some(remark));
+
+        self.remoting_client
+            .invoke_oneway(addr.to_string(), request, timeout_millis)
+            .await;
+        Ok(())
+    }
 }

rocketmq-client/src/lib.rs

@@ -16,6 +16,8 @@
  */
 #![allow(dead_code)]
 #![allow(unused_variables)]
+#![recursion_limit = "256"]
+
 extern crate core;
 
 use crate::error::MQClientError;

rocketmq-client/src/producer.rs

@@ -28,4 +28,6 @@ pub mod send_callback;
 pub mod send_result;
 pub mod send_status;
 pub mod transaction_listener;
+pub mod transaction_mq_produce_builder;
+pub mod transaction_mq_producer;
 pub mod transaction_send_result;