[ISSUE #9791] Prefer original producer when checking transaction state

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

@@ -316,6 +316,33 @@ public Channel getAvailableChannel(String groupId) {
         return lastActiveChannel;
     }
 
+    /**
+     * Get an available channel for the given group, preferring the producer that originally sent the message.
+     * Falls back to round-robin if the preferred producer is not found or not available.
+     *
+     * @param groupId producer group
+     * @param preferredClientId the clientId of the original producer (from half message properties), may be null
+     * @return an available channel, or null if none found
+     */
+    public Channel getAvailableChannel(String groupId, String preferredClientId) {
+        if (groupId == null) {
+            return null;
+        }
+        if (preferredClientId != null) {
+            ConcurrentMap<Channel, ClientChannelInfo> channelMap = groupChannelTable.get(groupId);
+            if (channelMap != null) {
+                for (Map.Entry<Channel, ClientChannelInfo> entry : channelMap.entrySet()) {
+                    if (preferredClientId.equals(entry.getValue().getClientId())
+                        && entry.getKey().isActive() && entry.getKey().isWritable()) {
+                        return entry.getKey();
+                    }
+                }
+            }
+        }
+        // Fall back to round-robin selection
+        return getAvailableChannel(groupId);
+    }
+
     public Channel findChannel(String clientId) {
         return clientChannelTable.get(clientId);
     }

broker/src/main/java/org/apache/rocketmq/broker/processor/EndTransactionProcessor.java

@@ -291,6 +291,9 @@ private MessageExtBrokerInner endMessageTransaction(MessageExt msgExt) {
         MessageAccessor.setProperties(msgInner, MessageDecoder.string2messageProperties(MessageDecoder.messageProperties2String(msgExt.getProperties())));
         MessageAccessor.clearProperty(msgInner, MessageConst.PROPERTY_REAL_TOPIC);
         MessageAccessor.clearProperty(msgInner, MessageConst.PROPERTY_REAL_QUEUE_ID);
+        // Clear the internal routing hint before encoding so it never reaches the wire data
+        // of the committed message that consumers observe.
+        MessageAccessor.clearProperty(msgInner, MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID);
         msgInner.setPropertiesString(MessageDecoder.messageProperties2String(msgInner.getProperties()));
         return msgInner;
     }

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 producerClientId = msgExt.getUserProperty(MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID);
+        Channel channel = brokerController.getProducerManager().getAvailableChannel(groupId, producerClientId);
         if (channel != null) {
             brokerController.getBroker2Client().checkProducerTransactionState(groupId, channel, checkTransactionStateRequestHeader, msgExt);
         } else {

broker/src/main/java/org/apache/rocketmq/broker/transaction/rocksdb/TransactionalMessageRocksDBService.java

@@ -234,7 +234,8 @@ private void sendCheckMessage(MessageExt msgExt) {
             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 producerClientId = msgExt.getUserProperty(MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID);
+            Channel channel = brokerController.getProducerManager().getAvailableChannel(groupId, producerClientId);
             if (channel != null) {
                 brokerController.getBroker2Client().checkProducerTransactionState(groupId, channel, checkTransactionStateRequestHeader, msgExt);
             } else {

broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java

@@ -225,4 +225,44 @@ public void testGetAvailableChannel() {
         assertThat(c).isNull();
     }
 
+    @Test
+    public void testGetAvailableChannelWithPreferredClientId() {
+        producerManager.registerProducer(group, clientInfo);
+        when(channel.isActive()).thenReturn(true);
+        when(channel.isWritable()).thenReturn(true);
+
+        // Match: preferred clientId matches registered producer
+        Channel c = producerManager.getAvailableChannel(group, "clientId");
+        assertThat(c).isSameAs(channel);
+    }
+
+    @Test
+    public void testGetAvailableChannelWithPreferredClientIdNotFound() {
+        producerManager.registerProducer(group, clientInfo);
+        when(channel.isActive()).thenReturn(true);
+        when(channel.isWritable()).thenReturn(true);
+
+        // No match: falls back to round-robin (returns some channel from group)
+        Channel c = producerManager.getAvailableChannel(group, "nonExistentClientId");
+        assertThat(c).isNotNull(); // should fall back to round-robin
+    }
+
+    @Test
+    public void testGetAvailableChannelWithNullPreferredClientId() {
+        producerManager.registerProducer(group, clientInfo);
+        when(channel.isActive()).thenReturn(true);
+        when(channel.isWritable()).thenReturn(true);
+
+        // null clientId: should behave exactly like original getAvailableChannel
+        Channel c = producerManager.getAvailableChannel(group, null);
+        assertThat(c).isNotNull();
+    }
+
+    @Test
+    public void testGetAvailableChannelWithNullGroupId() {
+        // null groupId with non-null preferredClientId should return null (no NPE)
+        Channel c = producerManager.getAvailableChannel(null, "someClientId");
+        assertThat(c).isNull();
+    }
+
 }

broker/src/test/java/org/apache/rocketmq/broker/processor/EndTransactionProcessorTest.java

@@ -45,6 +45,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Spy;
 import org.mockito.junit.MockitoJUnitRunner;
@@ -53,6 +54,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 @RunWith(MockitoJUnitRunner.class)
@@ -110,6 +112,29 @@ public void testProcessRequest() throws RemotingCommandException {
         assertThat(brokerController.getBrokerStatsManager().getStatsItem(Stats.TOPIC_PUT_SIZE, TOPIC).getValue().sum()).isEqualTo(1L);
     }
 
+    @Test
+    public void testProcessRequest_CommitStripsProducerRoutingHintFromWireData() throws RemotingCommandException {
+        MessageExt halfMessage = createDefaultMessageExt();
+        MessageAccessor.putProperty(halfMessage, MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID, "127.0.0.1@12345");
+        OperationResult result = new OperationResult();
+        result.setPrepareMessage(halfMessage);
+        result.setResponseCode(ResponseCode.SUCCESS);
+        when(transactionMsgService.commitMessage(any(EndTransactionRequestHeader.class))).thenReturn(result);
+        when(messageStore.putMessage(any(MessageExtBrokerInner.class)))
+            .thenReturn(new PutMessageResult(PutMessageStatus.PUT_OK, createAppendMessageResult(AppendMessageStatus.PUT_OK)));
+        RemotingCommand request = createEndTransactionMsgCommand(MessageSysFlag.TRANSACTION_COMMIT_TYPE, false);
+
+        RemotingCommand response = endTransactionProcessor.processRequest(handlerContext, request);
+
+        assertThat(response.getCode()).isEqualTo(ResponseCode.SUCCESS);
+        ArgumentCaptor<MessageExtBrokerInner> captor = ArgumentCaptor.forClass(MessageExtBrokerInner.class);
+        verify(messageStore).putMessage(captor.capture());
+        MessageExtBrokerInner committed = captor.getValue();
+        // The hint must be stripped from both the property map and the encoded wire data.
+        assertThat(committed.getProperty(MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID)).isNull();
+        assertThat(committed.getPropertiesString()).doesNotContain(MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID);
+    }
+
     @Test
     public void testProcessRequest_CheckMessage() throws RemotingCommandException {
         when(transactionMsgService.commitMessage(any(EndTransactionRequestHeader.class))).thenReturn(createResponse(ResponseCode.SUCCESS));

broker/src/test/java/org/apache/rocketmq/broker/transaction/queue/TransactionalMessageUtilTest.java

@@ -39,6 +39,7 @@ public void testBuildTransactionalMessageFromHalfMessage() {
         halfMessage.setTransactionId("tranId");
         MessageAccessor.putProperty(halfMessage, MessageConst.PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX, "tranId");
         MessageAccessor.putProperty(halfMessage, MessageConst.PROPERTY_PRODUCER_GROUP, "trans-producer-grp");
+        MessageAccessor.putProperty(halfMessage, MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID, "127.0.0.1@12345");
 
         MessageExtBrokerInner msgExtInner = TransactionalMessageUtil.buildTransactionalMessageFromHalfMessage(halfMessage);
 
@@ -50,6 +51,9 @@ public void testBuildTransactionalMessageFromHalfMessage() {
         assertEquals(msgExtInner.getMsgId(), halfMessage.getMsgId());
         assertTrue(MessageSysFlag.check(msgExtInner.getSysFlag(), MessageSysFlag.TRANSACTION_PREPARED_TYPE));
         assertEquals(msgExtInner.getProperty(MessageConst.PROPERTY_PRODUCER_GROUP), halfMessage.getProperty(MessageConst.PROPERTY_PRODUCER_GROUP));
+        // The rebuilt message is still a half message (PREPARED=true re-enters the half topic on the
+        // target broker), so the producer routing hint must survive for transaction check routing.
+        assertEquals("127.0.0.1@12345", msgExtInner.getProperty(MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID));
     }
 
     @Test

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

@@ -1433,6 +1433,7 @@ public void sendOneway(Message msg, MessageQueueSelector selector, Object arg)
     public TransactionSendResult sendMessageInTransaction(final Message msg,
         final TransactionListener localTransactionListener, final Object arg)
         throws MQClientException {
+        this.makeSureStateOK();
         TransactionListener transactionListener = getCheckListener();
         if (null == localTransactionListener && null == transactionListener) {
             throw new MQClientException("tranExecutor is null", null);
@@ -1444,6 +1445,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_TRANSACTION_PRODUCER_CLIENT_ID, this.mQClientFactory.getClientId());
         try {
             sendResult = this.send(msg);
         } catch (Exception e) {

common/src/main/java/org/apache/rocketmq/common/message/MessageConst.java

@@ -52,6 +52,7 @@ public class MessageConst {
     public static final String PROPERTY_TRANSACTION_PREPARED_QUEUE_OFFSET = "TRAN_PREPARED_QUEUE_OFFSET";
     public static final String PROPERTY_TRANSACTION_ID = "__transactionId__";
     public static final String PROPERTY_TRANSACTION_CHECK_TIMES = "TRANSACTION_CHECK_TIMES";
+    public static final String PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID = "__TXN_PRODUCER_CID__";
     public static final String PROPERTY_INSTANCE_ID = "INSTANCE_ID";
     public static final String PROPERTY_CORRELATION_ID = "CORRELATION_ID";
     public static final String PROPERTY_MESSAGE_REPLY_TO_CLIENT = "REPLY_TO_CLIENT";
@@ -173,5 +174,6 @@ public class MessageConst {
         STRING_HASH_SET.add(PROPERTY_CRC32);
         STRING_HASH_SET.add(PROPERTY_PRIORITY);
         STRING_HASH_SET.add(PROPERTY_LITE_TOPIC);
+        STRING_HASH_SET.add(PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID);
     }
 }

proxy/src/main/java/org/apache/rocketmq/proxy/service/client/ProxyClientRemotingProcessor.java

@@ -67,11 +67,12 @@ public RemotingCommand checkTransactionState(ChannelHandlerContext ctx,
         if (messageExt != null) {
             final String group = messageExt.getProperty(MessageConst.PROPERTY_PRODUCER_GROUP);
             if (group != null) {
+                final String producerClientId = messageExt.getUserProperty(MessageConst.PROPERTY_TRANSACTION_PRODUCER_CLIENT_ID);
                 CheckTransactionStateRequestHeader requestHeader =
                     (CheckTransactionStateRequestHeader) request.decodeCommandCustomHeader(CheckTransactionStateRequestHeader.class);
                 request.writeCustomHeader(requestHeader);
                 request.addExtField(ProxyUtils.BROKER_ADDR, NetworkUtil.socketAddress2String(ctx.channel().remoteAddress()));
-                Channel channel = this.producerManager.getAvailableChannel(group);
+                Channel channel = this.producerManager.getAvailableChannel(group, producerClientId);
                 if (channel != null) {
                     channel.writeAndFlush(request);
                 } else {

proxy/src/test/java/org/apache/rocketmq/proxy/service/mqclient/ProxyClientRemotingProcessorTest.java

@@ -56,6 +56,7 @@
 import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -85,7 +86,7 @@ public void testTransactionCheck() throws Exception {
 
         GrpcClientChannel grpcClientChannel = new GrpcClientChannel(proxyRelayService, grpcClientSettingsManager, null,
             ProxyContext.create().setRemoteAddress("127.0.0.1:8888").setLocalAddress("127.0.0.1:10911"), "clientId");
-        when(producerManager.getAvailableChannel(anyString()))
+        when(producerManager.getAvailableChannel(anyString(), nullable(String.class)))
             .thenReturn(grpcClientChannel);
 
         ProxyClientRemotingProcessor processor = new ProxyClientRemotingProcessor(producerManager, null);