diff --git a/shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandler.java b/shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandler.java index 49ddc8a058a3..b8d79f4c14bd 100644 --- a/shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandler.java +++ b/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); } } diff --git a/shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/Publish.java b/shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/Publish.java index bd7f48ef9667..57adb9303055 100644 --- a/shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/Publish.java +++ b/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); } }); diff --git a/shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandlerTest.java b/shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/MqttTransportHandlerTest.java new file mode 100644 index 000000000000..d001ab2e4a59 --- /dev/null +++ b/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(); + } +} diff --git a/shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/PublishTest.java b/shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/PublishTest.java index 6d25ff035205..253f537008fd 100644 --- a/shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/PublishTest.java +++ b/shenyu-protocol/shenyu-protocol-mqtt/src/test/java/org/apache/shenyu/protocol/mqtt/PublishTest.java @@ -18,6 +18,7 @@ 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; @@ -25,17 +26,24 @@ 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 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);