fix: release inbound MqttPublishMessage payload ByteBuf and retain per subscriber write

shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandler.java

@@ -20,6 +20,7 @@
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.mqtt.MqttMessage;
+import io.netty.util.ReferenceCountUtil;
 import io.netty.util.concurrent.Future;
 import io.netty.util.concurrent.GenericFutureListener;
 
@@ -30,11 +31,15 @@ public class MqttTransportHandler extends ChannelInboundHandlerAdapter implement
 
     @Override
     public void channelRead(final ChannelHandlerContext ctx, final Object msg) throws Exception {
-        if (msg instanceof MqttMessage) {
-            MqttFactory mqttFactory = new MqttFactory((MqttMessage) msg, ctx);
-            mqttFactory.connect();
-        } else {
-            ctx.close();
+        try {
+            if (msg instanceof MqttMessage) {
+                MqttFactory mqttFactory = new MqttFactory((MqttMessage) msg, ctx);
+                mqttFactory.connect();
+            } else {
+                ctx.close();
+            }
+        } finally {
+            ReferenceCountUtil.release(msg);
         }
     }
 

shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/Publish.java

@@ -18,7 +18,6 @@
 package org.apache.shenyu.protocol.mqtt;
 
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.mqtt.MqttFixedHeader;
@@ -28,6 +27,7 @@
 import io.netty.handler.codec.mqtt.MqttPubAckMessage;
 import io.netty.handler.codec.mqtt.MqttMessageType;
 import io.netty.handler.codec.mqtt.MqttPublishVariableHeader;
+import io.netty.util.ReferenceCountUtil;
 import org.apache.shenyu.common.utils.Singleton;
 import org.apache.shenyu.protocol.mqtt.repositories.SubscribeRepository;
 import org.apache.shenyu.protocol.mqtt.repositories.TopicRepository;
@@ -50,8 +50,6 @@ public void publish(final ChannelHandlerContext ctx, final MqttPublishMessage ms
         String topic = msg.variableHeader().topicName();
         ByteBuf payload = msg.payload();
         String message = byteBufToString(payload);
-        //// todo qos
-        MqttQoS mqttQoS = msg.fixedHeader().qosLevel();
         if (msg.fixedHeader().isRetain()) {
             if (payload.isReadable()) {
                 Singleton.INST.get(TopicRepository.class).add(topic, message);
@@ -60,8 +58,18 @@ public void publish(final ChannelHandlerContext ctx, final MqttPublishMessage ms
             }
         }
         int packetId = msg.variableHeader().packetId();
-        CompletableFuture.runAsync(() -> send(topic, payload, packetId));
+        // The inbound message is released by MqttTransportHandler once publish returns, retain the payload for the asynchronous send.
+        payload.retain();
+        CompletableFuture.runAsync(() -> {
+            try {
+                send(topic, payload, packetId);
+            } finally {
+                ReferenceCountUtil.safeRelease(payload);
+            }
+        });
 
+        //// todo qos
+        MqttQoS mqttQoS = msg.fixedHeader().qosLevel();
         switch (mqttQoS.value()) {
             case 0:
                 break;
@@ -125,7 +133,7 @@ private void send(final String topic, final ByteBuf payload, final int packetId)
             if (channel.isActive()) {
                 MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, MqttQoS.AT_MOST_ONCE, false, 0);
                 MqttPublishVariableHeader mqttPublishVariableHeader = new MqttPublishVariableHeader(topic, packetId);
-                MqttPublishMessage mqttPublishMessage = new MqttPublishMessage(mqttFixedHeader, mqttPublishVariableHeader, Unpooled.wrappedBuffer(payload));
+                MqttPublishMessage mqttPublishMessage = new MqttPublishMessage(mqttFixedHeader, mqttPublishVariableHeader, payload.retainedDuplicate());
                 channel.writeAndFlush(mqttPublishMessage);
             }
         });

shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandlerTest.java

@@ -0,0 +1,76 @@
+/*
+ * 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.shenyu.protocol.mqtt;
+
+import io.netty.buffer.Unpooled;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.mqtt.MqttFixedHeader;
+import io.netty.handler.codec.mqtt.MqttMessageType;
+import io.netty.handler.codec.mqtt.MqttPublishMessage;
+import io.netty.handler.codec.mqtt.MqttPublishVariableHeader;
+import io.netty.handler.codec.mqtt.MqttQoS;
+import io.netty.util.CharsetUtil;
+import io.netty.util.IllegalReferenceCountException;
+import org.apache.shenyu.common.utils.Singleton;
+import org.apache.shenyu.protocol.mqtt.repositories.SubscribeRepository;
+import org.apache.shenyu.protocol.mqtt.repositories.TopicRepository;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+
+import static org.awaitility.Awaitility.await;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+/**
+ * Test cases for {@link MqttTransportHandler}.
+ */
+public final class MqttTransportHandlerTest {
+
+    @BeforeAll
+    static void setUp() {
+        Singleton.INST.single(TopicRepository.class, new TopicRepository());
+        Singleton.INST.single(SubscribeRepository.class, new SubscribeRepository());
+    }
+
+    @Test
+    public void channelReadReleasesInboundMessage() throws Exception {
+        MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, MqttQoS.AT_MOST_ONCE, false, 0);
+        MqttPublishVariableHeader variableHeader = new MqttPublishVariableHeader("test/topic", 1);
+        MqttPublishMessage msg = new MqttPublishMessage(fixedHeader, variableHeader, Unpooled.copiedBuffer("hello", CharsetUtil.UTF_8));
+        new MqttTransportHandler().channelRead(mock(ChannelHandlerContext.class), msg);
+        await().atMost(Duration.ofSeconds(5))
+                .until(() -> {
+                    try {
+                        msg.payload().refCnt();
+                        return false;
+                    } catch (IllegalReferenceCountException e) {
+                        // refCnt() throws once the payload has been fully released.
+                        return true;
+                    }
+                });
+    }
+
+    @Test
+    public void channelReadClosesChannelForNonMqttMessage() throws Exception {
+        ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+        new MqttTransportHandler().channelRead(ctx, new Object());
+        verify(ctx).close();
+    }
+}

shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/PublishTest.java

@@ -18,24 +18,32 @@
 package org.apache.shenyu.protocol.mqtt;
 
 import io.netty.buffer.Unpooled;
+import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.mqtt.MqttFixedHeader;
 import io.netty.handler.codec.mqtt.MqttMessageType;
 import io.netty.handler.codec.mqtt.MqttPublishMessage;
 import io.netty.handler.codec.mqtt.MqttPublishVariableHeader;
 import io.netty.handler.codec.mqtt.MqttQoS;
 import io.netty.util.CharsetUtil;
+import io.netty.util.ReferenceCountUtil;
 import org.apache.shenyu.common.utils.Singleton;
 import org.apache.shenyu.protocol.mqtt.repositories.SubscribeRepository;
 import org.apache.shenyu.protocol.mqtt.repositories.TopicRepository;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 
 import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
 
 import static org.awaitility.Awaitility.await;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 /**
  * Test cases for {@link Publish}.
@@ -80,6 +88,28 @@ public void zeroByteRetainedPublishClearsRetainedMessage() {
         assertNull(topicRepository.get(CLEARED_TOPIC));
     }
 
+    @Test
+    public void publishDeliversPayloadToEachSubscriber() {
+        Channel channel1 = mock(Channel.class);
+        when(channel1.isActive()).thenReturn(true);
+        Channel channel2 = mock(Channel.class);
+        when(channel2.isActive()).thenReturn(true);
+        SubscribeRepository subscribeRepository = Singleton.INST.get(SubscribeRepository.class);
+        subscribeRepository.add(Collections.singletonList("test/fanout"), Arrays.asList(channel1, channel2));
+        await().atMost(Duration.ofSeconds(5))
+                .until(() -> subscribeRepository.get("test/fanout").contains(channel1) && subscribeRepository.get("test/fanout").contains(channel2));
+        MqttPublishMessage msg = publishMessage("test/fanout", "hello", false);
+        new Publish().publish(mock(ChannelHandlerContext.class), msg);
+        // one reference held by the inbound message plus one per active subscriber after the send completes.
+        await().atMost(Duration.ofSeconds(5))
+                .until(() -> msg.payload().refCnt() == 3);
+        ArgumentCaptor<MqttPublishMessage> captor = ArgumentCaptor.forClass(MqttPublishMessage.class);
+        verify(channel1).writeAndFlush(captor.capture());
+        verify(channel2).writeAndFlush(captor.capture());
+        captor.getAllValues().forEach(message -> assertEquals("hello", message.payload().toString(CharsetUtil.UTF_8)));
+        captor.getAllValues().forEach(ReferenceCountUtil::safeRelease);
+    }
+
     private MqttPublishMessage publishMessage(final String topic, final String payload, final boolean retain) {
         MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, MqttQoS.AT_MOST_ONCE, retain, 0);
         MqttPublishVariableHeader variableHeader = new MqttPublishVariableHeader(topic, 1);