diff --git a/pom.xml b/pom.xml index 9caa47358d..b29d63409b 100644 --- a/pom.xml +++ b/pom.xml @@ -434,7 +434,7 @@ org.assertj assertj-core - 2.6.0 + 3.22.0 test diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/common/ProxyExceptionCode.java b/proxy/src/main/java/org/apache/rocketmq/proxy/common/ProxyExceptionCode.java index ae056b2168..a297e50aed 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/common/ProxyExceptionCode.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/common/ProxyExceptionCode.java @@ -20,5 +20,6 @@ public enum ProxyExceptionCode { FORBIDDEN, RECEIPT_HANDLE_EXPIRED, INVALID_BROKER_NAME, - INVALID_RECEIPT_HANDLE + INVALID_RECEIPT_HANDLE, + ILLEGAL_MESSAGE, } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/config/ProxyConfig.java b/proxy/src/main/java/org/apache/rocketmq/proxy/config/ProxyConfig.java index 2127a6c0f7..6bf4ecf653 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/config/ProxyConfig.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/config/ProxyConfig.java @@ -56,6 +56,7 @@ public class ProxyConfig { private int grpcMaxInboundMessageSize = 130 * 1024 * 1024; private int channelExpiredInSeconds = 60; + private int contextExpiredInSeconds = 30; private int rocketmqMQClientNum = 6; @@ -251,6 +252,14 @@ public class ProxyConfig { this.channelExpiredInSeconds = channelExpiredInSeconds; } + public int getContextExpiredInSeconds() { + return contextExpiredInSeconds; + } + + public void setContextExpiredInSeconds(int contextExpiredInSeconds) { + this.contextExpiredInSeconds = contextExpiredInSeconds; + } + public int getRocketmqMQClientNum() { return rocketmqMQClientNum; } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/common/GrpcProxyException.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/common/GrpcProxyException.java index 702ca2c75d..7cd5f5a43d 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/common/GrpcProxyException.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/common/GrpcProxyException.java @@ -31,8 +31,10 @@ public class GrpcProxyException extends RuntimeException { static { CODE_MAPPING.put(ProxyExceptionCode.INVALID_BROKER_NAME, Code.INVALID_RECEIPT_HANDLE); + CODE_MAPPING.put(ProxyExceptionCode.INVALID_RECEIPT_HANDLE, Code.INVALID_RECEIPT_HANDLE); CODE_MAPPING.put(ProxyExceptionCode.RECEIPT_HANDLE_EXPIRED, Code.RECEIPT_HANDLE_EXPIRED); CODE_MAPPING.put(ProxyExceptionCode.FORBIDDEN, Code.FORBIDDEN); + CODE_MAPPING.put(ProxyExceptionCode.ILLEGAL_MESSAGE, Code.ILLEGAL_MESSAGE); } public GrpcProxyException(Code code, String message) { diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/LocalServiceManager.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/LocalServiceManager.java index c55d0acb87..42c9e249d2 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/service/LocalServiceManager.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/LocalServiceManager.java @@ -17,18 +17,23 @@ package org.apache.rocketmq.proxy.service; 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.broker.client.ConsumerManager; import org.apache.rocketmq.broker.client.ProducerManager; +import org.apache.rocketmq.common.ThreadFactoryImpl; import org.apache.rocketmq.proxy.common.AbstractStartAndShutdown; +import org.apache.rocketmq.proxy.common.StartAndShutdown; +import org.apache.rocketmq.proxy.service.channel.ChannelManager; import org.apache.rocketmq.proxy.service.message.LocalMessageService; import org.apache.rocketmq.proxy.service.message.MessageService; -import org.apache.rocketmq.proxy.service.relay.LocalProxyRelayService; -import org.apache.rocketmq.proxy.service.relay.ProxyRelayService; import org.apache.rocketmq.proxy.service.metadata.LocalMetadataService; import org.apache.rocketmq.proxy.service.metadata.MetadataService; import org.apache.rocketmq.proxy.service.mqclient.DoNothingClientRemotingProcessor; import org.apache.rocketmq.proxy.service.mqclient.MQClientAPIFactory; +import org.apache.rocketmq.proxy.service.relay.LocalProxyRelayService; +import org.apache.rocketmq.proxy.service.relay.ProxyRelayService; import org.apache.rocketmq.proxy.service.route.LocalTopicRouteService; import org.apache.rocketmq.proxy.service.route.TopicRouteService; import org.apache.rocketmq.proxy.service.transaction.LocalTransactionService; @@ -45,16 +50,21 @@ public class LocalServiceManager extends AbstractStartAndShutdown implements Ser private final MetadataService metadataService; private final MQClientAPIFactory mqClientAPIFactory; + private final ChannelManager channelManager; + + private final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor( + new ThreadFactoryImpl("LocalServiceManagerScheduledThread")); public LocalServiceManager(BrokerController brokerController, RPCHook rpcHook) { this.brokerController = brokerController; - this.messageService = new LocalMessageService(brokerController, rpcHook); + this.channelManager = new ChannelManager(); + this.messageService = new LocalMessageService(brokerController, channelManager, rpcHook); this.mqClientAPIFactory = new MQClientAPIFactory( "TopicRouteServiceClient_", 1, new DoNothingClientRemotingProcessor(null), rpcHook, - Executors.newSingleThreadScheduledExecutor() + scheduledExecutorService ); this.topicRouteService = new LocalTopicRouteService(brokerController, mqClientAPIFactory); this.transactionService = new LocalTransactionService(); @@ -66,6 +76,7 @@ public class LocalServiceManager extends AbstractStartAndShutdown implements Ser protected void init() { this.appendStartAndShutdown(this.mqClientAPIFactory); this.appendStartAndShutdown(this.topicRouteService); + this.appendStartAndShutdown(new LocalServiceManagerStartAndShutdown()); } @Override @@ -102,4 +113,14 @@ public class LocalServiceManager extends AbstractStartAndShutdown implements Ser public MetadataService getMetadataService() { return this.metadataService; } + + private class LocalServiceManagerStartAndShutdown implements StartAndShutdown { + @Override public void start() throws Exception { + LocalServiceManager.this.scheduledExecutorService.scheduleWithFixedDelay(channelManager::scanAndCleanChannels, 5, 5, TimeUnit.MINUTES); + } + + @Override public void shutdown() throws Exception { + LocalServiceManager.this.scheduledExecutorService.shutdown(); + } + } } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/ChannelManager.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/ChannelManager.java new file mode 100644 index 0000000000..f4b4e1896e --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/ChannelManager.java @@ -0,0 +1,92 @@ +/* + * 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.service.channel; + +import com.google.common.base.Strings; +import java.util.Iterator; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.rocketmq.common.constant.LoggerName; +import org.apache.rocketmq.proxy.common.ContextVariable; +import org.apache.rocketmq.proxy.common.ProxyContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ChannelManager { + private static final Logger log = LoggerFactory.getLogger(LoggerName.PROXY_LOGGER_NAME); + private final ConcurrentMap clientIdChannelMap = new ConcurrentHashMap<>(); + + public SimpleChannel createChannel(ProxyContext context) { + final String clientId = anonymousChannelId(context); + if (Strings.isNullOrEmpty(clientId)) { + log.warn("ClientId is unexpected null or empty"); + return createChannelInner(context); + } + + SimpleChannel channel = clientIdChannelMap.computeIfAbsent(clientId, k -> createChannelInner(context)); + channel.updateLastAccessTime(); + return channel; + } + + public SimpleChannel createInvocationChannel(ProxyContext context) { + final String clientId = anonymousChannelId(InvocationChannel.class.getName(), context); + final String clientHost = context.getVal(ContextVariable.REMOTE_ADDRESS); + final String localAddress = context.getVal(ContextVariable.LOCAL_ADDRESS); + if (Strings.isNullOrEmpty(clientId)) { + log.warn("ClientId is unexpected null or empty"); + return new InvocationChannel(clientHost, localAddress); + } + + SimpleChannel channel = clientIdChannelMap.computeIfAbsent(clientId, k -> new InvocationChannel(clientHost, localAddress)); + channel.updateLastAccessTime(); + return channel; + } + + private String anonymousChannelId(ProxyContext context) { + final String clientHost = context.getVal(ContextVariable.REMOTE_ADDRESS); + final String localAddress = context.getVal(ContextVariable.LOCAL_ADDRESS); + return clientHost + "@" + localAddress; + } + + private String anonymousChannelId(String key, ProxyContext context) { + final String clientHost = context.getVal(ContextVariable.REMOTE_ADDRESS); + final String localAddress = context.getVal(ContextVariable.LOCAL_ADDRESS); + return key + "@" + clientHost + "@" + localAddress; + } + + private SimpleChannel createChannelInner(ProxyContext context) { + return new SimpleChannel(context.getVal(ContextVariable.REMOTE_ADDRESS), context.getVal(ContextVariable.LOCAL_ADDRESS)); + } + + public void scanAndCleanChannels() { + try { + Iterator> iterator = clientIdChannelMap.entrySet().iterator(); + while (iterator.hasNext()) { + Map.Entry entry = iterator.next(); + if (!entry.getValue().isActive()) { + iterator.remove(); + } else { + entry.getValue().clearExpireContext(); + } + } + } catch (Throwable e) { + log.error("Unexpected exception", e); + } + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationChannel.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationChannel.java new file mode 100644 index 0000000000..00e8cea99c --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationChannel.java @@ -0,0 +1,80 @@ +/* + * 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.service.channel; + +import io.netty.channel.ChannelFuture; +import java.util.Iterator; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.rocketmq.proxy.config.ConfigurationManager; +import org.apache.rocketmq.remoting.protocol.RemotingCommand; + +public class InvocationChannel extends SimpleChannel { + protected final ConcurrentMap inFlightRequestMap; + + public InvocationChannel(String remoteAddress, String localAddress) { + super(remoteAddress, localAddress); + this.inFlightRequestMap = new ConcurrentHashMap<>(); + } + + @Override + public ChannelFuture writeAndFlush(Object msg) { + if (msg instanceof RemotingCommand) { + RemotingCommand responseCommand = (RemotingCommand) msg; + InvocationContextInterface context = inFlightRequestMap.remove(responseCommand.getOpaque()); + if (null != context) { + context.handle(responseCommand); + } + inFlightRequestMap.remove(responseCommand.getOpaque()); + } + return super.writeAndFlush(msg); + } + + @Override + public boolean isWritable() { + return inFlightRequestMap.size() > 0; + } + + @Override + public void registerInvocationContext(int opaque, InvocationContextInterface context) { + inFlightRequestMap.put(opaque, context); + } + + @Override + public void eraseInvocationContext(int opaque) { + inFlightRequestMap.remove(opaque); + } + + @Override + public void clearExpireContext() { + Iterator> iterator = inFlightRequestMap.entrySet().iterator(); + int count = 0; + while (iterator.hasNext()) { + Map.Entry entry = iterator.next(); + if (entry.getValue().expired(ConfigurationManager.getProxyConfig().getChannelExpiredInSeconds())) { + iterator.remove(); + count++; + log.debug("An expired request is found, request: {}", entry.getValue()); + } + } + if (count > 0) { + log.warn("[BUG] {} expired in-flight requests is cleaned.", count); + } + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationContext.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationContext.java new file mode 100644 index 0000000000..9fb488eb9b --- /dev/null +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationContext.java @@ -0,0 +1,43 @@ +/* + * 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.service.channel; + +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import org.apache.rocketmq.remoting.protocol.RemotingCommand; + +public class InvocationContext implements InvocationContextInterface { + private final CompletableFuture response; + private final long timestamp = System.currentTimeMillis(); + + public InvocationContext(CompletableFuture resp) { + this.response = resp; + } + + public boolean expired(long expiredTimeSec) { + return System.currentTimeMillis() - timestamp >= Duration.ofSeconds(expiredTimeSec).toMillis(); + } + + public CompletableFuture getResponse() { + return response; + } + + public void handle(RemotingCommand remotingCommand) { + response.complete(remotingCommand); + } +} diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/common/Cleaner.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationContextInterface.java similarity index 76% rename from proxy/src/main/java/org/apache/rocketmq/proxy/common/Cleaner.java rename to proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationContextInterface.java index a02b08a913..0db9516486 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/common/Cleaner.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/InvocationContextInterface.java @@ -15,8 +15,12 @@ * limitations under the License. */ -package org.apache.rocketmq.proxy.common; +package org.apache.rocketmq.proxy.service.channel; -public interface Cleaner { - void clean(); +import org.apache.rocketmq.remoting.protocol.RemotingCommand; + +public interface InvocationContextInterface { + void handle(RemotingCommand remotingCommand); + + boolean expired(long expiredTimeSec); } diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/SimpleChannel.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/SimpleChannel.java index 3e2ea4562a..9f010526b8 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/SimpleChannel.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/channel/SimpleChannel.java @@ -22,6 +22,7 @@ import io.netty.channel.AbstractChannel; import io.netty.channel.Channel; import io.netty.channel.ChannelConfig; import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelMetadata; import io.netty.channel.ChannelOutboundBuffer; import io.netty.channel.DefaultChannelPromise; @@ -45,6 +46,7 @@ public class SimpleChannel extends AbstractChannel { protected final String localAddress; protected long lastAccessTime; + protected ChannelHandlerContext channelHandlerContext; /** * Creates a new instance. @@ -58,6 +60,7 @@ public class SimpleChannel extends AbstractChannel { lastAccessTime = System.currentTimeMillis(); this.remoteAddress = remoteAddress; this.localAddress = localAddress; + this.channelHandlerContext = new SimpleChannelHandlerContext(this); } public SimpleChannel(String remoteAddress, String localAddress) { @@ -165,4 +168,24 @@ public class SimpleChannel extends AbstractChannel { promise.setSuccess(); return promise; } + + public void updateLastAccessTime() { + this.lastAccessTime = System.currentTimeMillis(); + } + + public void registerInvocationContext(int opaque, InvocationContextInterface context) { + + } + + public void eraseInvocationContext(int opaque) { + + } + + public void clearExpireContext() { + + } + + public ChannelHandlerContext getChannelHandlerContext() { + return channelHandlerContext; + } } \ No newline at end of file diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/service/message/LocalMessageService.java b/proxy/src/main/java/org/apache/rocketmq/proxy/service/message/LocalMessageService.java index 055f9e85e1..c18e8fe4d1 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/service/message/LocalMessageService.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/service/message/LocalMessageService.java @@ -16,37 +16,121 @@ */ package org.apache.rocketmq.proxy.service.message; +import io.netty.channel.ChannelHandlerContext; +import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.concurrent.CompletableFuture; import org.apache.rocketmq.broker.BrokerController; import org.apache.rocketmq.client.consumer.AckResult; import org.apache.rocketmq.client.consumer.PopResult; import org.apache.rocketmq.client.producer.SendResult; +import org.apache.rocketmq.client.producer.SendStatus; +import org.apache.rocketmq.common.constant.LoggerName; import org.apache.rocketmq.common.consumer.ReceiptHandle; import org.apache.rocketmq.common.message.Message; +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.header.AckMessageRequestHeader; import org.apache.rocketmq.common.protocol.header.ChangeInvisibleTimeRequestHeader; import org.apache.rocketmq.common.protocol.header.ConsumerSendMsgBackRequestHeader; import org.apache.rocketmq.common.protocol.header.EndTransactionRequestHeader; import org.apache.rocketmq.common.protocol.header.PopMessageRequestHeader; import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader; +import org.apache.rocketmq.common.protocol.header.SendMessageResponseHeader; import org.apache.rocketmq.proxy.common.ProxyContext; +import org.apache.rocketmq.proxy.common.ProxyException; +import org.apache.rocketmq.proxy.common.ProxyExceptionCode; +import org.apache.rocketmq.proxy.service.channel.ChannelManager; +import org.apache.rocketmq.proxy.service.channel.InvocationContext; +import org.apache.rocketmq.proxy.service.channel.SimpleChannel; import org.apache.rocketmq.proxy.service.route.SelectableMessageQueue; import org.apache.rocketmq.proxy.service.transaction.TransactionId; import org.apache.rocketmq.remoting.RPCHook; import org.apache.rocketmq.remoting.protocol.RemotingCommand; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class LocalMessageService implements MessageService { + private static final Logger log = LoggerFactory.getLogger(LoggerName.PROXY_LOGGER_NAME); + private final BrokerController brokerController; + private final ChannelManager channelManager; - private BrokerController brokerController; - - public LocalMessageService(BrokerController brokerController, RPCHook rpcHook) { + public LocalMessageService(BrokerController brokerController, ChannelManager channelManager, RPCHook rpcHook) { this.brokerController = brokerController; + this.channelManager = channelManager; } @Override public CompletableFuture> sendMessage(ProxyContext ctx, SelectableMessageQueue messageQueue, List msgList, SendMessageRequestHeader requestHeader, long timeoutMillis) { - return null; + byte[] body; + String messageId; + if (msgList.size() > 1) { + requestHeader.setBatch(true); + MessageBatch msgBatch = MessageBatch.generateFromList(msgList); + MessageClientIDSetter.setUniqID(msgBatch); + body = msgBatch.encode(); + msgBatch.setBody(body); + messageId = MessageClientIDSetter.getUniqID(msgBatch); + } else { + Message message = msgList.get(0); + body = message.getBody(); + messageId = MessageClientIDSetter.getUniqID(message); + } + RemotingCommand request = LocalRemotingCommand.createRequestCommand(RequestCode.SEND_MESSAGE, requestHeader); + request.setBody(body); + CompletableFuture future = new CompletableFuture<>(); + SimpleChannel channel = channelManager.createInvocationChannel(ctx); + InvocationContext invocationContext = new InvocationContext(future); + channel.registerInvocationContext(request.getOpaque(), invocationContext); + ChannelHandlerContext simpleChannelHandlerContext = channel.getChannelHandlerContext(); + try { + RemotingCommand response = brokerController.getSendMessageProcessor().processRequest(simpleChannelHandlerContext, request); + if (response != null) { + invocationContext.handle(response); + } + } catch (Exception e) { + future.completeExceptionally(e); + log.error("Failed to process send message command", e); + } finally { + channel.eraseInvocationContext(request.getOpaque()); + } + return future.thenApply(r -> { + SendResult sendResult = new SendResult(); + SendMessageResponseHeader responseHeader = (SendMessageResponseHeader) r.readCustomHeader(); + SendStatus sendStatus; + switch (r.getCode()) { + case ResponseCode.FLUSH_DISK_TIMEOUT: { + sendStatus = SendStatus.FLUSH_DISK_TIMEOUT; + break; + } + case ResponseCode.FLUSH_SLAVE_TIMEOUT: { + sendStatus = SendStatus.FLUSH_SLAVE_TIMEOUT; + break; + } + case ResponseCode.SLAVE_NOT_AVAILABLE: { + sendStatus = SendStatus.SLAVE_NOT_AVAILABLE; + break; + } + case ResponseCode.SUCCESS: { + sendStatus = SendStatus.SEND_OK; + break; + } + default: { + throw new ProxyException(ProxyExceptionCode.ILLEGAL_MESSAGE, r.getRemark()); + } + } + sendResult.setSendStatus(sendStatus); + sendResult.setMsgId(messageId); + sendResult.setMessageQueue(new MessageQueue(requestHeader.getTopic(), brokerController.getBrokerConfig().getBrokerName(), requestHeader.getQueueId())); + sendResult.setQueueOffset(responseHeader.getQueueOffset()); + sendResult.setTransactionId(responseHeader.getTransactionId()); + sendResult.setOffsetMsgId(responseHeader.getMsgId()); + return Collections.singletonList(sendResult); + }); } @Override diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/service/message/LocalMessageServiceTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/service/message/LocalMessageServiceTest.java new file mode 100644 index 0000000000..7053e08e26 --- /dev/null +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/service/message/LocalMessageServiceTest.java @@ -0,0 +1,208 @@ +/* + * 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.service.message; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import org.apache.rocketmq.broker.BrokerController; +import org.apache.rocketmq.broker.processor.SendMessageProcessor; +import org.apache.rocketmq.client.producer.SendResult; +import org.apache.rocketmq.client.producer.SendStatus; +import org.apache.rocketmq.common.BrokerConfig; +import org.apache.rocketmq.common.message.Message; +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.header.SendMessageRequestHeader; +import org.apache.rocketmq.common.protocol.header.SendMessageResponseHeader; +import org.apache.rocketmq.proxy.common.ContextVariable; +import org.apache.rocketmq.proxy.common.ProxyContext; +import org.apache.rocketmq.proxy.common.ProxyException; +import org.apache.rocketmq.proxy.common.ProxyExceptionCode; +import org.apache.rocketmq.proxy.config.ConfigurationManager; +import org.apache.rocketmq.proxy.config.InitConfigAndLoggerTest; +import org.apache.rocketmq.proxy.service.channel.ChannelManager; +import org.apache.rocketmq.proxy.service.channel.SimpleChannelHandlerContext; +import org.apache.rocketmq.remoting.exception.RemotingCommandException; +import org.apache.rocketmq.remoting.protocol.RemotingCommand; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowableOfType; + +@RunWith(MockitoJUnitRunner.class) +public class LocalMessageServiceTest extends InitConfigAndLoggerTest { + private LocalMessageService localMessageService; + @Mock + private SendMessageProcessor sendMessageProcessorMock; + @Mock + private BrokerController brokerControllerMock; + + private ProxyContext proxyContext; + + private ChannelManager channelManager; + + private String topic = "topic"; + + private int queueId = 0; + + private long queueOffset = 0L; + + private String transactionId = "transactionId"; + + private String offsetMessageId = "offsetMessageId"; + + @Before + public void setUp() throws Throwable { + super.before(); + ConfigurationManager.getProxyConfig().setNameSrvAddr("1.1.1.1"); + channelManager = new ChannelManager(); + Mockito.when(brokerControllerMock.getSendMessageProcessor()).thenReturn(sendMessageProcessorMock); + Mockito.when(brokerControllerMock.getBrokerConfig()).thenReturn(new BrokerConfig()); + localMessageService = new LocalMessageService(brokerControllerMock, channelManager, null); + proxyContext = ProxyContext.create().withVal(ContextVariable.REMOTE_ADDRESS, "0.0.0.1") + .withVal(ContextVariable.LOCAL_ADDRESS, "0.0.0.2"); + } + + @Test + public void testSendMessageWriteAndFlush() throws Exception { + Message message = new Message(topic, "body".getBytes(StandardCharsets.UTF_8)); + MessageClientIDSetter.setUniqID(message); + List messagesList = Collections.singletonList(message); + SendMessageRequestHeader requestHeader = new SendMessageRequestHeader(); + requestHeader.setTopic(topic); + requestHeader.setQueueId(queueId); + Mockito.when(sendMessageProcessorMock.processRequest(Mockito.any(SimpleChannelHandlerContext.class), Mockito.argThat(argument -> { + boolean first = argument.getCode() == RequestCode.SEND_MESSAGE; + boolean second = Arrays.equals(argument.getBody(), message.getBody()); + return first & second; + }))).thenAnswer(invocation -> { + SimpleChannelHandlerContext simpleChannelHandlerContext = invocation.getArgument(0); + RemotingCommand request = invocation.getArgument(1); + RemotingCommand response = RemotingCommand.createResponseCommand(SendMessageResponseHeader.class); + response.setOpaque(request.getOpaque()); + response.setCode(ResponseCode.SUCCESS); + response.setBody(message.getBody()); + SendMessageResponseHeader sendMessageResponseHeader = (SendMessageResponseHeader) response.readCustomHeader(); + sendMessageResponseHeader.setQueueId(queueId); + sendMessageResponseHeader.setQueueOffset(queueOffset); + sendMessageResponseHeader.setMsgId(offsetMessageId); + sendMessageResponseHeader.setTransactionId(transactionId); + simpleChannelHandlerContext.writeAndFlush(response); + return null; + }); + + CompletableFuture future = localMessageService.sendMessage(proxyContext, null, messagesList, requestHeader, 1000L); + SendResult sendResult = future.get(); + assertThat(sendResult.getSendStatus()).isEqualTo(SendStatus.SEND_OK); + assertThat(sendResult.getMsgId()).isEqualTo(MessageClientIDSetter.getUniqID(message)); + assertThat(sendResult.getMessageQueue()) + .isEqualTo(new MessageQueue(topic, brokerControllerMock.getBrokerConfig().getBrokerName(), queueId)); + assertThat(sendResult.getQueueOffset()).isEqualTo(queueOffset); + assertThat(sendResult.getTransactionId()).isEqualTo(transactionId); + assertThat(sendResult.getOffsetMsgId()).isEqualTo(offsetMessageId); + } + + @Test + public void testSendBatchMessageWriteAndFlush() throws Exception { + Message message1 = new Message(topic, "body1".getBytes(StandardCharsets.UTF_8)); + Message message2 = new Message(topic, "body2".getBytes(StandardCharsets.UTF_8)); + MessageClientIDSetter.setUniqID(message1); + MessageClientIDSetter.setUniqID(message2); + List messagesList = Arrays.asList(message1, message2); + MessageBatch msgBatch = MessageBatch.generateFromList(messagesList); + MessageClientIDSetter.setUniqID(msgBatch); + byte[] body = msgBatch.encode(); + msgBatch.setBody(body); + SendMessageRequestHeader requestHeader = new SendMessageRequestHeader(); + requestHeader.setTopic(topic); + requestHeader.setQueueId(queueId); + Mockito.when(sendMessageProcessorMock.processRequest(Mockito.any(SimpleChannelHandlerContext.class), Mockito.argThat(argument -> { + boolean first = argument.getCode() == RequestCode.SEND_MESSAGE; + boolean second = Arrays.equals(argument.getBody(), body); + return first & second; + }))).thenAnswer(invocation -> { + SimpleChannelHandlerContext simpleChannelHandlerContext = invocation.getArgument(0); + RemotingCommand request = invocation.getArgument(1); + RemotingCommand response = RemotingCommand.createResponseCommand(SendMessageResponseHeader.class); + response.setOpaque(request.getOpaque()); + response.setCode(ResponseCode.SUCCESS); + response.setBody(body); + SendMessageResponseHeader sendMessageResponseHeader = (SendMessageResponseHeader) response.readCustomHeader(); + sendMessageResponseHeader.setQueueId(queueId); + sendMessageResponseHeader.setQueueOffset(queueOffset); + sendMessageResponseHeader.setMsgId(offsetMessageId); + sendMessageResponseHeader.setTransactionId(transactionId); + simpleChannelHandlerContext.writeAndFlush(response); + return null; + }); + + CompletableFuture future = localMessageService.sendMessage(proxyContext, null, messagesList, requestHeader, 1000L); + SendResult sendResult = future.get(); + assertThat(sendResult.getSendStatus()).isEqualTo(SendStatus.SEND_OK); + assertThat(sendResult.getMessageQueue()) + .isEqualTo(new MessageQueue(topic, brokerControllerMock.getBrokerConfig().getBrokerName(), queueId)); + assertThat(sendResult.getQueueOffset()).isEqualTo(queueOffset); + assertThat(sendResult.getTransactionId()).isEqualTo(transactionId); + assertThat(sendResult.getOffsetMsgId()).isEqualTo(offsetMessageId); + } + + @Test + public void testSendMessageError() throws Exception { + RemotingCommand response = RemotingCommand.createResponseCommand(SendMessageResponseHeader.class); + response.setCode(ResponseCode.SYSTEM_ERROR); + Message message = new Message("topic", "body".getBytes(StandardCharsets.UTF_8)); + MessageClientIDSetter.setUniqID(message); + List messagesList = Collections.singletonList(message); + SendMessageRequestHeader sendMessageRequestHeader = new SendMessageRequestHeader(); + sendMessageRequestHeader.setTopic(topic); + sendMessageRequestHeader.setQueueId(queueId); + + Mockito.when(sendMessageProcessorMock.processRequest(Mockito.any(SimpleChannelHandlerContext.class), Mockito.any(RemotingCommand.class))) + .thenReturn(response); + + CompletableFuture future = localMessageService.sendMessage(proxyContext, null, messagesList, sendMessageRequestHeader, 1000L); + ExecutionException exception = catchThrowableOfType(future::get, ExecutionException.class); + assertThat(exception.getCause()).isInstanceOf(ProxyException.class); + assertThat(((ProxyException) exception.getCause()).getCode()).isEqualTo(ProxyExceptionCode.ILLEGAL_MESSAGE); + } + + @Test + public void testSendMessageWithException() throws Exception { + Mockito.when(sendMessageProcessorMock.processRequest(Mockito.any(SimpleChannelHandlerContext.class), Mockito.any(RemotingCommand.class))) + .thenThrow(new RemotingCommandException("test")); + Message message = new Message("topic", "body".getBytes(StandardCharsets.UTF_8)); + MessageClientIDSetter.setUniqID(message); + List messagesList = Collections.singletonList(message); + SendMessageRequestHeader sendMessageRequestHeader = new SendMessageRequestHeader(); + CompletableFuture future = localMessageService.sendMessage(proxyContext, null, messagesList, sendMessageRequestHeader, 1000L); + ExecutionException exception = catchThrowableOfType(future::get, ExecutionException.class); + assertThat(exception.getCause()).isInstanceOf(RemotingCommandException.class); + } +} \ No newline at end of file