From 0e5dee8b4f6e9fea039e01470a48bb63d0b1901a Mon Sep 17 00:00:00 2001 From: "kaiyi.lk" Date: Tue, 15 Mar 2022 14:45:56 +0800 Subject: [PATCH] [ISSUE #3949] queryRoute, queryAssignment, sendMessage for cluster mode --- .../rocketmq/proxy/client/ClientManager.java | 10 +- .../proxy/client/TopicRouteCache.java | 20 ++ .../client/route/MessageQueueWrapper.java | 8 + .../client/route/SelectableMessageQueue.java | 8 + .../proxy/client/route/TopicRouteWrapper.java | 6 +- .../common/AbstractStartAndShutdown.java | 3 +- .../rocketmq/proxy/grpc/common/Converter.java | 11 + .../proxy/grpc/common/ProxyException.java | 36 +++ .../proxy/grpc/common/ProxyResponseCode.java | 29 ++ .../proxy/grpc/common/ResponseBuilder.java | 15 + .../grpc/interceptor/HeaderInterceptor.java | 4 +- .../grpc/service/ClusterGrpcService.java | 37 ++- .../grpc/service/cluster/BaseService.java | 28 ++ .../grpc/service/cluster/ConsumerService.java | 34 ++ .../grpc/service/cluster/ProducerService.java | 153 +++++++++ .../grpc/service/cluster/RouteService.java | 225 +++++++++++++ .../grpc/service/cluster/BaseServiceTest.java | 62 ++++ .../service/cluster/ProducerServiceTest.java | 304 ++++++++++++++++++ .../service/cluster/RouteServiceTest.java | 197 ++++++++++++ 19 files changed, 1164 insertions(+), 26 deletions(-) create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyException.java create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyResponseCode.java create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseService.java create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ConsumerService.java create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerService.java create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteService.java create mode 100644 proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseServiceTest.java create mode 100644 proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerServiceTest.java create mode 100644 proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteServiceTest.java diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/client/ClientManager.java b/proxy/src/main/java/org/apache/rocketmq/proxy/client/ClientManager.java index 8850da5f5c..5c571620cd 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/client/ClientManager.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/client/ClientManager.java @@ -38,11 +38,11 @@ public class ClientManager extends AbstractStartAndShutdown { this.topicRouteCache = new TopicRouteCache(this.defaultClient); - this.appendStartAndShutdown(this.clientFactory) - .appendStartAndShutdown(this.defaultClient) - .appendStartAndShutdown(this.producerClient) - .appendStartAndShutdown(this.readConsumerClient) - .appendStartAndShutdown(this.writeConsumerClient); + this.appendStartAndShutdown(this.clientFactory); + this.appendStartAndShutdown(this.defaultClient); + this.appendStartAndShutdown(this.producerClient); + this.appendStartAndShutdown(this.readConsumerClient); + this.appendStartAndShutdown(this.writeConsumerClient); } public DefaultClient getDefaultClient() { diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/client/TopicRouteCache.java b/proxy/src/main/java/org/apache/rocketmq/proxy/client/TopicRouteCache.java index e013cc4e4e..c7f911f685 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/client/TopicRouteCache.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/client/TopicRouteCache.java @@ -18,12 +18,15 @@ package org.apache.rocketmq.proxy.client; import com.google.common.cache.CacheBuilder; import com.google.common.cache.LoadingCache; +import com.google.common.hash.Hashing; +import java.util.List; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import org.apache.rocketmq.client.exception.MQClientException; import org.apache.rocketmq.common.protocol.ResponseCode; import org.apache.rocketmq.common.protocol.route.TopicRouteData; import org.apache.rocketmq.common.thread.ThreadPoolMonitor; +import org.apache.rocketmq.proxy.client.route.AddressableMessageQueue; import org.apache.rocketmq.proxy.client.route.MessageQueueWrapper; import org.apache.rocketmq.proxy.common.RetainCacheLoader; import org.apache.rocketmq.proxy.common.RocketMQHelper; @@ -64,6 +67,23 @@ public class TopicRouteCache { return getCacheMessageQueueWrapper(this.topicCache, topicName); } + public AddressableMessageQueue selectOneWriteQueue(String topic, AddressableMessageQueue last) throws Exception { + if (last == null) { + return getMessageQueue(topic).getWrite().selectOne(false); + } + return getMessageQueue(topic).getWrite().selectNextOne(last); + } + + public AddressableMessageQueue selectOneWriteQueue(String topic, String brokerName, int queueId) throws Exception { + return getMessageQueue(topic).getWrite().selectOne(brokerName, queueId); + } + + public AddressableMessageQueue selectOneWriteQueueByKey(String topic, String shardingKey, AddressableMessageQueue last) throws Exception { + List writeQueues = getMessageQueue(topic).getWrite().getQueues(); + int bucket = Hashing.consistentHash(shardingKey.hashCode(), writeQueues.size()); + return writeQueues.get(bucket); + } + protected static MessageQueueWrapper getCacheMessageQueueWrapper(LoadingCache topicCache, String key) throws Exception { MessageQueueWrapper res = topicCache.get(key); if (res.isEmptyCachedQueue()) { diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/MessageQueueWrapper.java b/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/MessageQueueWrapper.java index 5198d7bbfe..d93eaab979 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/MessageQueueWrapper.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/MessageQueueWrapper.java @@ -44,6 +44,14 @@ public class MessageQueueWrapper { return this == EMPTY_CACHED_QUEUE; } + public SelectableMessageQueue getRead() { + return read; + } + + public SelectableMessageQueue getWrite() { + return write; + } + @Override public String toString() { return "MessageQueueWrapper{" + diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/SelectableMessageQueue.java b/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/SelectableMessageQueue.java index 390b19a577..728071e4f3 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/SelectableMessageQueue.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/SelectableMessageQueue.java @@ -210,6 +210,14 @@ public class SelectableMessageQueue { // return newOne; // } + public List getQueues() { + return queues; + } + + public List getBrokers() { + return brokers; + } + @Override public boolean equals(Object o) { if (this == o) { diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/TopicRouteWrapper.java b/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/TopicRouteWrapper.java index 951ca5c5b7..9a3c69cbc5 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/TopicRouteWrapper.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/client/route/TopicRouteWrapper.java @@ -35,8 +35,10 @@ public class TopicRouteWrapper { this.topicRouteData = topicRouteData; this.topicName = topicName; - for (BrokerData brokerData : this.topicRouteData.getBrokerDatas()) { - brokerNameRouteData.put(brokerData.getBrokerName(), brokerData); + if (this.topicRouteData.getBrokerDatas() != null) { + for (BrokerData brokerData : this.topicRouteData.getBrokerDatas()) { + brokerNameRouteData.put(brokerData.getBrokerName(), brokerData); + } } } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/common/AbstractStartAndShutdown.java b/proxy/src/main/java/org/apache/rocketmq/proxy/common/AbstractStartAndShutdown.java index 250d3b3d61..35e361fae1 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/common/AbstractStartAndShutdown.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/common/AbstractStartAndShutdown.java @@ -23,9 +23,8 @@ public abstract class AbstractStartAndShutdown implements StartAndShutdown { protected List startAndShutdownList = new CopyOnWriteArrayList<>(); - public AbstractStartAndShutdown appendStartAndShutdown(StartAndShutdown startAndShutdown) { + protected void appendStartAndShutdown(StartAndShutdown startAndShutdown) { this.startAndShutdownList.add(startAndShutdown); - return this; } @Override diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/Converter.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/Converter.java index c04181da96..5edadf45a4 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/Converter.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/Converter.java @@ -210,6 +210,17 @@ public class Converter { return messageWithHeader.getProperties(); } + public static org.apache.rocketmq.common.message.Message buildMessage(Message protoMessage) { + String topic = getResourceNameWithNamespace(protoMessage.getTopic()); + + org.apache.rocketmq.common.message.Message message = + new org.apache.rocketmq.common.message.Message(topic, protoMessage.getBody().toByteArray()); + Map messageProperty = buildMessageProperty(protoMessage); + + MessageAccessor.setProperties(message, messageProperty); + return message; + } + public static String buildExpressionType(FilterType filterType) { switch (filterType) { case SQL: diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyException.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyException.java new file mode 100644 index 0000000000..093b9424f4 --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyException.java @@ -0,0 +1,36 @@ +/* + * 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.common; + +public class ProxyException extends RuntimeException { + + private final ProxyResponseCode code; + + public ProxyException(ProxyResponseCode proxyResponseCode, String errorMessage) { + super(errorMessage); + this.code = proxyResponseCode; + } + + public ProxyException(ProxyResponseCode proxyResponseCode, String message, Throwable cause) { + super(message, cause); + this.code = proxyResponseCode; + } + + public ProxyResponseCode getCode() { + return code; + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyResponseCode.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyResponseCode.java new file mode 100644 index 0000000000..2134cb6eb3 --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ProxyResponseCode.java @@ -0,0 +1,29 @@ +/* + * 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.common; + +public enum ProxyResponseCode { + SYS_ERR, + PARAMETER_ERR, + AUTH_PERMISSION_CHECK_ERROR, + + NO_TOPIC_ROUTE, + SUBSCRIPTION_NOT_CONSISTENT, + BROKER_NOT_EXIST, + QUERY_NOT_FOUND, + SEND_MSG_FAILED; +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ResponseBuilder.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ResponseBuilder.java index 25eddefd7b..11d347550c 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ResponseBuilder.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/common/ResponseBuilder.java @@ -23,6 +23,8 @@ import apache.rocketmq.v1.SendMessageResponse; import com.google.rpc.Code; import com.google.rpc.Status; import org.apache.commons.lang3.StringUtils; +import org.apache.rocketmq.client.producer.SendResult; +import org.apache.rocketmq.client.producer.SendStatus; import org.apache.rocketmq.common.protocol.ResponseCode; import org.apache.rocketmq.common.protocol.header.SendMessageResponseHeader; import org.apache.rocketmq.remoting.protocol.RemotingCommand; @@ -71,6 +73,19 @@ public class ResponseBuilder { .build(); } + public static SendMessageResponse buildSendMessageResponse(SendResult sendResult) { + if (sendResult.getSendStatus() != SendStatus.SEND_OK) { + return SendMessageResponse.newBuilder() + .setCommon(buildCommon(Code.INTERNAL, "send message failed, sendStatus=" + sendResult.getSendStatus())) + .build(); + } + return SendMessageResponse.newBuilder() + .setCommon(buildCommon(Code.OK, Code.OK.name())) + .setMessageId(StringUtils.defaultString(sendResult.getMsgId())) + .setTransactionId(StringUtils.defaultString(sendResult.getTransactionId())) + .build(); + } + public static Code buildCode(int responseCode) { Code code; switch (responseCode) { diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/interceptor/HeaderInterceptor.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/interceptor/HeaderInterceptor.java index d09ccdadb5..64688dd728 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/interceptor/HeaderInterceptor.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/interceptor/HeaderInterceptor.java @@ -29,8 +29,8 @@ import org.apache.rocketmq.proxy.grpc.common.InterceptorConstants; public class HeaderInterceptor implements ServerInterceptor { @Override - public ServerCall.Listener interceptCall(ServerCall call, Metadata headers, - ServerCallHandler next) { + public ServerCall.Listener interceptCall(ServerCall call, Metadata headers, + ServerCallHandler next) { SocketAddress remoteSocketAddress = call.getAttributes() .get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR); String remoteAddress = parseSocketAddress(remoteSocketAddress); diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/ClusterGrpcService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/ClusterGrpcService.java index ce739b6792..798a918523 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/ClusterGrpcService.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/ClusterGrpcService.java @@ -54,19 +54,32 @@ import apache.rocketmq.v1.SendMessageResponse; import io.grpc.Context; import java.util.concurrent.CompletableFuture; import org.apache.rocketmq.common.constant.LoggerName; +import org.apache.rocketmq.proxy.client.ClientManager; +import org.apache.rocketmq.proxy.common.AbstractStartAndShutdown; +import org.apache.rocketmq.proxy.grpc.service.cluster.ProducerService; +import org.apache.rocketmq.proxy.grpc.service.cluster.RouteService; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -public class ClusterGrpcService implements GrpcForwardService { +public class ClusterGrpcService extends AbstractStartAndShutdown implements GrpcForwardService { private static final Logger LOGGER = LoggerFactory.getLogger(LoggerName.GRPC_LOGGER_NAME); - public ClusterGrpcService() { + private final ClientManager clientManager; + private final ProducerService producerService; + private final RouteService routeService; + public ClusterGrpcService() { + this.clientManager = new ClientManager(checkData -> { + }); + this.producerService = new ProducerService(clientManager); + this.routeService = new RouteService(clientManager); + + this.appendStartAndShutdown(this.clientManager); } @Override public CompletableFuture queryRoute(Context ctx, QueryRouteRequest request) { - return null; + return this.routeService.queryRoute(ctx, request); } @Override @@ -81,15 +94,16 @@ public class ClusterGrpcService implements GrpcForwardService { @Override public CompletableFuture sendMessage(Context ctx, SendMessageRequest request) { - return null; + return this.producerService.sendMessage(ctx, request); } @Override public CompletableFuture queryAssignment(Context ctx, QueryAssignmentRequest request) { - return null; + return this.routeService.queryAssignment(ctx, request); } - @Override public CompletableFuture receiveMessage(Context ctx, ReceiveMessageRequest request) { + @Override + public CompletableFuture receiveMessage(Context ctx, ReceiveMessageRequest request) { return null; } @@ -128,7 +142,8 @@ public class ClusterGrpcService implements GrpcForwardService { return null; } - @Override public CompletableFuture reportMessageConsumptionResult(Context ctx, + @Override + public CompletableFuture reportMessageConsumptionResult(Context ctx, ReportMessageConsumptionResultRequest request) { return null; } @@ -142,12 +157,4 @@ public class ClusterGrpcService implements GrpcForwardService { ChangeInvisibleDurationRequest request) { return null; } - - @Override - public void start() throws Exception { - } - - @Override - public void shutdown() throws Exception { - } } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseService.java new file mode 100644 index 0000000000..acb9691322 --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseService.java @@ -0,0 +1,28 @@ +/* + * 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.service.cluster; + +import org.apache.rocketmq.proxy.client.ClientManager; + +public class BaseService { + + protected final ClientManager clientManager; + + public BaseService(ClientManager clientManager) { + this.clientManager = clientManager; + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ConsumerService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ConsumerService.java new file mode 100644 index 0000000000..8894c16fb7 --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ConsumerService.java @@ -0,0 +1,34 @@ +/* + * 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.service.cluster; + +import apache.rocketmq.v1.ReceiveMessageRequest; +import apache.rocketmq.v1.ReceiveMessageResponse; +import io.grpc.Context; +import java.util.concurrent.CompletableFuture; +import org.apache.rocketmq.proxy.client.ClientManager; + +public class ConsumerService extends BaseService { + + public ConsumerService(ClientManager clientManager) { + super(clientManager); + } + + public CompletableFuture receiveMessage(Context ctx, ReceiveMessageRequest request) { + + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerService.java new file mode 100644 index 0000000000..94c316778d --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerService.java @@ -0,0 +1,153 @@ +/* + * 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.service.cluster; + +import apache.rocketmq.v1.SendMessageRequest; +import apache.rocketmq.v1.SendMessageResponse; +import io.grpc.Context; +import java.util.concurrent.CompletableFuture; +import org.apache.commons.lang3.StringUtils; +import org.apache.rocketmq.client.producer.SendResult; +import org.apache.rocketmq.common.message.Message; +import org.apache.rocketmq.common.message.MessageConst; +import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader; +import org.apache.rocketmq.proxy.client.ClientManager; +import org.apache.rocketmq.proxy.client.route.AddressableMessageQueue; +import org.apache.rocketmq.proxy.common.utils.ProxyUtils; +import org.apache.rocketmq.proxy.grpc.common.Converter; +import org.apache.rocketmq.proxy.grpc.common.ProxyException; +import org.apache.rocketmq.proxy.grpc.common.ProxyResponseCode; +import org.apache.rocketmq.proxy.grpc.common.ResponseBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ProducerService extends BaseService { + + private static final Logger log = LoggerFactory.getLogger(ProducerService.class); + + private volatile ProducerServiceHook producerServiceHook = null; + private volatile MessageQueueSelector messageQueueSelector = new DefaultMessageQueueSelector(); + + public ProducerService(ClientManager clientManager) { + super(clientManager); + } + + public interface MessageQueueSelector { + AddressableMessageQueue selectQueue(Context ctx, SendMessageRequest request, SendMessageRequestHeader requestHeader, + org.apache.rocketmq.common.message.Message message); + } + + public class DefaultMessageQueueSelector implements MessageQueueSelector { + + @Override + public AddressableMessageQueue selectQueue(Context ctx, SendMessageRequest request, SendMessageRequestHeader requestHeader, + org.apache.rocketmq.common.message.Message message) { + try { + String topic = requestHeader.getTopic(); + String brokerName = ""; + if (request.hasPartition()) { + brokerName = request.getPartition().getBroker().getName(); + } + Integer queueId = requestHeader.getQueueId(); + String shardingKey = message.getProperty(MessageConst.PROPERTY_SHARDING_KEY); + AddressableMessageQueue addressableMessageQueue; + if (!StringUtils.isBlank(brokerName) && queueId != null) { + // Grpc client sendSelect situation + addressableMessageQueue = selectTargetQueue(topic, brokerName, queueId); + } else if (shardingKey != null) { + // With shardingKey + addressableMessageQueue = selectOrderQueue(topic, shardingKey); + } else { + addressableMessageQueue = selectNormalQueue(topic); + } + return addressableMessageQueue; + } catch (Exception e) { + log.error("error when select queue in DefaultMessageQueueSelector. request: {}", request, e); + return null; + } + } + + protected AddressableMessageQueue selectNormalQueue(String topic) throws Exception { + return clientManager.getTopicRouteCache().selectOneWriteQueue(topic, null); + } + + protected AddressableMessageQueue selectTargetQueue(String topic, String brokerName, int queueId) throws Exception { + return clientManager.getTopicRouteCache().selectOneWriteQueue(topic, brokerName, queueId); + } + + protected AddressableMessageQueue selectOrderQueue(String topic, String shardingKey) throws Exception { + return clientManager.getTopicRouteCache().selectOneWriteQueueByKey(topic, shardingKey, null); + } + } + + public interface ProducerServiceHook { + + void beforeSend(Context ctx, AddressableMessageQueue addressableMessageQueue, Message msg, SendMessageRequestHeader requestHeader); + + void afterSend(Context ctx, AddressableMessageQueue addressableMessageQueue, Message msg, SendMessageRequestHeader requestHeader, + SendResult sendResult); + } + + public void setProducerServiceHook(ProducerServiceHook hook) { + this.producerServiceHook = hook; + } + + public void setMessageQueueSelector(MessageQueueSelector messageQueueSelector) { + this.messageQueueSelector = messageQueueSelector; + } + + public CompletableFuture sendMessage(Context ctx, SendMessageRequest request) { + org.apache.rocketmq.common.message.Message message = Converter.buildMessage(request.getMessage()); + CompletableFuture future = new CompletableFuture<>(); + + try { + SendMessageRequestHeader requestHeader = Converter.buildSendMessageRequestHeader(request); + AddressableMessageQueue addressableMessageQueue = messageQueueSelector.selectQueue(ctx, request, requestHeader, message); + + String topic = requestHeader.getTopic(); + if (addressableMessageQueue == null) { + throw new ProxyException(ProxyResponseCode.NO_TOPIC_ROUTE, + "no writeable topic route for topic " + topic); + } + + if (producerServiceHook != null) { + producerServiceHook.beforeSend(ctx, addressableMessageQueue, message, requestHeader); + } + CompletableFuture sendResultCompletableFuture = this.clientManager.getProducerClient().sendMessage( + addressableMessageQueue.getBrokerAddr(), + addressableMessageQueue.getBrokerName(), + message, + requestHeader, + ProxyUtils.DEFAULT_MQ_CLIENT_TIMEOUT + ); + sendResultCompletableFuture + .thenAccept(result -> { + if (producerServiceHook != null) { + producerServiceHook.afterSend(ctx, addressableMessageQueue, message, requestHeader, result); + } + future.complete(ResponseBuilder.buildSendMessageResponse(result)); + }) + .exceptionally(e -> { + future.completeExceptionally(e); + return null; + }); + } catch (Throwable t) { + future.completeExceptionally(t); + } + return future; + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteService.java new file mode 100644 index 0000000000..4bcea12cc0 --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteService.java @@ -0,0 +1,225 @@ +/* + * 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.service.cluster; + +import apache.rocketmq.v1.Assignment; +import apache.rocketmq.v1.Broker; +import apache.rocketmq.v1.Endpoints; +import apache.rocketmq.v1.Partition; +import apache.rocketmq.v1.Permission; +import apache.rocketmq.v1.QueryAssignmentRequest; +import apache.rocketmq.v1.QueryAssignmentResponse; +import apache.rocketmq.v1.QueryRouteRequest; +import apache.rocketmq.v1.QueryRouteResponse; +import apache.rocketmq.v1.Resource; +import com.google.rpc.Code; +import io.grpc.Context; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import org.apache.rocketmq.common.constant.PermName; +import org.apache.rocketmq.common.protocol.route.QueueData; +import org.apache.rocketmq.common.protocol.route.TopicRouteData; +import org.apache.rocketmq.proxy.client.ClientManager; +import org.apache.rocketmq.proxy.client.route.AddressableMessageQueue; +import org.apache.rocketmq.proxy.client.route.MessageQueueWrapper; +import org.apache.rocketmq.proxy.common.RocketMQHelper; +import org.apache.rocketmq.proxy.grpc.common.Converter; +import org.apache.rocketmq.proxy.grpc.common.ResponseBuilder; + +public class RouteService extends BaseService { + + private volatile RouteAssignmentQueueSelector assignmentQueueSelector = new DefaultRouteAssignmentQueueSelector(); + private volatile QueryRouteHook queryRouteHook = null; + private volatile QueryAssignmentHook queryAssignmentHook = null; + + public RouteService(ClientManager clientManager) { + super(clientManager); + } + + public interface QueryRouteHook { + QueryRouteResponse beforeResponse(Context ctx, QueryRouteRequest request, QueryRouteResponse response); + } + + public interface QueryAssignmentHook { + QueryAssignmentResponse beforeResponse(Context ctx, QueryAssignmentRequest request, + QueryAssignmentResponse response); + } + + public interface RouteAssignmentQueueSelector { + List getAssignment(QueryAssignmentRequest request) throws Exception; + } + + public class DefaultRouteAssignmentQueueSelector implements RouteAssignmentQueueSelector { + + @Override + public List getAssignment(QueryAssignmentRequest request) throws Exception { + MessageQueueWrapper messageQueueWrapper = clientManager.getTopicRouteCache() + .getMessageQueue(Converter.getResourceNameWithNamespace(request.getTopic())); + return messageQueueWrapper.getRead().getBrokers(); + } + } + + public void setQueryRouteHook(QueryRouteHook queryRouteHook) { + this.queryRouteHook = queryRouteHook; + } + + public void setAssignmentQueueSelector(RouteAssignmentQueueSelector assignmentQueueSelector) { + this.assignmentQueueSelector = assignmentQueueSelector; + } + + public void setQueryAssignmentHook(QueryAssignmentHook queryAssignmentHook) { + this.queryAssignmentHook = queryAssignmentHook; + } + + public CompletableFuture queryRoute(Context ctx, QueryRouteRequest request) { + CompletableFuture future = new CompletableFuture<>(); + CompletableFuture resFuture = future.thenApply(r -> { + if (this.queryRouteHook != null) { + return this.queryRouteHook.beforeResponse(ctx, request, r); + } + return r; + }); + + try { + Endpoints resEndpoints = request.getEndpoints(); + if (resEndpoints.getDefaultInstanceForType().equals(resEndpoints)) { + future.complete(QueryRouteResponse.newBuilder() + .setCommon(ResponseBuilder.buildCommon(Code.INVALID_ARGUMENT, "endpoint " + + request.getEndpoints() + " is invalidate")) + .build()); + return resFuture; + } + + MessageQueueWrapper messageQueueWrapper = this.clientManager.getTopicRouteCache() + .getMessageQueue(Converter.getResourceNameWithNamespace(request.getTopic())); + TopicRouteData topicRouteData = messageQueueWrapper.getTopicRouteData(); + List queueDataList = topicRouteData.getQueueDatas(); + + List partitionList = new ArrayList<>(); + for (QueueData queueData : queueDataList) { + Broker broker = Broker.newBuilder() + .setName(queueData.getBrokerName()) + .setId(0) + .setEndpoints(resEndpoints) + .build(); + + partitionList.addAll(genPartitionFromQueueData(queueData, request.getTopic(), broker)); + } + QueryRouteResponse response = QueryRouteResponse.newBuilder() + .setCommon(ResponseBuilder.buildCommon(Code.OK, Code.OK.name())) + .addAllPartitions(partitionList) + .build(); + future.complete(response); + } catch (Throwable t) { + if (RocketMQHelper.isTopicNotExistError(t)) { + future.complete(QueryRouteResponse.newBuilder() + .setCommon(ResponseBuilder.buildCommon(Code.NOT_FOUND, t.getMessage())) + .build()); + } else { + future.completeExceptionally(t); + } + } + return resFuture; + } + + protected static List genPartitionFromQueueData(QueueData queueData, Resource topic, Broker broker) { + List partitionList = new ArrayList<>(); + + int r = 0; + int w = 0; + int rw = 0; + if (PermName.isWriteable(queueData.getPerm()) && PermName.isReadable(queueData.getPerm())) { + rw = Math.min(queueData.getWriteQueueNums(), queueData.getReadQueueNums()); + r = queueData.getReadQueueNums() - rw; + w = queueData.getWriteQueueNums() - rw; + } else if (PermName.isWriteable(queueData.getPerm())) { + w = queueData.getWriteQueueNums(); + } else if (PermName.isReadable(queueData.getPerm())) { + r = queueData.getReadQueueNums(); + } + + for (int i = 0; i < (rw + r + w); i++) { + Partition.Builder builder = Partition.newBuilder() + .setBroker(broker) + .setTopic(topic) + .setId(i); + if (i < r) { + builder.setPermission(Permission.READ); + } else if (i < w) { + builder.setPermission(Permission.WRITE); + } else { + builder.setPermission(Permission.READ_WRITE); + } + partitionList.add(builder.build()); + } + return partitionList; + } + + public CompletableFuture queryAssignment(Context ctx, QueryAssignmentRequest request) { + CompletableFuture future = new CompletableFuture<>(); + CompletableFuture resFuture = future.thenApply(r -> { + if (this.queryAssignmentHook != null) { + return this.queryAssignmentHook.beforeResponse(ctx, request, r); + } + return r; + }); + try { + Endpoints resEndpoints = request.getEndpoints(); + if (resEndpoints.getDefaultInstanceForType().equals(resEndpoints)) { + future.complete(QueryAssignmentResponse.newBuilder() + .setCommon(ResponseBuilder.buildCommon(Code.INVALID_ARGUMENT, "endpoint " + + request.getEndpoints() + " is invalidate")) + .build()); + return resFuture; + } + + List assignments = new ArrayList<>(); + List messageQueueList = this.assignmentQueueSelector.getAssignment(request); + + for (AddressableMessageQueue messageQueue : messageQueueList) { + Broker broker = Broker.newBuilder() + .setName(messageQueue.getBrokerName()) + .setId(0) + .setEndpoints(resEndpoints) + .build(); + + Partition defaultPartition = Partition.newBuilder() + .setTopic(request.getTopic()) + .setId(-1) + .setPermission(Permission.READ_WRITE) + .setBroker(broker) + .build(); + + assignments.add(Assignment.newBuilder() + .setPartition(defaultPartition) + .build()); + } + QueryAssignmentResponse response = QueryAssignmentResponse.newBuilder() + .addAllAssignments(assignments) + .setCommon(ResponseBuilder.buildCommon(Code.OK, Code.OK.name())) + .build(); + if (this.queryAssignmentHook != null) { + this.queryAssignmentHook.beforeResponse(ctx, request, response); + } + future.complete(response); + } catch (Throwable t) { + future.completeExceptionally(t); + } + return resFuture; + } +} diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseServiceTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseServiceTest.java new file mode 100644 index 0000000000..fd220378a0 --- /dev/null +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/BaseServiceTest.java @@ -0,0 +1,62 @@ +/* + * 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.service.cluster; + +import org.apache.rocketmq.proxy.client.ClientManager; +import org.apache.rocketmq.proxy.client.DefaultClient; +import org.apache.rocketmq.proxy.client.ProducerClient; +import org.apache.rocketmq.proxy.client.ReadConsumerClient; +import org.apache.rocketmq.proxy.client.TopicRouteCache; +import org.apache.rocketmq.proxy.client.WriteConsumerClient; +import org.junit.Before; +import org.junit.Ignore; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; + +import static org.mockito.Mockito.when; + +@Ignore +@RunWith(MockitoJUnitRunner.Silent.class) +public abstract class BaseServiceTest { + + @Mock + protected ClientManager clientManager; + @Mock + protected DefaultClient defaultClient; + @Mock + protected ProducerClient producerClient; + @Mock + protected ReadConsumerClient readConsumerClient; + @Mock + protected WriteConsumerClient writeConsumerClient; + @Mock + protected TopicRouteCache topicRouteCache; + + @Before + public void before() throws Throwable { + when(clientManager.getDefaultClient()).thenReturn(defaultClient); + when(clientManager.getProducerClient()).thenReturn(producerClient); + when(clientManager.getReadConsumerClient()).thenReturn(readConsumerClient); + when(clientManager.getWriteConsumerClient()).thenReturn(writeConsumerClient); + when(clientManager.getTopicRouteCache()).thenReturn(topicRouteCache); + + beforeEach(); + } + + public abstract void beforeEach() throws Throwable; +} diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerServiceTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerServiceTest.java new file mode 100644 index 0000000000..4165ad8d69 --- /dev/null +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/ProducerServiceTest.java @@ -0,0 +1,304 @@ +/* + * 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.service.cluster; + +import apache.rocketmq.v1.Broker; +import apache.rocketmq.v1.Message; +import apache.rocketmq.v1.Partition; +import apache.rocketmq.v1.Resource; +import apache.rocketmq.v1.SendMessageRequest; +import apache.rocketmq.v1.SendMessageResponse; +import apache.rocketmq.v1.SystemAttribute; +import com.google.protobuf.ByteString; +import com.google.rpc.Code; +import io.grpc.Context; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.rocketmq.client.producer.SendResult; +import org.apache.rocketmq.client.producer.SendStatus; +import org.apache.rocketmq.common.message.MessageConst; +import org.apache.rocketmq.common.message.MessageQueue; +import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader; +import org.apache.rocketmq.proxy.client.route.AddressableMessageQueue; +import org.apache.rocketmq.proxy.grpc.common.ProxyException; +import org.apache.rocketmq.proxy.grpc.common.ProxyResponseCode; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.when; + +public class ProducerServiceTest extends BaseServiceTest { + + @Override + public void beforeEach() throws Throwable { + AddressableMessageQueue queue = new AddressableMessageQueue( + new MessageQueue("topic", "selectOrderQueue", 0), + "selectOrderQueueAddr"); + when(topicRouteCache.selectOneWriteQueueByKey(anyString(), anyString(), isNull())) + .thenReturn(queue); + + queue = new AddressableMessageQueue( + new MessageQueue("topic", "selectTargetQueue", 0), + "selectTargetQueueAddr"); + when(topicRouteCache.selectOneWriteQueue(anyString(), anyString(), anyInt())) + .thenReturn(queue); + + queue = new AddressableMessageQueue( + new MessageQueue("topic", "selectNormalQueue", 0), + "selectNormalQueueAddr"); + when(topicRouteCache.selectOneWriteQueue(anyString(), isNull())) + .thenReturn(queue); + } + + @Test + public void testSendOrderMessageWithShardingKey() { + CompletableFuture sendResultFuture = new CompletableFuture<>(); + when(producerClient.sendMessage(anyString(), anyString(), any(), any(), anyLong())) + .thenReturn(sendResultFuture); + sendResultFuture.complete(new SendResult(SendStatus.SEND_OK, "msgId", new MessageQueue(), + 1L, "txId", "offsetMsgId", "regionId")); + + ProducerService producerService = new ProducerService(this.clientManager); + + AtomicReference selectQueueRef = new AtomicReference<>(); + AtomicReference messageRef = new AtomicReference<>(); + producerService.setProducerServiceHook(new ProducerService.ProducerServiceHook() { + @Override + public void beforeSend(Context ctx, AddressableMessageQueue addressableMessageQueue, + org.apache.rocketmq.common.message.Message msg, SendMessageRequestHeader requestHeader) { + selectQueueRef.set(addressableMessageQueue); + } + + @Override + public void afterSend(Context ctx, AddressableMessageQueue addressableMessageQueue, + org.apache.rocketmq.common.message.Message msg, SendMessageRequestHeader requestHeader, + SendResult sendResult) { + + } + }); + + CompletableFuture future = producerService.sendMessage(Context.current(), SendMessageRequest.newBuilder() + .setMessage(Message.newBuilder() + .setTopic(Resource.newBuilder() + .setResourceNamespace("namespace") + .setName("topic") + .build()) + .putUserAttribute(MessageConst.PROPERTY_SHARDING_KEY, "key") + .setSystemAttribute(SystemAttribute.newBuilder() + .setMessageId("msgId") + .build()) + .setBody(ByteString.copyFrom("hello", StandardCharsets.UTF_8)) + .build()) + .build()); + + try { + SendMessageResponse response = future.get(); + + assertEquals(Code.OK.getNumber(), response.getCommon().getStatus().getCode()); + assertEquals("msgId", response.getMessageId()); + assertEquals("selectOrderQueue", selectQueueRef.get().getBrokerName()); + assertEquals("selectOrderQueueAddr", selectQueueRef.get().getBrokerAddr()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testSendNormalMessage() { + CompletableFuture sendResultFuture = new CompletableFuture<>(); + when(producerClient.sendMessage(anyString(), anyString(), any(), any(), anyLong())) + .thenReturn(sendResultFuture); + sendResultFuture.complete(new SendResult(SendStatus.SEND_OK, "msgId", new MessageQueue(), + 1L, "txId", "offsetMsgId", "regionId")); + + ProducerService producerService = new ProducerService(this.clientManager); + + AtomicReference selectQueueRef = new AtomicReference<>(); + AtomicReference messageRef = new AtomicReference<>(); + producerService.setProducerServiceHook(new ProducerService.ProducerServiceHook() { + @Override + public void beforeSend(Context ctx, AddressableMessageQueue addressableMessageQueue, + org.apache.rocketmq.common.message.Message msg, SendMessageRequestHeader requestHeader) { + selectQueueRef.set(addressableMessageQueue); + } + + @Override + public void afterSend(Context ctx, AddressableMessageQueue addressableMessageQueue, + org.apache.rocketmq.common.message.Message msg, SendMessageRequestHeader requestHeader, + SendResult sendResult) { + + } + }); + + CompletableFuture future = producerService.sendMessage(Context.current(), SendMessageRequest.newBuilder() + .setMessage(Message.newBuilder() + .setTopic(Resource.newBuilder() + .setResourceNamespace("namespace") + .setName("topic") + .build()) + .setSystemAttribute(SystemAttribute.newBuilder() + .setMessageId("msgId") + .build()) + .setBody(ByteString.copyFrom("hello", StandardCharsets.UTF_8)) + .build()) + .build()); + + try { + SendMessageResponse response = future.get(); + + assertEquals(Code.OK.getNumber(), response.getCommon().getStatus().getCode()); + assertEquals("msgId", response.getMessageId()); + assertEquals("selectNormalQueue", selectQueueRef.get().getBrokerName()); + assertEquals("selectNormalQueueAddr", selectQueueRef.get().getBrokerAddr()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testSendOrderMessageSelectQueue() { + CompletableFuture sendResultFuture = new CompletableFuture<>(); + when(producerClient.sendMessage(anyString(), anyString(), any(), any(), anyLong())) + .thenReturn(sendResultFuture); + sendResultFuture.complete(new SendResult(SendStatus.SEND_OK, "msgId", new MessageQueue(), + 1L, "txId", "offsetMsgId", "regionId")); + + ProducerService producerService = new ProducerService(this.clientManager); + + AtomicReference selectQueueRef = new AtomicReference<>(); + AtomicReference messageRef = new AtomicReference<>(); + producerService.setProducerServiceHook(new ProducerService.ProducerServiceHook() { + @Override + public void beforeSend(Context ctx, AddressableMessageQueue addressableMessageQueue, + org.apache.rocketmq.common.message.Message msg, SendMessageRequestHeader requestHeader) { + selectQueueRef.set(addressableMessageQueue); + } + + @Override + public void afterSend(Context ctx, AddressableMessageQueue addressableMessageQueue, + org.apache.rocketmq.common.message.Message msg, SendMessageRequestHeader requestHeader, + SendResult sendResult) { + + } + }); + + CompletableFuture future = producerService.sendMessage(Context.current(), SendMessageRequest.newBuilder() + .setMessage(Message.newBuilder() + .setTopic(Resource.newBuilder() + .setResourceNamespace("namespace") + .setName("topic") + .build()) + .setSystemAttribute(SystemAttribute.newBuilder() + .setMessageId("msgId") + .setPartitionId(1) + .build()) + .setBody(ByteString.copyFrom("hello", StandardCharsets.UTF_8)) + .build()) + .setPartition(Partition.newBuilder() + .setBroker(Broker.newBuilder() + .setName("brokerName") + .build()) + .build()) + .build()); + + try { + SendMessageResponse response = future.get(); + + assertEquals(Code.OK.getNumber(), response.getCommon().getStatus().getCode()); + assertEquals("msgId", response.getMessageId()); + assertEquals("selectTargetQueue", selectQueueRef.get().getBrokerName()); + assertEquals("selectTargetQueueAddr", selectQueueRef.get().getBrokerAddr()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testSendMessageNoQueueSelect() { + ProducerService producerService = new ProducerService(this.clientManager); + + producerService.setMessageQueueSelector((ctx, request, requestHeader, message) -> null); + + CompletableFuture future = producerService.sendMessage(Context.current(), SendMessageRequest.newBuilder() + .setMessage(Message.newBuilder() + .setTopic(Resource.newBuilder() + .setResourceNamespace("namespace") + .setName("topic") + .build()) + .setSystemAttribute(SystemAttribute.newBuilder() + .setMessageId("msgId") + .build()) + .setBody(ByteString.copyFrom("hello", StandardCharsets.UTF_8)) + .build()) + .build()); + + try { + SendMessageResponse response = future.get(); + assertNull(response); + } catch (Exception e) { + assertNotNull(e); + assertTrue(e instanceof ExecutionException); + assertTrue(e.getCause() instanceof ProxyException); + assertEquals(ProxyResponseCode.NO_TOPIC_ROUTE, ((ProxyException)e.getCause()).getCode()); + } + } + + @Test + public void testSendMessageWithError() { + RuntimeException ex = new RuntimeException(); + + CompletableFuture sendResultFuture = new CompletableFuture<>(); + when(producerClient.sendMessage(anyString(), anyString(), any(), any(), anyLong())) + .thenReturn(sendResultFuture); + sendResultFuture.completeExceptionally(ex); + + ProducerService producerService = new ProducerService(this.clientManager); + + CompletableFuture future = producerService.sendMessage(Context.current(), SendMessageRequest.newBuilder() + .setMessage(Message.newBuilder() + .setTopic(Resource.newBuilder() + .setResourceNamespace("namespace") + .setName("topic") + .build()) + .setSystemAttribute(SystemAttribute.newBuilder() + .setMessageId("msgId") + .build()) + .setBody(ByteString.copyFrom("hello", StandardCharsets.UTF_8)) + .build()) + .build()); + + try { + SendMessageResponse response = future.get(); + assertNull(response); + } catch (Exception e) { + assertNotNull(e); + assertTrue(e instanceof ExecutionException); + assertSame(ex, e.getCause()); + } + } +} \ No newline at end of file diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteServiceTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteServiceTest.java new file mode 100644 index 0000000000..39ee582e6a --- /dev/null +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/cluster/RouteServiceTest.java @@ -0,0 +1,197 @@ +package org.apache.rocketmq.proxy.grpc.service.cluster; + +import apache.rocketmq.v1.Address; +import apache.rocketmq.v1.AddressScheme; +import apache.rocketmq.v1.Endpoints; +import apache.rocketmq.v1.QueryAssignmentRequest; +import apache.rocketmq.v1.QueryAssignmentResponse; +import apache.rocketmq.v1.QueryRouteRequest; +import apache.rocketmq.v1.QueryRouteResponse; +import apache.rocketmq.v1.Resource; +import com.google.rpc.Code; +import io.grpc.Context; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import org.apache.rocketmq.client.exception.MQClientException; +import org.apache.rocketmq.common.protocol.ResponseCode; +import org.apache.rocketmq.common.protocol.route.BrokerData; +import org.apache.rocketmq.common.protocol.route.QueueData; +import org.apache.rocketmq.common.protocol.route.TopicRouteData; +import org.apache.rocketmq.proxy.client.route.MessageQueueWrapper; +import org.apache.rocketmq.proxy.grpc.common.ResponseBuilder; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.when; + +public class RouteServiceTest extends BaseServiceTest { + + @Override + public void beforeEach() throws Throwable { + TopicRouteData routeData = new TopicRouteData(); + + List brokerDataList = new ArrayList<>(); + BrokerData brokerData = new BrokerData(); + brokerData.setCluster("cluster"); + brokerData.setBrokerName("brokerName"); + HashMap brokerAddrs = new HashMap() {{ + put(0L, "127.0.0.1:10911"); + }}; + brokerData.setBrokerAddrs(brokerAddrs); + brokerDataList.add(brokerData); + + List queueDataList = new ArrayList<>(); + QueueData queueData = new QueueData(); + queueData.setPerm(6); + queueData.setWriteQueueNums(8); + queueData.setReadQueueNums(8); + queueData.setBrokerName("brokerName"); + queueDataList.add(queueData); + + routeData.setBrokerDatas(brokerDataList); + routeData.setQueueDatas(queueDataList); + + MessageQueueWrapper messageQueueWrapper = new MessageQueueWrapper("topic", routeData); + when(this.topicRouteCache.getMessageQueue("topic")).thenReturn(messageQueueWrapper); + + when(this.topicRouteCache.getMessageQueue("notExistTopic")).thenThrow(new MQClientException(ResponseCode.TOPIC_NOT_EXIST, "")); + } + + @Test + public void testQueryRouteWithInvalidEndpoints() { + RouteService routeService = new RouteService(this.clientManager); + + CompletableFuture future = routeService.queryRoute(Context.current(), QueryRouteRequest.newBuilder() + .build()); + + try { + QueryRouteResponse response = future.get(); + assertEquals(Code.INVALID_ARGUMENT.getNumber(), response.getCommon().getStatus().getCode()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testQueryRoute() { + RouteService routeService = new RouteService(this.clientManager); + + CompletableFuture future = routeService.queryRoute(Context.current(), QueryRouteRequest.newBuilder() + .setEndpoints(Endpoints.newBuilder() + .addAddresses(Address.newBuilder() + .setPort(80) + .setHost("host") + .build()) + .setScheme(AddressScheme.DOMAIN_NAME) + .build()) + .setTopic(Resource.newBuilder() + .setName("topic") + .build()) + .build()); + + try { + QueryRouteResponse response = future.get(); + assertEquals(Code.OK.getNumber(), response.getCommon().getStatus().getCode()); + assertEquals(8, response.getPartitionsCount()); + assertEquals("host", response.getPartitions(0).getBroker() + .getEndpoints().getAddresses(0).getHost()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testQueryRouteWhenTopicNotExist() { + RouteService routeService = new RouteService(this.clientManager); + + CompletableFuture future = routeService.queryRoute(Context.current(), QueryRouteRequest.newBuilder() + .setEndpoints(Endpoints.newBuilder() + .addAddresses(Address.newBuilder() + .setPort(80) + .setHost("host") + .build()) + .setScheme(AddressScheme.DOMAIN_NAME) + .build()) + .setTopic(Resource.newBuilder() + .setName("notExistTopic") + .build()) + .build()); + + try { + QueryRouteResponse response = future.get(); + assertEquals(Code.NOT_FOUND.getNumber(), response.getCommon().getStatus().getCode()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testQueryRouteHook() { + RouteService routeService = new RouteService(this.clientManager); + + routeService.setQueryRouteHook((ctx, request, response) -> QueryRouteResponse.newBuilder() + .setCommon(ResponseBuilder.buildCommon(Code.NOT_FOUND, Code.NOT_FOUND.name())) + .build()); + + CompletableFuture future = routeService.queryRoute(Context.current(), QueryRouteRequest.newBuilder() + .build()); + + try { + QueryRouteResponse response = future.get(); + assertEquals(Code.NOT_FOUND.getNumber(), response.getCommon().getStatus().getCode()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testQueryAssignmentInvalidEndpoints() { + RouteService routeService = new RouteService(this.clientManager); + + CompletableFuture future = routeService.queryAssignment(Context.current(), QueryAssignmentRequest.newBuilder() + .build()); + + try { + QueryAssignmentResponse response = future.get(); + assertEquals(Code.INVALID_ARGUMENT.getNumber(), response.getCommon().getStatus().getCode()); + } catch (Exception e) { + assertNull(e); + } + } + + @Test + public void testQueryAssignment() { + RouteService routeService = new RouteService(this.clientManager); + + CompletableFuture future = routeService.queryAssignment(Context.current(), QueryAssignmentRequest.newBuilder() + .setEndpoints(Endpoints.newBuilder() + .addAddresses(Address.newBuilder() + .setPort(80) + .setHost("host") + .build()) + .setScheme(AddressScheme.DOMAIN_NAME) + .build()) + .setTopic(Resource.newBuilder() + .setName("topic") + .build()) + .setGroup(Resource.newBuilder() + .setName("group") + .build()) + .setClientId("clientId") + .build()); + + try { + QueryAssignmentResponse response = future.get(); + assertEquals(Code.OK.getNumber(), response.getCommon().getStatus().getCode()); + assertEquals(1, response.getAssignmentsCount()); + assertEquals("brokerName", response.getAssignments(0).getPartition().getBroker().getName()); + assertEquals("host", response.getAssignments(0).getPartition().getBroker().getEndpoints().getAddresses(0).getHost()); + } catch (Exception e) { + assertNull(e); + } + } +} \ No newline at end of file