diff --git a/broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java b/broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java index bc8400c19a2..9cdca3bac53 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/client/ProducerManager.java +++ b/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 channelMap = groupChannelTable.get(groupId); + if (channelMap != null) { + for (Map.Entry 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); } diff --git a/broker/src/main/java/org/apache/rocketmq/broker/processor/EndTransactionProcessor.java b/broker/src/main/java/org/apache/rocketmq/broker/processor/EndTransactionProcessor.java index 2be2e188023..b4d3f46531b 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/processor/EndTransactionProcessor.java +++ b/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; } diff --git a/broker/src/main/java/org/apache/rocketmq/broker/transaction/AbstractTransactionalMessageCheckListener.java b/broker/src/main/java/org/apache/rocketmq/broker/transaction/AbstractTransactionalMessageCheckListener.java index d1b77355b03..6a357a6e782 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/transaction/AbstractTransactionalMessageCheckListener.java +++ b/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 { diff --git a/broker/src/main/java/org/apache/rocketmq/broker/transaction/rocksdb/TransactionalMessageRocksDBService.java b/broker/src/main/java/org/apache/rocketmq/broker/transaction/rocksdb/TransactionalMessageRocksDBService.java index dbd3575d69c..b0ff35d45a7 100644 --- a/broker/src/main/java/org/apache/rocketmq/broker/transaction/rocksdb/TransactionalMessageRocksDBService.java +++ b/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 { diff --git a/broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java b/broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java index 451b0e044c7..39ffdf88ca9 100644 --- a/broker/src/test/java/org/apache/rocketmq/broker/client/ProducerManagerTest.java +++ b/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(); + } + } \ No newline at end of file diff --git a/broker/src/test/java/org/apache/rocketmq/broker/processor/EndTransactionProcessorTest.java b/broker/src/test/java/org/apache/rocketmq/broker/processor/EndTransactionProcessorTest.java index 42af2dbe47e..7828f006378 100644 --- a/broker/src/test/java/org/apache/rocketmq/broker/processor/EndTransactionProcessorTest.java +++ b/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 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)); diff --git a/broker/src/test/java/org/apache/rocketmq/broker/transaction/queue/TransactionalMessageUtilTest.java b/broker/src/test/java/org/apache/rocketmq/broker/transaction/queue/TransactionalMessageUtilTest.java index 722a306848e..bc6fd353669 100644 --- a/broker/src/test/java/org/apache/rocketmq/broker/transaction/queue/TransactionalMessageUtilTest.java +++ b/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 diff --git a/client/src/main/java/org/apache/rocketmq/client/impl/producer/DefaultMQProducerImpl.java b/client/src/main/java/org/apache/rocketmq/client/impl/producer/DefaultMQProducerImpl.java index 9ad5fcef4dc..1c1eec9a43c 100644 --- a/client/src/main/java/org/apache/rocketmq/client/impl/producer/DefaultMQProducerImpl.java +++ b/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) { diff --git a/common/src/main/java/org/apache/rocketmq/common/message/MessageConst.java b/common/src/main/java/org/apache/rocketmq/common/message/MessageConst.java index 77ab3f2cb9f..b8c95f8dc16 100644 --- a/common/src/main/java/org/apache/rocketmq/common/message/MessageConst.java +++ b/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); } } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/client/ProxyClientRemotingProcessor.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/client/ProxyClientRemotingProcessor.java index 10a8f3df50d..09317a2e500 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/service/client/ProxyClientRemotingProcessor.java +++ b/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 { diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/service/mqclient/ProxyClientRemotingProcessorTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/service/mqclient/ProxyClientRemotingProcessorTest.java index 441d3c04012..b7dd024a38a 100644 --- a/proxy/src/test/java/org/apache/rocketmq/proxy/service/mqclient/ProxyClientRemotingProcessorTest.java +++ b/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);