[ISSUE #3949] Add sendMessage in LocalMessageService

pom.xml

@@ -434,7 +434,7 @@
         <dependency>
             <groupId>org.assertj</groupId>
             <artifactId>assertj-core</artifactId>
-            <version>2.6.0</version>
+            <version>3.22.0</version>
             <scope>test</scope>
         </dependency>
         <dependency>

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,
 }

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;
     }

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) {

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();
+        }
+    }
 }

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<String, SimpleChannel> 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<Map.Entry<String, SimpleChannel>> iterator = clientIdChannelMap.entrySet().iterator();
+            while (iterator.hasNext()) {
+                Map.Entry<String, SimpleChannel> entry = iterator.next();
+                if (!entry.getValue().isActive()) {
+                    iterator.remove();
+                } else {
+                    entry.getValue().clearExpireContext();
+                }
+            }
+        } catch (Throwable e) {
+            log.error("Unexpected exception", e);
+        }
+    }
+}

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<Integer, InvocationContextInterface> 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<Map.Entry<Integer, InvocationContextInterface>> iterator = inFlightRequestMap.entrySet().iterator();
+        int count = 0;
+        while (iterator.hasNext()) {
+            Map.Entry<Integer, InvocationContextInterface> 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);
+        }
+    }
+}

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<RemotingCommand> response;
+    private final long timestamp = System.currentTimeMillis();
+
+    public InvocationContext(CompletableFuture<RemotingCommand> resp) {
+        this.response = resp;
+    }
+
+    public boolean expired(long expiredTimeSec) {
+        return System.currentTimeMillis() - timestamp >= Duration.ofSeconds(expiredTimeSec).toMillis();
+    }
+
+    public CompletableFuture<RemotingCommand> getResponse() {
+        return response;
+    }
+
+    public void handle(RemotingCommand remotingCommand) {
+        response.complete(remotingCommand);
+    }
+}

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);
 }

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;
+    }
 }

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<List<SendResult>> sendMessage(ProxyContext ctx, SelectableMessageQueue messageQueue,
         List<? extends Message> 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<RemotingCommand> 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

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<Message> 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<SendResult> 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<Message> 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<SendResult> 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<Message> 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<SendResult> 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<Message> messagesList = Collections.singletonList(message);
+        SendMessageRequestHeader sendMessageRequestHeader = new SendMessageRequestHeader();
+        CompletableFuture<SendResult> future = localMessageService.sendMessage(proxyContext, null, messagesList, sendMessageRequestHeader, 1000L);
+        ExecutionException exception = catchThrowableOfType(future::get, ExecutionException.class);
+        assertThat(exception.getCause()).isInstanceOf(RemotingCommandException.class);
+    }
+}