[ISSUE apache#9347] Optimization of Prioritizing Original Sender in Transactional Message Checkbacks

Author: CLFutureX

broker/src/main/java/org/apache/rocketmq/broker/client/ClientChannelInfo.java

@@ -61,6 +61,14 @@ public void setLastUpdateTimestamp(long lastUpdateTimestamp) {
         this.lastUpdateTimestamp = lastUpdateTimestamp;
     }
 
+    public boolean isActive() {
+        return getChannel().isActive();
+    }
+
+    public boolean isWritable() {
+        return getChannel().isWritable();
+    }
+
     @Override
     public int hashCode() {
         final int prime = 31;

broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java

@@ -206,13 +206,17 @@ public synchronized void unregisterProducer(final String group, final ClientChan
     }
 
     public Channel getAvailableChannel(String groupId) {
+        return getAvailableChannel(groupId, null);
+    }
+
+    public Channel getAvailableChannel(String groupId, String clientId) {
         if (groupId == null) {
             return null;
         }
-        List<Channel> channelList;
+        List<ClientChannelInfo> channelList;
         ConcurrentHashMap<Channel, ClientChannelInfo> channelClientChannelInfoHashMap = groupChannelTable.get(groupId);
         if (channelClientChannelInfoHashMap != null) {
-            channelList = new ArrayList<>(channelClientChannelInfoHashMap.keySet());
+            channelList = new ArrayList<>(channelClientChannelInfoHashMap.values());
         } else {
             log.warn("Check transaction failed, channel table is empty. groupId={}", groupId);
             return null;
@@ -224,25 +228,35 @@ public Channel getAvailableChannel(String groupId) {
             return null;
         }
 
-        Channel lastActiveChannel = null;
+        Channel firstChannel = null;
+        Channel secondChannel = null;
 
         int index = positiveAtomicCounter.incrementAndGet() % size;
-        Channel channel = channelList.get(index);
+        ClientChannelInfo clientChannel = channelList.get(index);
         int count = 0;
-        boolean isOk = channel.isActive() && channel.isWritable();
+        boolean isOk = clientChannel.isActive() && clientChannel.isWritable();
+        boolean isSendClient = clientChannel.getClientId().equals(clientId);
+
         while (count++ < GET_AVAILABLE_CHANNEL_RETRY_COUNT) {
+            if (isOk && isSendClient) {
+                return clientChannel.getChannel();
+            }
             if (isOk) {
-                return channel;
+                firstChannel = clientChannel.getChannel();
+            }
+            if (clientChannel.isActive()) {
+                secondChannel = clientChannel.getChannel();
             }
-            if (channel.isActive()) {
-                lastActiveChannel = channel;
+            if (clientChannel.isActive() && isSendClient) {
+                secondChannel = clientChannel.getChannel();
             }
             index = (++index) % size;
-            channel = channelList.get(index);
-            isOk = channel.isActive() && channel.isWritable();
+            clientChannel = channelList.get(index);
+            isOk = clientChannel.isActive() && clientChannel.isWritable();
+            isSendClient = clientChannel.getClientId().equals(clientId);
         }
 
-        return lastActiveChannel;
+        return firstChannel != null ? firstChannel : secondChannel;
     }
 
     public Channel findChannel(String clientId) {

broker/src/main/java/org/apache/rocketmq/broker/transaction/AbstractTransactionalMessageCheckListener.java

@@ -61,7 +61,8 @@ public void sendCheckMessage(MessageExt msgExt) throws Exception {
         msgExt.setQueueId(Integer.parseInt(msgExt.getUserProperty(MessageConst.PROPERTY_REAL_QUEUE_ID)));
         msgExt.setStoreSize(0);
         String groupId = msgExt.getProperty(MessageConst.PROPERTY_PRODUCER_GROUP);
-        Channel channel = brokerController.getProducerManager().getAvailableChannel(groupId);
+        String clientId = msgExt.getProperty(MessageConst.PROPERTY_INSTANCE_ID);
+        Channel channel = brokerController.getProducerManager().getAvailableChannel(groupId, clientId);
         if (channel != null) {
             brokerController.getBroker2Client().checkProducerTransactionState(groupId, channel, checkTransactionStateRequestHeader, msgExt);
         } else {

client/src/main/java/org/apache/rocketmq/client/impl/producer/DefaultMQProducerImpl.java

@@ -1392,6 +1392,7 @@ public TransactionSendResult sendMessageInTransaction(final Message msg,
         SendResult sendResult = null;
         MessageAccessor.putProperty(msg, MessageConst.PROPERTY_TRANSACTION_PREPARED, "true");
         MessageAccessor.putProperty(msg, MessageConst.PROPERTY_PRODUCER_GROUP, this.defaultMQProducer.getProducerGroup());
+        MessageAccessor.putProperty(msg, MessageConst.PROPERTY_INSTANCE_ID, this.mQClientFactory.getClientId());
         try {
             sendResult = this.send(msg);
         } catch (Exception e) {