From f5df79939721be4b437e04342e2cd57677beb509 Mon Sep 17 00:00:00 2001 From: zhouxiang Date: Tue, 15 Mar 2022 11:56:11 +0800 Subject: [PATCH] [ISSUE #3949] Add heartbeat in LocalGrpcService --- .../grpc/adapter/channel/ChannelType.java | 29 +++++ .../rocketmq/proxy/grpc/common/Converter.java | 113 ++++++++++++++++++ .../proxy/grpc/common/ResponseBuilder.java | 7 ++ .../proxy/grpc/service/LocalGrpcService.java | 34 ++++-- .../grpc/service/LocalGrpcServiceTest.java | 46 ++++++- 5 files changed, 216 insertions(+), 13 deletions(-) create mode 100644 proxy/src/main/java/org/apache/rocketmq/proxy/grpc/adapter/channel/ChannelType.java diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/adapter/channel/ChannelType.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/adapter/channel/ChannelType.java new file mode 100644 index 0000000000..5a883b6a35 --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/adapter/channel/ChannelType.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.adapter.channel; + +public enum ChannelType { + /** + * The channel connect to current proxy directly + */ + LOCAL, + /** + * The channel sync from other proxy + */ + REMOTE +} \ No newline at end of file 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 bcdabe1663..9ff4673763 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 @@ -17,27 +17,50 @@ package org.apache.rocketmq.proxy.grpc.common; +import apache.rocketmq.v1.ConsumeMessageType; +import apache.rocketmq.v1.ConsumeModel; +import apache.rocketmq.v1.ConsumePolicy; +import apache.rocketmq.v1.ConsumerData; import apache.rocketmq.v1.Encoding; +import apache.rocketmq.v1.FilterExpression; +import apache.rocketmq.v1.FilterType; +import apache.rocketmq.v1.HeartbeatRequest; import apache.rocketmq.v1.Message; import apache.rocketmq.v1.MessageType; +import apache.rocketmq.v1.ProducerData; import apache.rocketmq.v1.Resource; import apache.rocketmq.v1.SendMessageRequest; +import apache.rocketmq.v1.SubscriptionEntry; import apache.rocketmq.v1.SystemAttribute; import com.google.common.collect.Maps; import com.google.protobuf.Duration; import com.google.protobuf.Timestamp; import com.google.protobuf.util.Durations; import com.google.protobuf.util.Timestamps; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; +import org.apache.rocketmq.common.constant.LoggerName; +import org.apache.rocketmq.common.consumer.ConsumeFromWhere; +import org.apache.rocketmq.common.filter.ExpressionType; +import org.apache.rocketmq.common.filter.FilterAPI; import org.apache.rocketmq.common.message.MessageAccessor; import org.apache.rocketmq.common.message.MessageConst; import org.apache.rocketmq.common.message.MessageDecoder; import org.apache.rocketmq.common.protocol.NamespaceUtil; import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader; +import org.apache.rocketmq.common.protocol.heartbeat.ConsumeType; +import org.apache.rocketmq.common.protocol.heartbeat.HeartbeatData; +import org.apache.rocketmq.common.protocol.heartbeat.MessageModel; +import org.apache.rocketmq.common.protocol.heartbeat.SubscriptionData; import org.apache.rocketmq.common.sysflag.MessageSysFlag; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class Converter { + private static final Logger LOGGER = LoggerFactory.getLogger(LoggerName.GRPC_LOGGER_NAME); + public static String getResourceNameWithNamespace(Resource resource) { return NamespaceUtil.wrapNamespace(resource.getResourceNamespace(), resource.getName()); } @@ -149,4 +172,94 @@ public class Converter { } return messageWithHeader.getProperties(); } + + public static String buildExpressionType(FilterType filterType) { + switch (filterType) { + case SQL: + return ExpressionType.SQL92; + case TAG: + default: + return ExpressionType.TAG; + } + } + + public static HeartbeatData buildHeartbeatData(HeartbeatRequest request) { + HeartbeatData heartbeatData = new HeartbeatData(); + heartbeatData.setClientID(request.getClientId()); + Set producerDataSet = new HashSet<>(); + producerDataSet.add(buildProducerData(request.getProducerData())); + heartbeatData.setProducerDataSet(producerDataSet); + Set consumerDataSet = new HashSet<>(); + consumerDataSet.add(buildConsumerData(request.getConsumerData())); + heartbeatData.setConsumerDataSet(consumerDataSet); + return heartbeatData; + } + + public static org.apache.rocketmq.common.protocol.heartbeat.ProducerData buildProducerData(ProducerData producerData) { + org.apache.rocketmq.common.protocol.heartbeat.ProducerData buildProducerData = new org.apache.rocketmq.common.protocol.heartbeat.ProducerData(); + buildProducerData.setGroupName(getResourceNameWithNamespace(producerData.getGroup())); + return buildProducerData; + } + + public static org.apache.rocketmq.common.protocol.heartbeat.ConsumerData buildConsumerData(ConsumerData consumerData) { + org.apache.rocketmq.common.protocol.heartbeat.ConsumerData buildConsumerData = new org.apache.rocketmq.common.protocol.heartbeat.ConsumerData(); + buildConsumerData.setGroupName(getResourceNameWithNamespace(consumerData.getGroup())); + buildConsumerData.setConsumeType(buildConsumeType(consumerData.getConsumeType())); + buildConsumerData.setMessageModel(buildMessageModel(consumerData.getConsumeModel())); + buildConsumerData.setConsumeFromWhere(buildConsumeFromWhere(consumerData.getConsumePolicy())); + Set subscriptionDataSet = buildSubscriptionDataSet(consumerData.getSubscriptionsList()); + buildConsumerData.setSubscriptionDataSet(subscriptionDataSet); + return buildConsumerData; + } + + public static ConsumeType buildConsumeType(ConsumeMessageType consumeMessageType) { + switch (consumeMessageType) { + case ACTIVE: + return ConsumeType.CONSUME_ACTIVELY; + case PASSIVE: + default: + return ConsumeType.CONSUME_PASSIVELY; + } + } + + public static MessageModel buildMessageModel(ConsumeModel consumeModel) { + switch (consumeModel) { + case BROADCASTING: + return MessageModel.BROADCASTING; + case CLUSTERING: + default: + return MessageModel.CLUSTERING; + } + } + + public static ConsumeFromWhere buildConsumeFromWhere(ConsumePolicy policy) { + switch (policy) { + case PLAYBACK: + return ConsumeFromWhere.CONSUME_FROM_FIRST_OFFSET; + case DISCARD: + return ConsumeFromWhere.CONSUME_FROM_MAX_OFFSET; + case TARGET_TIMESTAMP: + return ConsumeFromWhere.CONSUME_FROM_TIMESTAMP; + case RESUME: + default: + return ConsumeFromWhere.CONSUME_FROM_LAST_OFFSET; + } + } + + public static Set buildSubscriptionDataSet(List subscriptionEntryList) { + Set subscriptionDataSet = new HashSet<>(); + for (SubscriptionEntry sub : subscriptionEntryList) { + String topicName = Converter.getResourceNameWithNamespace(sub.getTopic()); + FilterExpression filterExpression = sub.getExpression(); + String expression = filterExpression.getExpression(); + String expressionType = Converter.buildExpressionType(filterExpression.getType()); + try { + SubscriptionData subscriptionData = FilterAPI.build(topicName, expression, expressionType); + subscriptionDataSet.add(subscriptionData); + } catch (Exception e) { + throw new IllegalArgumentException("Build subscription failed when apply heartbeat", e); + } + } + return subscriptionDataSet; + } } 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 d8f87c6b89..25eddefd7b 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 @@ -17,6 +17,7 @@ package org.apache.rocketmq.proxy.grpc.common; +import apache.rocketmq.v1.HeartbeatResponse; import apache.rocketmq.v1.ResponseCommon; import apache.rocketmq.v1.SendMessageResponse; import com.google.rpc.Code; @@ -49,6 +50,12 @@ public class ResponseBuilder { .build(); } + public static HeartbeatResponse buildHeartbeatResponse(RemotingCommand command) { + return HeartbeatResponse.newBuilder() + .setCommon(buildCommon(command.getCode(), command.getRemark())) + .build(); + } + public static SendMessageResponse buildSendMessageResponse(RemotingCommand command) { SendMessageResponseHeader responseHeader = (SendMessageResponseHeader) command.readCustomHeader(); String messageId = ""; diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcService.java index a2e140d91d..4277c95a0e 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcService.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcService.java @@ -59,17 +59,22 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import org.apache.rocketmq.broker.BrokerController; +import org.apache.rocketmq.common.MQVersion; import org.apache.rocketmq.common.ThreadFactoryImpl; import org.apache.rocketmq.common.constant.LoggerName; import org.apache.rocketmq.common.protocol.RequestCode; import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader; -import org.apache.rocketmq.proxy.grpc.adapter.InvocationContext; +import org.apache.rocketmq.common.protocol.heartbeat.HeartbeatData; import org.apache.rocketmq.proxy.channel.ChannelManager; -import org.apache.rocketmq.proxy.grpc.adapter.channel.SendMessageChannel; +import org.apache.rocketmq.proxy.channel.SimpleChannel; import org.apache.rocketmq.proxy.channel.SimpleChannelHandlerContext; +import org.apache.rocketmq.proxy.grpc.adapter.InvocationContext; +import org.apache.rocketmq.proxy.grpc.adapter.channel.SendMessageChannel; import org.apache.rocketmq.proxy.grpc.adapter.handler.SendMessageResponseHandler; import org.apache.rocketmq.proxy.grpc.common.Converter; +import org.apache.rocketmq.proxy.grpc.common.InterceptorConstants; import org.apache.rocketmq.proxy.grpc.common.ResponseBuilder; +import org.apache.rocketmq.remoting.protocol.LanguageCode; import org.apache.rocketmq.remoting.protocol.RemotingCommand; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -80,11 +85,11 @@ public class LocalGrpcService implements GrpcForwardService { private final BrokerController brokerController; private final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor( new ThreadFactoryImpl("LocalGrpcServiceScheduledThread")); - private final ChannelManager sendChannelManager; + private final ChannelManager channelManager; public LocalGrpcService(BrokerController brokerController) { this.brokerController = brokerController; - this.sendChannelManager = new ChannelManager(); + this.channelManager = new ChannelManager(); } @Override public CompletableFuture queryRoute(Context ctx, QueryRouteRequest request) { @@ -92,8 +97,23 @@ public class LocalGrpcService implements GrpcForwardService { } @Override public CompletableFuture heartbeat(Context ctx, HeartbeatRequest request) { + LanguageCode languageCode; + String language = InterceptorConstants.METADATA.get(Context.current()).get(InterceptorConstants.LANGUAGE); + languageCode = LanguageCode.valueOf(language); + HeartbeatData heartbeatData = Converter.buildHeartbeatData(request); - return null; + SimpleChannel channel = channelManager.createChannel(); + SimpleChannelHandlerContext simpleChannelHandlerContext = new SimpleChannelHandlerContext(channel); + RemotingCommand command = RemotingCommand.createRequestCommand(RequestCode.HEART_BEAT, null); + command.setLanguage(languageCode); + command.setVersion(MQVersion.Version.V5_0_0.ordinal()); + command.setBody(heartbeatData.encode()); + command.makeCustomHeaderToNet(); + + RemotingCommand response = this.brokerController.getClientManageProcessor() + .heartBeat(simpleChannelHandlerContext, command); + HeartbeatResponse heartbeatResponse = ResponseBuilder.buildHeartbeatResponse(response); + return CompletableFuture.completedFuture(heartbeatResponse); } @Override public CompletableFuture healthCheck(Context ctx, HealthCheckRequest request) { @@ -113,7 +133,7 @@ public class LocalGrpcService implements GrpcForwardService { command.makeCustomHeaderToNet(); SendMessageResponseHandler handler = new SendMessageResponseHandler(message.getSystemAttribute().getMessageId()); - SendMessageChannel channel = SendMessageChannel.create(sendChannelManager.createChannel(), handler); + SendMessageChannel channel = SendMessageChannel.create(channelManager.createChannel(), handler); SimpleChannelHandlerContext channelHandlerContext = new SimpleChannelHandlerContext(channel); CompletableFuture future = new CompletableFuture<>(); InvocationContext context @@ -205,6 +225,6 @@ public class LocalGrpcService implements GrpcForwardService { } private void scanAndCleanChannels() { - this.sendChannelManager.scanAndCleanChannels(); + this.channelManager.scanAndCleanChannels(); } } diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcServiceTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcServiceTest.java index 781274c0a0..a16bf2cc42 100644 --- a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcServiceTest.java +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/service/LocalGrpcServiceTest.java @@ -17,6 +17,8 @@ package org.apache.rocketmq.proxy.grpc.service; +import apache.rocketmq.v1.HeartbeatRequest; +import apache.rocketmq.v1.HeartbeatResponse; import apache.rocketmq.v1.Message; import apache.rocketmq.v1.SendMessageRequest; import apache.rocketmq.v1.SendMessageResponse; @@ -27,6 +29,7 @@ import io.grpc.Metadata; import io.netty.channel.ChannelHandlerContext; import java.util.concurrent.CompletableFuture; import org.apache.rocketmq.broker.BrokerController; +import org.apache.rocketmq.broker.processor.ClientManageProcessor; import org.apache.rocketmq.broker.processor.SendMessageProcessor; import org.apache.rocketmq.common.protocol.ResponseCode; import org.apache.rocketmq.proxy.configuration.InitConfigurationTest; @@ -46,14 +49,48 @@ import static org.assertj.core.api.Assertions.assertThat; public class LocalGrpcServiceTest extends InitConfigurationTest { private LocalGrpcService localGrpcService; @Mock - SendMessageProcessor sendMessageProcessorMock; + private SendMessageProcessor sendMessageProcessorMock; + @Mock + private BrokerController brokerControllerMock; + + private Metadata metadata; @Before public void setUp() throws Exception { - super.before(); - BrokerController brokerControllerMock = Mockito.mock(BrokerController.class); + String mockProxyHome = "/mock/rmq/proxy/home"; + URL mockProxyHomeURL = getClass().getClassLoader().getResource("rmq-proxy-home"); + if (mockProxyHomeURL != null) { + mockProxyHome = mockProxyHomeURL.toURI().getPath(); + } + System.setProperty(RMQ_PROXY_HOME, mockProxyHome); + ConfigurationManager.initEnv(); + ConfigurationManager.intConfig(); + ConfigurationManager.initEnv(); + ConfigurationManager.intConfig(); Mockito.when(brokerControllerMock.getSendMessageProcessor()).thenReturn(sendMessageProcessorMock); localGrpcService = new LocalGrpcService(brokerControllerMock); + metadata = new Metadata(); + metadata.put(InterceptorConstants.REMOTE_ADDRESS, "1.1.1.1"); + metadata.put(InterceptorConstants.LOCAL_ADDRESS, "0.0.0.0"); + metadata.put(InterceptorConstants.LANGUAGE, "JAVA"); + } + + @Test + public void testHeartbeat() { + RemotingCommand response = RemotingCommand.createResponseCommand(ResponseCode.SUCCESS, null); + ClientManageProcessor clientManageProcessorMock = Mockito.mock(ClientManageProcessor.class); + Mockito.when(clientManageProcessorMock.heartBeat(Mockito.any(ChannelHandlerContext.class), Mockito.any(RemotingCommand.class))) + .thenReturn(response); + Mockito.when(brokerControllerMock.getClientManageProcessor()).thenReturn(clientManageProcessorMock); + HeartbeatRequest request = HeartbeatRequest.newBuilder().getDefaultInstanceForType(); + CompletableFuture grpcFuture = localGrpcService.heartbeat( + Context.current().withValue(InterceptorConstants.METADATA, metadata).attach(), request); + grpcFuture.thenAccept(r -> { + assertThat(r.getCommon().getStatus().getCode()) + .isEqualTo(Code.OK.getNumber()); + assertThat(r.getCommon().getStatus().getMessage()) + .isEqualTo(null); + }); } @Test @@ -70,9 +107,6 @@ public class LocalGrpcServiceTest extends InitConfigurationTest { .build()) .build()) .build(); - Metadata metadata = new Metadata(); - metadata.put(InterceptorConstants.REMOTE_ADDRESS, "1.1.1.1"); - metadata.put(InterceptorConstants.LOCAL_ADDRESS, "0.0.0.0"); CompletableFuture grpcFuture = localGrpcService.sendMessage( Context.current().withValue(InterceptorConstants.METADATA, metadata).attach(), request);