mirror of
https://github.com/apache/rocketmq.git
synced 2026-09-24 16:04:00 +08:00
[ISSUE #3949] v2 support
* Add LocalWriteQueueSelector * Adjust for protocol * Add unit test
This commit is contained in:
zhouxiang
+2
-2
@@ -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<SelectableMessageQueue> queues = new ArrayList<>();
|
||||
// one queue for one broker, with queueId : -1
|
||||
// one queue for brokers with queueId : -1
|
||||
private final List<SelectableMessageQueue> brokerActingQueues = new ArrayList<>();
|
||||
private final Map<String, SelectableMessageQueue> brokerNameQueueMap = new ConcurrentHashMap<>();
|
||||
private final AtomicInteger queueIndex;
|
||||
|
||||
+2
-1
@@ -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<ReceiveMessageRequest, ReceiveMessageResponse> {
|
||||
public class ReceiveMessageChannel extends InvocationChannel<ReceiveMessageRequest, Iterator<ReceiveMessageResponse>> {
|
||||
public ReceiveMessageChannel(ReceiveMessageResponseHandler handler) {
|
||||
super(handler);
|
||||
}
|
||||
|
||||
+13
-13
@@ -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<ReceiveMessageRequest, ReceiveMessageResponse> {
|
||||
public class ReceiveMessageResponseHandler implements ResponseHandler<ReceiveMessageRequest, Iterator<ReceiveMessageResponse>> {
|
||||
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<ReceiveMes
|
||||
|
||||
@Override
|
||||
public void handle(RemotingCommand responseCommand,
|
||||
InvocationContext<ReceiveMessageRequest, ReceiveMessageResponse> context) {
|
||||
CompletableFuture<ReceiveMessageResponse> future = context.getResponse();
|
||||
InvocationContext<ReceiveMessageRequest, Iterator<ReceiveMessageResponse>> context) {
|
||||
CompletableFuture<Iterator<ReceiveMessageResponse>> 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<ReceiveMessageResponse> responseList = new ArrayList<>();
|
||||
|
||||
ReceiveMessageResponse response;
|
||||
if (responseCommand.getCode() == RemotingSysResponseCode.SUCCESS) {
|
||||
ByteBuffer byteBuffer = ByteBuffer.wrap(responseCommand.getBody());
|
||||
List<MessageExt> msgFoundList = MessageDecoder.decodes(byteBuffer);
|
||||
@@ -145,14 +140,19 @@ public class ReceiveMessageResponseHandler implements ResponseHandler<ReceiveMes
|
||||
k -> 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);
|
||||
|
||||
+26
-9
@@ -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<SendMessageResponse> 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<org.apache.rocketmq.common.message.Message> messageList = GrpcConverter.buildMessage(request.getMessagesList(), topicName);
|
||||
List<org.apache.rocketmq.common.message.Message> 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<ReceiveMessageResponse> receiveMessage(Context ctx, ReceiveMessageRequest request) {
|
||||
public void receiveMessage(Context ctx, ReceiveMessageRequest request, StreamObserver<ReceiveMessageResponse> 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<ReceiveMessageResponse> future = new CompletableFuture<>();
|
||||
InvocationContext<ReceiveMessageRequest, ReceiveMessageResponse> context
|
||||
CompletableFuture<Iterator<ReceiveMessageResponse>> future = new CompletableFuture<>();
|
||||
InvocationContext<ReceiveMessageRequest, Iterator<ReceiveMessageResponse>> 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<NackMessageResponse> 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();
|
||||
|
||||
|
||||
+101
@@ -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);
|
||||
}
|
||||
}
|
||||
+32
-24
@@ -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<ReceiveMessageResponse> receiveStreamObserver = Mockito.mock(ServerCallStreamObserver.class);
|
||||
|
||||
private Metadata metadata;
|
||||
|
||||
private StreamObserver<TelemetryCommand> 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<ReceiveMessageResponse> 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<ReceiveMessageResponse> 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<ReceiveMessageResponse> 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
|
||||
|
||||
+78
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user