diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/connector/route/MessageQueueSelector.java b/proxy/src/main/java/org/apache/rocketmq/proxy/connector/route/MessageQueueSelector.java index eff85b1472..6e40325ad7 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/connector/route/MessageQueueSelector.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/connector/route/MessageQueueSelector.java @@ -35,9 +35,9 @@ import org.apache.rocketmq.common.protocol.route.QueueData; public class MessageQueueSelector { private static final int BROKER_ACTING_QUEUE_ID = -1; - // multiple queues for one broker, with queueId : normal + // multiple queues for brokers with queueId : normal private final List queues = new ArrayList<>(); - // one queue for one broker, with queueId : -1 + // one queue for brokers with queueId : -1 private final List brokerActingQueues = new ArrayList<>(); private final Map brokerNameQueueMap = new ConcurrentHashMap<>(); private final AtomicInteger queueIndex; diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/channel/ReceiveMessageChannel.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/channel/ReceiveMessageChannel.java index a561c1e444..8593b66023 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/channel/ReceiveMessageChannel.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/channel/ReceiveMessageChannel.java @@ -19,10 +19,11 @@ package org.apache.rocketmq.proxy.grpc.v2.adapter.channel; import apache.rocketmq.v2.ReceiveMessageRequest; import apache.rocketmq.v2.ReceiveMessageResponse; +import java.util.Iterator; import org.apache.rocketmq.proxy.channel.InvocationChannel; import org.apache.rocketmq.proxy.grpc.v2.adapter.handler.ReceiveMessageResponseHandler; -public class ReceiveMessageChannel extends InvocationChannel { +public class ReceiveMessageChannel extends InvocationChannel> { public ReceiveMessageChannel(ReceiveMessageResponseHandler handler) { super(handler); } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/handler/ReceiveMessageResponseHandler.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/handler/ReceiveMessageResponseHandler.java index edfd795c4e..23ffb5ff6b 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/handler/ReceiveMessageResponseHandler.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/adapter/handler/ReceiveMessageResponseHandler.java @@ -21,11 +21,10 @@ import apache.rocketmq.v2.ReceiveMessageRequest; import apache.rocketmq.v2.ReceiveMessageResponse; import apache.rocketmq.v2.Resource; import com.google.common.base.Stopwatch; -import com.google.protobuf.util.Durations; -import com.google.protobuf.util.Timestamps; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; @@ -44,7 +43,7 @@ import org.apache.rocketmq.proxy.grpc.v2.adapter.ResponseBuilder; import org.apache.rocketmq.remoting.protocol.RemotingCommand; import org.apache.rocketmq.remoting.protocol.RemotingSysResponseCode; -public class ReceiveMessageResponseHandler implements ResponseHandler { +public class ReceiveMessageResponseHandler implements ResponseHandler> { private static final InternalLogger log = InternalLoggerFactory.getLogger(LoggerName.PROXY_LOGGER_NAME); private final String brokerName; private final boolean fifo; @@ -56,20 +55,16 @@ public class ReceiveMessageResponseHandler implements ResponseHandler context) { - CompletableFuture future = context.getResponse(); + InvocationContext> context) { + CompletableFuture> future = context.getResponse(); long currentTimeInMillis = System.currentTimeMillis(); long popCosts = currentTimeInMillis - context.getTimestamp(); try { Stopwatch stopWatch = Stopwatch.createStarted(); - ReceiveMessageResponse.Builder builder = ReceiveMessageResponse.newBuilder(); PopMessageResponseHeader responseHeader = (PopMessageResponseHeader) responseCommand.readCustomHeader(); - builder.setStatus(ResponseBuilder.buildStatus(responseCommand.getCode(), responseCommand.getRemark())); - builder.setInvisibleDuration(Durations.fromMillis(responseHeader.getInvisibleTime())) - .setDeliveryTimestamp(Timestamps.fromMillis(responseHeader.getPopTime())); + List responseList = new ArrayList<>(); - ReceiveMessageResponse response; if (responseCommand.getCode() == RemotingSysResponseCode.SUCCESS) { ByteBuffer byteBuffer = ByteBuffer.wrap(responseCommand.getBody()); List msgFoundList = MessageDecoder.decodes(byteBuffer); @@ -145,14 +140,19 @@ public class ReceiveMessageResponseHandler implements ResponseHandler String.valueOf(responseHeader.getPopTime())); } + responseList.add(ReceiveMessageResponse.newBuilder() + .setStatus(ResponseBuilder.buildStatus(responseCommand.getCode(), responseCommand.getRemark())) + .build()); + for (MessageExt messageExt : msgFoundList) { - builder.addMessages(GrpcConverter.buildMessage(messageExt)); + ReceiveMessageResponse response = ReceiveMessageResponse.newBuilder() + .setMessage(GrpcConverter.buildMessage(messageExt)).build(); + responseList.add(response); } } - response = builder.build(); long elapsed = stopWatch.stop().elapsed(TimeUnit.MILLISECONDS); log.debug("Translating remoting response to gRPC response costs {}ms. Duration request received: {}", elapsed, popCosts); - future.complete(response); + future.complete(responseList.iterator()); } catch (Exception e) { log.error("Unexpected exception raised when handle pop remoting command", e); future.completeExceptionally(e); diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcService.java index faf09fef56..920c5b2938 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcService.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcService.java @@ -41,6 +41,7 @@ import apache.rocketmq.v2.QueryRouteResponse; import apache.rocketmq.v2.ReceiveMessageRequest; import apache.rocketmq.v2.ReceiveMessageResponse; import apache.rocketmq.v2.Resource; +import apache.rocketmq.v2.RetryPolicy; import apache.rocketmq.v2.SendMessageRequest; import apache.rocketmq.v2.SendMessageResponse; import apache.rocketmq.v2.Settings; @@ -51,6 +52,7 @@ import io.grpc.Context; import io.grpc.stub.StreamObserver; import io.netty.channel.Channel; import java.util.ArrayList; +import java.util.Iterator; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; @@ -68,6 +70,7 @@ import org.apache.rocketmq.common.constant.LoggerName; import org.apache.rocketmq.common.consumer.ReceiptHandle; import org.apache.rocketmq.common.message.MessageBatch; import org.apache.rocketmq.common.message.MessageClientIDSetter; +import org.apache.rocketmq.common.message.MessageQueue; import org.apache.rocketmq.common.protocol.RequestCode; import org.apache.rocketmq.common.protocol.ResponseCode; import org.apache.rocketmq.common.protocol.body.ConsumeMessageDirectlyResult; @@ -97,12 +100,14 @@ import org.apache.rocketmq.proxy.grpc.v2.adapter.GrpcConverter; import org.apache.rocketmq.proxy.grpc.v2.adapter.ProxyException; import org.apache.rocketmq.proxy.grpc.v2.adapter.ProxyMode; import org.apache.rocketmq.proxy.grpc.v2.adapter.ResponseBuilder; +import org.apache.rocketmq.proxy.grpc.v2.adapter.ResponseWriter; import org.apache.rocketmq.proxy.grpc.v2.adapter.channel.GrpcClientChannel; import org.apache.rocketmq.proxy.grpc.v2.adapter.channel.ReceiveMessageChannel; import org.apache.rocketmq.proxy.grpc.v2.adapter.channel.SendMessageChannel; import org.apache.rocketmq.proxy.grpc.v2.adapter.handler.ReceiveMessageResponseHandler; import org.apache.rocketmq.proxy.grpc.v2.adapter.handler.SendMessageResponseHandler; import org.apache.rocketmq.proxy.grpc.v2.service.cluster.RouteService; +import org.apache.rocketmq.proxy.grpc.v2.service.local.LocalWriteQueueSelector; import org.apache.rocketmq.remoting.RemotingServer; import org.apache.rocketmq.remoting.netty.NettyRemotingAbstract; import org.apache.rocketmq.remoting.protocol.LanguageCode; @@ -120,6 +125,7 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo private final RouteService routeService; private final ClientSettingsService clientSettingsService; private final DelayPolicy delayPolicy; + private final LocalWriteQueueSelector localWriteQueueSelector; public LocalGrpcService(BrokerController brokerController) { this(brokerController, new TelemetryCommandManager()); @@ -140,6 +146,8 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo this.routeService = new RouteService(ProxyMode.LOCAL, connectorManager, grpcClientManager); this.clientSettingsService = new ClientSettingsService(this.channelManager, this.grpcClientManager, this.telemetryCommandManager); this.delayPolicy = DelayPolicy.build(brokerController.getMessageStoreConfig().getMessageDelayLevel()); + this.localWriteQueueSelector = new LocalWriteQueueSelector(brokerController.getBrokerConfig().getBrokerName(), + brokerController.getTopicConfigManager(), connectorManager.getTopicRouteCache()); this.brokerController.getConsumerManager().appendConsumerIdsChangeListener(new ConsumerIdsChangeListenerImpl()); this.brokerController.getProducerManager().appendProducerChangeListener(new ProducerChangeListenerImpl()); @@ -208,10 +216,11 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo @Override public CompletableFuture sendMessage(Context ctx, SendMessageRequest request) { - String topicName = GrpcConverter.wrapResourceWithNamespace(request.getMessageQueue().getTopic()); - SendMessageRequestHeader requestHeader = GrpcConverter.buildSendMessageRequestHeader(request, topicName); + MessageQueue messageQueue = localWriteQueueSelector.selectQueue(ctx, request).getMessageQueue(); + String topicName = messageQueue.getTopic(); + SendMessageRequestHeader requestHeader = GrpcConverter.buildSendMessageRequestHeader(request, topicName, messageQueue.getQueueId()); RemotingCommand command = RemotingCommand.createRequestCommand(RequestCode.SEND_MESSAGE, requestHeader); - List messageList = GrpcConverter.buildMessage(request.getMessagesList(), topicName); + List messageList = GrpcConverter.buildMessage(request.getMessagesList(), request.getMessages(0).getTopic()); String messageId; if (messageList.size() == 1) { org.apache.rocketmq.common.message.Message message = messageList.get(0); @@ -254,7 +263,7 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo } @Override - public CompletableFuture receiveMessage(Context ctx, ReceiveMessageRequest request) { + public void receiveMessage(Context ctx, ReceiveMessageRequest request, StreamObserver responseObserver) { long pollTime = ctx.getDeadline().timeRemaining(TimeUnit.MILLISECONDS); // TODO: get fifo config from subscriptionGroupManager boolean fifo = false; @@ -265,8 +274,8 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo ReceiveMessageResponseHandler handler = new ReceiveMessageResponseHandler(brokerController.getBrokerConfig().getBrokerName(), fifo); ReceiveMessageChannel channel = channelManager.createChannel(() -> new ReceiveMessageChannel(handler), ReceiveMessageChannel.class); SimpleChannelHandlerContext channelHandlerContext = new SimpleChannelHandlerContext(channel); - CompletableFuture future = new CompletableFuture<>(); - InvocationContext context + CompletableFuture> future = new CompletableFuture<>(); + InvocationContext> context = new InvocationContext<>(request, future); channel.registerInvocationContext(command.getOpaque(), context); try { @@ -280,7 +289,14 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo channel.eraseInvocationContext(command.getOpaque()); future.completeExceptionally(e); } - return future; + future.thenAccept(r -> ResponseWriter.write(responseObserver, r)) + .exceptionally(e -> { + ResponseWriter.write( + responseObserver, + ReceiveMessageResponse.newBuilder().setStatus(ResponseBuilder.buildStatus(e)).build() + ); + return null; + }); } @Override @@ -329,7 +345,8 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo SimpleChannelHandlerContext channelHandlerContext = new SimpleChannelHandlerContext(channel); CompletableFuture future = new CompletableFuture<>(); - int maxReconsumeTimes = grpcClientManager.getClientSettings(ctx).getSubscription().getBackoffPolicy().getMaxAttempts(); + RetryPolicy retryPolicy = grpcClientManager.getClientSettings(ctx).getSubscription().getBackoffPolicy(); + int maxReconsumeTimes = retryPolicy.getMaxAttempts(); if (request.getDeliveryAttempt() >= maxReconsumeTimes) { ConsumerSendMsgBackRequestHeader requestHeader = GrpcConverter.buildConsumerSendMsgBackToDLQRequestHeader(request, maxReconsumeTimes); RemotingCommand command = RemotingCommand.createRequestCommand(RequestCode.CONSUMER_SEND_MSG_BACK, requestHeader); @@ -347,7 +364,7 @@ public class LocalGrpcService extends AbstractStartAndShutdown implements GrpcFo future.completeExceptionally(e); } } else { - ChangeInvisibleTimeRequestHeader requestHeader = GrpcConverter.buildChangeInvisibleTimeRequestHeader(request, delayPolicy); + ChangeInvisibleTimeRequestHeader requestHeader = GrpcConverter.buildChangeInvisibleTimeRequestHeader(request, retryPolicy); RemotingCommand command = RemotingCommand.createRequestCommand(RequestCode.CHANGE_MESSAGE_INVISIBLETIME, requestHeader); command.makeCustomHeaderToNet(); diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/local/LocalWriteQueueSelector.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/local/LocalWriteQueueSelector.java new file mode 100644 index 0000000000..8880b83e9c --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/service/local/LocalWriteQueueSelector.java @@ -0,0 +1,101 @@ +/* + * 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. + */ + +package org.apache.rocketmq.proxy.grpc.v2.service.local; + +import apache.rocketmq.v2.Code; +import apache.rocketmq.v2.Message; +import apache.rocketmq.v2.SendMessageRequest; +import com.google.common.hash.Hashing; +import io.grpc.Context; +import java.util.Random; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.commons.lang3.StringUtils; +import org.apache.rocketmq.broker.topic.TopicConfigManager; +import org.apache.rocketmq.common.TopicConfig; +import org.apache.rocketmq.common.constant.LoggerName; +import org.apache.rocketmq.common.message.MessageQueue; +import org.apache.rocketmq.logging.InternalLogger; +import org.apache.rocketmq.logging.InternalLoggerFactory; +import org.apache.rocketmq.proxy.connector.route.SelectableMessageQueue; +import org.apache.rocketmq.proxy.connector.route.TopicRouteCache; +import org.apache.rocketmq.proxy.grpc.v2.adapter.GrpcConverter; +import org.apache.rocketmq.proxy.grpc.v2.adapter.ProxyException; +import org.apache.rocketmq.proxy.grpc.v2.service.cluster.DefaultWriteQueueSelector; + +public class LocalWriteQueueSelector extends DefaultWriteQueueSelector { + private static final InternalLogger log = InternalLoggerFactory.getLogger(LoggerName.PROXY_LOGGER_NAME); + private final String brokerName; + private final AtomicInteger queueIndex; + private final TopicConfigManager topicConfigManager; + + public LocalWriteQueueSelector(String brokerName, TopicConfigManager topicConfigManager, TopicRouteCache topicRouteCache) { + super(topicRouteCache); + this.brokerName = brokerName; + this.queueIndex = new AtomicInteger(Math.abs(new Random().nextInt())); + this.topicConfigManager = topicConfigManager; + } + + @Override + public SelectableMessageQueue selectQueue( + Context ctx, + SendMessageRequest request + ) { + try { + if (request.getMessagesCount() <= 0) { + throw new ProxyException(Code.MESSAGE_CORRUPTED, "no message to send"); + } + Message message = request.getMessages(0); + String topic = GrpcConverter.wrapResourceWithNamespace(message.getTopic()); + String shardingKey = null; + if (request.getMessagesCount() == 1) { + shardingKey = message.getSystemProperties().getMessageGroup(); + } + SelectableMessageQueue targetMessageQueue; + if (StringUtils.isNotEmpty(shardingKey)) { + // With shardingKey + targetMessageQueue = selectOrderQueue(topic, shardingKey); + } else { + targetMessageQueue = selectNormalQueue(topic); + } + return targetMessageQueue; + } catch (Exception e) { + log.error("error when select queue in DefaultMessageQueueSelector. request: {}", request, e); + return null; + } + } + + protected SelectableMessageQueue selectNormalQueue(String topic) throws Exception { + TopicConfig topicConfig = topicConfigManager.selectTopicConfig(topic); + if (topicConfig == null) { + throw new ProxyException(Code.TOPIC_NOT_FOUND, topic); + } + int selectQueueId = queueIndex.getAndIncrement() % topicConfig.getWriteQueueNums(); + MessageQueue mq = new MessageQueue(topic, brokerName, selectQueueId); + return new SelectableMessageQueue(mq, null); + } + + protected SelectableMessageQueue selectOrderQueue(String topic, String shardingKey) throws Exception { + TopicConfig topicConfig = topicConfigManager.selectTopicConfig(topic); + if (topicConfig == null) { + throw new ProxyException(Code.TOPIC_NOT_FOUND, topic); + } + int selectQueueId = Hashing.consistentHash(shardingKey.hashCode(), topicConfig.getWriteQueueNums()); + MessageQueue mq = new MessageQueue(topic, brokerName, selectQueueId); + return new SelectableMessageQueue(mq, null); + } +} diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcServiceTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcServiceTest.java index 8c93da865d..7639729f38 100644 --- a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcServiceTest.java +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/LocalGrpcServiceTest.java @@ -48,9 +48,9 @@ import apache.rocketmq.v2.SystemProperties; import apache.rocketmq.v2.TelemetryCommand; import apache.rocketmq.v2.ThreadStackTrace; import apache.rocketmq.v2.VerifyMessageResult; -import com.google.protobuf.util.Durations; import io.grpc.Context; import io.grpc.Metadata; +import io.grpc.stub.ServerCallStreamObserver; import io.grpc.stub.StreamObserver; import io.netty.channel.ChannelHandlerContext; import java.net.InetSocketAddress; @@ -71,8 +71,10 @@ import org.apache.rocketmq.broker.processor.EndTransactionProcessor; import org.apache.rocketmq.broker.processor.PopMessageProcessor; import org.apache.rocketmq.broker.processor.PullMessageProcessor; import org.apache.rocketmq.broker.processor.SendMessageProcessor; +import org.apache.rocketmq.broker.topic.TopicConfigManager; import org.apache.rocketmq.common.BrokerConfig; import org.apache.rocketmq.common.ThreadFactoryImpl; +import org.apache.rocketmq.common.TopicConfig; import org.apache.rocketmq.common.consumer.ReceiptHandle; import org.apache.rocketmq.common.message.MessageDecoder; import org.apache.rocketmq.common.message.MessageExt; @@ -83,6 +85,7 @@ import org.apache.rocketmq.common.protocol.header.ConsumerSendMsgBackRequestHead import org.apache.rocketmq.common.protocol.header.PopMessageResponseHeader; import org.apache.rocketmq.proxy.common.TelemetryCommandManager; import org.apache.rocketmq.proxy.common.TelemetryCommandRecord; +import org.apache.rocketmq.proxy.config.ConfigurationManager; import org.apache.rocketmq.proxy.config.InitConfigAndLoggerTest; import org.apache.rocketmq.proxy.connector.transaction.TransactionId; import org.apache.rocketmq.proxy.grpc.interceptor.InterceptorConstants; @@ -95,6 +98,7 @@ import org.apache.rocketmq.store.config.MessageStoreConfig; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.junit.MockitoJUnitRunner; @@ -117,10 +121,14 @@ public class LocalGrpcServiceTest extends InitConfigAndLoggerTest { private ConsumerManager consumerManagerMock; @Mock private ProducerManager producerManagerMock; + @Mock + private TopicConfigManager topicConfigManagerMock; @Mock private TelemetryCommandManager telemetryCommandManager; + StreamObserver receiveStreamObserver = Mockito.mock(ServerCallStreamObserver.class); + private Metadata metadata; private StreamObserver streamObserver; @@ -128,10 +136,13 @@ public class LocalGrpcServiceTest extends InitConfigAndLoggerTest { @Before public void setUp() throws Throwable { super.before(); + ConfigurationManager.getProxyConfig().setNameSrvAddr("1.1.1.1"); Mockito.when(brokerControllerMock.getSendMessageProcessor()).thenReturn(sendMessageProcessorMock); Mockito.when(brokerControllerMock.getPopMessageProcessor()).thenReturn(popMessageProcessorMock); Mockito.when(brokerControllerMock.getBrokerConfig()).thenReturn(new BrokerConfig()); Mockito.when(brokerControllerMock.getMessageStoreConfig()).thenReturn(new MessageStoreConfig()); + Mockito.when(brokerControllerMock.getTopicConfigManager()).thenReturn(topicConfigManagerMock); + Mockito.when(topicConfigManagerMock.selectTopicConfig(Mockito.anyString())).thenReturn(new TopicConfig("topic", 8, 8)); Mockito.doNothing().when(consumerManagerMock).appendConsumerIdsChangeListener(Mockito.any(ConsumerIdsChangeListener.class)); Mockito.doNothing().when(producerManagerMock).appendProducerChangeListener(Mockito.any(ProducerChangeListener.class)); Mockito.when(brokerControllerMock.getConsumerManager()).thenReturn(consumerManagerMock); @@ -320,29 +331,26 @@ public class LocalGrpcServiceTest extends InitConfigAndLoggerTest { .build()) .build()) .build(); - CompletableFuture grpcFuture = localGrpcService.receiveMessage( - Context.current() - .withValue(InterceptorConstants.METADATA, metadata) - .withDeadlineAfter(20, TimeUnit.SECONDS, Executors.newSingleThreadScheduledExecutor( - new ThreadFactoryImpl("test"))), request); - ReceiveMessageResponse r = grpcFuture.get(); - assertThat(r.getStatus().getCode()).isEqualTo(Code.OK); - assertThat(r.getMessagesCount()).isEqualTo(1); - assertThat(Durations.toMillis(r.getInvisibleDuration())).isEqualTo(invisibleTime); - assertThat(GrpcConverter.wrapResourceWithNamespace(r.getMessages(0).getTopic())).isEqualTo(topic); - assertThat(r.getMessages(0).getBody().toByteArray()).isEqualTo(body); - } - - @Test - public void testReceiveMessageSuccessWriteAndFlush() throws Exception { - Mockito.when(popMessageProcessorMock.processRequest(Mockito.any(ChannelHandlerContext.class), Mockito.any(RemotingCommand.class))) - .thenReturn(null); - ReceiveMessageRequest request = ReceiveMessageRequest.newBuilder().getDefaultInstanceForType(); - CompletableFuture grpcFuture = localGrpcService.receiveMessage( - Context.current() - .withDeadlineAfter(20, TimeUnit.SECONDS, Executors.newSingleThreadScheduledExecutor( - new ThreadFactoryImpl("test"))), request); - assertThat(grpcFuture.isDone()).isFalse(); + ReceiveMessageResponse receiveMessageResponse1 = ReceiveMessageResponse.newBuilder() + .setStatus(ResponseBuilder.buildStatus(ResponseCode.SUCCESS, null)) + .build(); + Message message = GrpcConverter.buildMessage(messageExt); + ReceiveMessageResponse receiveMessageResponse2 = ReceiveMessageResponse.newBuilder() + .setMessage(message.toBuilder() + .setSystemProperties( + message.getSystemProperties() + .toBuilder() + .setReceiptHandle("0 0 1000 0 0 zhouxiang_MBP16 0 0 0") + .build()) + .build()) + .build(); + Mockito.doNothing().when(receiveStreamObserver).onNext(Mockito.any()); + localGrpcService.receiveMessage(Context.current().withDeadlineAfter(20, TimeUnit.SECONDS, + Executors.newSingleThreadScheduledExecutor(new ThreadFactoryImpl("test"))), request, receiveStreamObserver); + ArgumentCaptor argument = ArgumentCaptor.forClass(ReceiveMessageResponse.class); + Mockito.verify(receiveStreamObserver, Mockito.times(2)).onNext(argument.capture()); + assertThat(argument.getAllValues().get(0)).isEqualTo(receiveMessageResponse1); + assertThat(argument.getAllValues().get(1)).isEqualTo(receiveMessageResponse2); } @Test diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/local/LocalWriteQueueSelectorTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/local/LocalWriteQueueSelectorTest.java new file mode 100644 index 0000000000..bacf434fd5 --- /dev/null +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/service/local/LocalWriteQueueSelectorTest.java @@ -0,0 +1,78 @@ +/* + * 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. + */ + +package org.apache.rocketmq.proxy.grpc.v2.service.local; + +import apache.rocketmq.v2.Message; +import apache.rocketmq.v2.Resource; +import apache.rocketmq.v2.SendMessageRequest; +import apache.rocketmq.v2.SystemProperties; +import io.grpc.Context; +import org.apache.rocketmq.broker.topic.TopicConfigManager; +import org.apache.rocketmq.common.TopicConfig; +import org.apache.rocketmq.proxy.connector.route.SelectableMessageQueue; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +import static org.assertj.core.api.Assertions.assertThat; + +public class LocalWriteQueueSelectorTest { + private LocalWriteQueueSelector localWriteQueueSelector; + private String topic = "test-topic"; + private String brokerName = "broker"; + private int writeQueueId = 8; + private String messageGroup = "message-group"; + + @Before + public void setup() { + TopicConfigManager topicConfigManager = Mockito.mock(TopicConfigManager.class); + TopicConfig topicConfig = new TopicConfig(topic, writeQueueId, writeQueueId); + Mockito.when(topicConfigManager.selectTopicConfig(topic)).thenReturn(topicConfig); + localWriteQueueSelector = new LocalWriteQueueSelector(brokerName, topicConfigManager, null); + } + + @Test + public void testSelectQueueWithNormalMessage() { + SendMessageRequest sendMessageRequest = SendMessageRequest.newBuilder() + .addMessages(Message.newBuilder().setTopic(Resource.newBuilder().setName(topic).build())) + .build(); + SelectableMessageQueue selectableMessageQueue = localWriteQueueSelector.selectQueue(Context.current(), sendMessageRequest); + assertThat(selectableMessageQueue.getBrokerName()).isEqualTo(brokerName); + assertThat(selectableMessageQueue.getTopic()).isEqualTo(topic); + int selectQueueId = selectableMessageQueue.getQueueId(); + selectableMessageQueue = localWriteQueueSelector.selectQueue(Context.current(), sendMessageRequest); + assertThat(selectableMessageQueue.getQueueId()).isEqualTo((selectQueueId + 1) % writeQueueId); + } + + @Test + public void testSelectQueueWithFifoMessage() { + SendMessageRequest sendMessageRequest = SendMessageRequest.newBuilder() + .addMessages(Message.newBuilder() + .setSystemProperties(SystemProperties.newBuilder() + .setMessageGroup(messageGroup) + .build()) + .setTopic(Resource.newBuilder().setName(topic).build())) + .build(); + SelectableMessageQueue selectableMessageQueue = localWriteQueueSelector.selectQueue(Context.current(), sendMessageRequest); + assertThat(selectableMessageQueue.getBrokerName()).isEqualTo(brokerName); + assertThat(selectableMessageQueue.getTopic()).isEqualTo(topic); + int selectQueueId = selectableMessageQueue.getQueueId(); + selectableMessageQueue = localWriteQueueSelector.selectQueue(Context.current(), sendMessageRequest); + assertThat(selectableMessageQueue.getQueueId()).isEqualTo(selectQueueId); + } +} \ No newline at end of file