From a998906f3c62179f72bb418350a32bdc5b08b44a Mon Sep 17 00:00:00 2001 From: liuhy Date: Sat, 25 Jul 2026 23:18:56 -0700 Subject: [PATCH] [ISSUE #10577] Keep reentrant POP receives broker sticky --- .../v2/consumer/ReceiveMessageActivity.java | 17 +++- .../consumer/ReceiveMessageActivityTest.java | 96 ++++++++++++++++--- 2 files changed, 97 insertions(+), 16 deletions(-) diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivity.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivity.java index 39f2995d6fc..8f8775e1cb2 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivity.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivity.java @@ -142,7 +142,9 @@ public void receiveMessage(ProxyContext ctx, ReceiveMessageRequest request, CompletableFuture popFuture = this.messagingProcessor.popMessage( ctx, new ReceiveMessageQueueSelector( - request.getMessageQueue().getBroker().getName() + request.getMessageQueue().getBroker().getName(), + // Reentrant FIFO POP must stay on the original broker instead of falling back. + fifo && request.hasAttemptId() ), group, topic, @@ -216,9 +218,19 @@ protected ReceiveMessageResponseStreamWriter createWriter(ProxyContext ctx, protected static class ReceiveMessageQueueSelector implements QueueSelector { private final String brokerName; + private final boolean brokerSticky; public ReceiveMessageQueueSelector(String brokerName) { + this(brokerName, false); + } + + /** + * @param brokerSticky if true, keep the target broker sticky and return null when the + * target broker is unavailable instead of falling back to another broker. + */ + public ReceiveMessageQueueSelector(String brokerName, boolean brokerSticky) { this.brokerName = brokerName; + this.brokerSticky = brokerSticky; } @Override @@ -229,6 +241,9 @@ public AddressableMessageQueue select(ProxyContext ctx, MessageQueueView message if (StringUtils.isNotBlank(brokerName)) { addressableMessageQueue = messageQueueSelector.getQueueByBrokerName(brokerName); + if (addressableMessageQueue != null || brokerSticky) { + return addressableMessageQueue; + } } if (addressableMessageQueue == null) { diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivityTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivityTest.java index 5341c259761..9b89a8b036e 100644 --- a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivityTest.java +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/consumer/ReceiveMessageActivityTest.java @@ -18,6 +18,7 @@ package org.apache.rocketmq.proxy.grpc.v2.consumer; import apache.rocketmq.v2.Code; +import apache.rocketmq.v2.Broker; import apache.rocketmq.v2.ClientType; import apache.rocketmq.v2.FilterExpression; import apache.rocketmq.v2.FilterType; @@ -26,6 +27,7 @@ import apache.rocketmq.v2.ReceiveMessageResponse; import apache.rocketmq.v2.Resource; import apache.rocketmq.v2.Settings; +import apache.rocketmq.v2.Subscription; import com.google.protobuf.Duration; import com.google.protobuf.util.Durations; import io.grpc.stub.ServerCallStreamObserver; @@ -50,6 +52,7 @@ import org.apache.rocketmq.proxy.common.ProxyContext; import org.apache.rocketmq.proxy.config.ConfigurationManager; import org.apache.rocketmq.proxy.grpc.v2.BaseActivityTest; +import org.apache.rocketmq.proxy.processor.QueueSelector; import org.apache.rocketmq.proxy.service.route.AddressableMessageQueue; import org.apache.rocketmq.proxy.service.route.MessageQueueView; import org.apache.rocketmq.remoting.protocol.route.BrokerData; @@ -61,6 +64,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; @@ -422,6 +427,53 @@ public void testReceiveMessage() { assertEquals(Code.MESSAGE_NOT_FOUND, getResponseCodeFromReceiveMessageResponseList(responseArgumentCaptor.getAllValues())); } + @Test + public void testReceiveFifoMessageWithAttemptIdUsesStickyBrokerSelector() throws Exception { + StreamObserver receiveStreamObserver = mock(ServerCallStreamObserver.class); + doNothing().when(receiveStreamObserver).onNext(any()); + + when(this.grpcClientSettingsManager.getClientSettings(any())).thenReturn(Settings.newBuilder() + .setSubscription(Subscription.newBuilder().setFifo(true).build()) + .build()); + + PopResult popResult = new PopResult(PopStatus.NO_NEW_MSG, Collections.emptyList()); + ArgumentCaptor queueSelectorCaptor = ArgumentCaptor.forClass(QueueSelector.class); + when(this.messagingProcessor.popMessage( + any(), + queueSelectorCaptor.capture(), + anyString(), + anyString(), + anyInt(), + anyLong(), + anyLong(), + anyInt(), + any(), + eq(true), + any(), + eq("attempt-id"), + anyLong())).thenReturn(CompletableFuture.completedFuture(popResult)); + + this.receiveMessageActivity.receiveMessage( + createContext(), + ReceiveMessageRequest.newBuilder() + .setGroup(Resource.newBuilder().setName(CONSUMER_GROUP).build()) + .setMessageQueue(MessageQueue.newBuilder() + .setTopic(Resource.newBuilder().setName(TOPIC).build()) + .setBroker(Broker.newBuilder().setName("missing-broker").build()) + .build()) + .setAttemptId("attempt-id") + .setAutoRenew(true) + .setFilterExpression(FilterExpression.newBuilder() + .setType(FilterType.TAG) + .setExpression("*") + .build()) + .build(), + receiveStreamObserver + ); + + assertNull(queueSelectorCaptor.getValue().select(ProxyContext.create(), createMessageQueueView())); + } + private Code getResponseCodeFromReceiveMessageResponseList(List responseList) { for (ReceiveMessageResponse response : responseList) { if (response.hasStatus()) { @@ -433,6 +485,34 @@ private Code getResponseCodeFromReceiveMessageResponseList(List queueDatas = new ArrayList<>(); for (int i = 0; i < 2; i++) { @@ -456,20 +536,6 @@ public void testReceiveMessageQueueSelector() throws Exception { } topicRouteData.setBrokerDatas(brokerDatas); - MessageQueueView messageQueueView = new MessageQueueView(TOPIC, topicRouteData, null); - ReceiveMessageActivity.ReceiveMessageQueueSelector selector = new ReceiveMessageActivity.ReceiveMessageQueueSelector(""); - - AddressableMessageQueue firstSelect = selector.select(ProxyContext.create(), messageQueueView); - AddressableMessageQueue secondSelect = selector.select(ProxyContext.create(), messageQueueView); - AddressableMessageQueue thirdSelect = selector.select(ProxyContext.create(), messageQueueView); - - assertEquals(firstSelect, thirdSelect); - assertNotEquals(firstSelect, secondSelect); - - for (int i = 0; i < 2; i++) { - ReceiveMessageActivity.ReceiveMessageQueueSelector selectorBrokerName = - new ReceiveMessageActivity.ReceiveMessageQueueSelector(BROKER_NAME + i); - assertEquals(BROKER_NAME + i, selectorBrokerName.select(ProxyContext.create(), messageQueueView).getBrokerName()); - } + return new MessageQueueView(TOPIC, topicRouteData, null); } }