fix: deliver mqtt messages at granted qos with per-subscriber packet id

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

@@ -19,6 +19,7 @@
 
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.mqtt.MqttConnectMessage;
+import io.netty.handler.codec.mqtt.MqttMessage;
 import io.netty.handler.codec.mqtt.MqttPublishMessage;
 import io.netty.handler.codec.mqtt.MqttSubscribeMessage;
 import io.netty.handler.codec.mqtt.MqttUnsubscribeMessage;
@@ -88,4 +89,13 @@ default void disconnect(final ChannelHandlerContext ctx) {
 
     }
 
+    /**
+     * Publish Release, third message of the QoS 2 protocol flow.
+     * @param ctx ctx
+     * @param msg msg
+     */
+    default void pubRel(final ChannelHandlerContext ctx, final MqttMessage msg) {
+
+    }
+
 }

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

@@ -21,6 +21,7 @@
 import io.netty.channel.ChannelHandlerContext;
 import org.apache.shenyu.common.utils.Singleton;
 import org.apache.shenyu.protocol.mqtt.repositories.ChannelRepository;
+import org.apache.shenyu.protocol.mqtt.utils.MqttPacketIdGenerator;
 
 /**
  * The DISCONNECT message is sent from the client to the server to indicate
@@ -45,5 +46,6 @@ public void disconnect(final ChannelHandlerContext ctx) {
     private void cleanChannel(final Channel channel) {
         //// todo ttl
         Singleton.INST.get(ChannelRepository.class).remove(channel);
+        MqttPacketIdGenerator.remove(channel);
     }
 }

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

@@ -20,6 +20,7 @@
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.mqtt.MqttConnectMessage;
+import io.netty.handler.codec.mqtt.MqttMessage;
 import io.netty.handler.codec.mqtt.MqttPublishMessage;
 import io.netty.handler.codec.mqtt.MqttSubscribeMessage;
 import io.netty.handler.codec.mqtt.MqttUnsubscribeMessage;
@@ -91,4 +92,10 @@ public void disconnect(final ChannelHandlerContext ctx) {
         //// todo polymorphism disconnect
         new Disconnect().disconnect(ctx);
     }
+
+    @Override
+    public void pubRel(final ChannelHandlerContext ctx, final MqttMessage msg) {
+        //// todo polymorphism pubRel
+        new PubRel().pubRel(ctx, msg);
+    }
 }

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

@@ -65,6 +65,9 @@ public void connect() {
             case PINGREQ:
                 messageType.pingReq(ctx);
                 break;
+            case PUBREL:
+                messageType.pubRel(ctx, msg);
+                break;
             case PUBACK:
             case DISCONNECT:
             default:

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

@@ -17,13 +17,17 @@
 
 package org.apache.shenyu.protocol.mqtt;
 
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.mqtt.MqttMessage;
 import io.netty.util.concurrent.Future;
 import io.netty.util.concurrent.GenericFutureListener;
 import org.apache.shenyu.common.utils.Singleton;
 import org.apache.shenyu.protocol.mqtt.repositories.ChannelRepository;
+import org.apache.shenyu.protocol.mqtt.repositories.SubscribeRepository;
+import org.apache.shenyu.protocol.mqtt.utils.MqttPacketIdGenerator;
 
 /**
  * mqtt transport handler.
@@ -48,7 +52,10 @@ public void channelInactive(final ChannelHandlerContext ctx) throws Exception {
 
     @Override
     public void operationComplete(final Future<? super Void> future) throws Exception {
-
+        Channel channel = ((ChannelFuture) future).channel();
+        Singleton.INST.get(ChannelRepository.class).remove(channel);
+        Singleton.INST.get(SubscribeRepository.class).remove(channel);
+        MqttPacketIdGenerator.remove(channel);
     }
 
 }

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

@@ -0,0 +1,41 @@
+/*
+ * 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.channel.ChannelHandlerContext;
+import io.netty.handler.codec.mqtt.MqttFixedHeader;
+import io.netty.handler.codec.mqtt.MqttMessage;
+import io.netty.handler.codec.mqtt.MqttMessageIdVariableHeader;
+import io.netty.handler.codec.mqtt.MqttQoS;
+
+import static io.netty.handler.codec.mqtt.MqttMessageType.PUBCOMP;
+
+/**
+ * The PUBREL message is the third message of the QoS 2 protocol flow,
+ * the server responds with PUBCOMP to release the packet id.
+ */
+public class PubRel extends MessageType {
+
+    @Override
+    public void pubRel(final ChannelHandlerContext ctx, final MqttMessage msg) {
+        MqttMessageIdVariableHeader variableHeader = (MqttMessageIdVariableHeader) msg.variableHeader();
+        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(PUBCOMP, false, MqttQoS.AT_MOST_ONCE, false, 0);
+        MqttMessage mqttPubCompMessage = new MqttMessage(mqttFixedHeader, variableHeader);
+        ctx.writeAndFlush(mqttPubCompMessage);
+    }
+}

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

@@ -22,6 +22,7 @@
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.mqtt.MqttFixedHeader;
+import io.netty.handler.codec.mqtt.MqttMessage;
 import io.netty.handler.codec.mqtt.MqttMessageIdVariableHeader;
 import io.netty.handler.codec.mqtt.MqttPublishMessage;
 import io.netty.handler.codec.mqtt.MqttQoS;
@@ -31,12 +32,14 @@
 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.apache.shenyu.protocol.mqtt.utils.MqttPacketIdGenerator;
 
-import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
 import static io.netty.channel.ChannelFutureListener.FIRE_EXCEPTION_ON_FAILURE;
 import static io.netty.handler.codec.mqtt.MqttMessageType.PUBACK;
+import static io.netty.handler.codec.mqtt.MqttMessageType.PUBREC;
 
 /**
  * Publish message.
@@ -52,7 +55,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()) {
@@ -62,7 +64,7 @@ public void publish(final ChannelHandlerContext ctx, final MqttPublishMessage ms
             }
         }
         int packetId = msg.variableHeader().packetId();
-        CompletableFuture.runAsync(() -> send(topic, payload, packetId));
+        CompletableFuture.runAsync(() -> send(topic, payload, mqttQoS));
 
         switch (mqttQoS.value()) {
             case 0:
@@ -82,32 +84,25 @@ public void publish(final ChannelHandlerContext ctx, final MqttPublishMessage ms
     }
 
     /**
-     * todo qos0.
-     */
-    private void qos0() {
-
-    }
-
-    /**
-     * todo qos1.
+     * send PUBACK to the publisher for a qos1 publish.
      */
     private void qos1(final ChannelHandlerContext ctx, final int packetId) {
-        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(PUBACK, false, MqttQoS.AT_LEAST_ONCE, false, 0);
+        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(PUBACK, false, MqttQoS.AT_MOST_ONCE, false, 0);
         MqttMessageIdVariableHeader mqttMsgIdVariableHeader = MqttMessageIdVariableHeader.from(packetId);
 
         MqttPubAckMessage mqttPubAckMessage = new MqttPubAckMessage(mqttFixedHeader, mqttMsgIdVariableHeader);
         ctx.writeAndFlush(mqttPubAckMessage);
     }
 
     /**
-     * todo qos2.
+     * send PUBREC to the publisher for a qos2 publish.
      */
     private void qos2(final ChannelHandlerContext ctx, final int packetId) {
-        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(PUBACK, false, MqttQoS.EXACTLY_ONCE, false, 0);
+        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(PUBREC, false, MqttQoS.AT_MOST_ONCE, false, 0);
         MqttMessageIdVariableHeader mqttMsgIdVariableHeader = MqttMessageIdVariableHeader.from(packetId);
 
-        MqttPubAckMessage mqttPubAckMessage = new MqttPubAckMessage(mqttFixedHeader, mqttMsgIdVariableHeader);
-        ctx.writeAndFlush(mqttPubAckMessage);
+        MqttMessage mqttPubRecMessage = new MqttMessage(mqttFixedHeader, mqttMsgIdVariableHeader);
+        ctx.writeAndFlush(mqttPubRecMessage);
     }
 
     private String byteBufToString(final ByteBuf byteBuf) {
@@ -120,16 +115,23 @@ private String byteBufToString(final ByteBuf byteBuf) {
         }
     }
 
-    private void send(final String topic, final ByteBuf payload, final int packetId) {
-        List<Channel> channels = Singleton.INST.get(SubscribeRepository.class).get(topic);
+    private void send(final String topic, final ByteBuf payload, final MqttQoS publishQoS) {
+        Map<Channel, MqttQoS> subscribers = Singleton.INST.get(SubscribeRepository.class).get(topic);
         //// todo thread pool
-        channels.parallelStream().forEach(channel -> {
+        subscribers.entrySet().parallelStream().forEach(entry -> {
+            Channel channel = entry.getKey();
             if (channel.isActive()) {
-                MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, MqttQoS.AT_MOST_ONCE, false, 0);
+                MqttQoS qos = minQoS(publishQoS, entry.getValue());
+                int packetId = MqttQoS.AT_MOST_ONCE == qos ? 0 : MqttPacketIdGenerator.next(channel);
+                MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, qos, false, 0);
                 MqttPublishVariableHeader mqttPublishVariableHeader = new MqttPublishVariableHeader(topic, packetId);
-                MqttPublishMessage mqttPublishMessage = new MqttPublishMessage(mqttFixedHeader, mqttPublishVariableHeader, Unpooled.wrappedBuffer(payload));
+                MqttPublishMessage mqttPublishMessage = new MqttPublishMessage(mqttFixedHeader, mqttPublishVariableHeader, Unpooled.wrappedBuffer(payload.retain()));
                 channel.writeAndFlush(mqttPublishMessage);
             }
         });
     }
+
+    private static MqttQoS minQoS(final MqttQoS publishQoS, final MqttQoS grantedQoS) {
+        return publishQoS.value() <= grantedQoS.value() ? publishQoS : grantedQoS;
+    }
 }

shenyu-protocol/shenyu-protocol-mqtt/src/main/java/org/apache/shenyu/protocol/mqtt/repositories/SubscribeRepository.java

@@ -18,35 +18,33 @@
 package org.apache.shenyu.protocol.mqtt.repositories;
 
 import io.netty.channel.Channel;
+import io.netty.handler.codec.mqtt.MqttQoS;
 import io.netty.handler.codec.mqtt.MqttTopicSubscription;
-import org.apache.commons.collections4.CollectionUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.concurrent.CopyOnWriteArraySet;
 
 /**
  * Topic and channel association.
  */
-public class SubscribeRepository implements BaseRepository<List<String>, List<Channel>> {
+public class SubscribeRepository implements BaseRepository<List<String>, Map<Channel, MqttQoS>> {
 
     private static final Logger LOG = LoggerFactory.getLogger(SubscribeRepository.class);
 
-    private static final Map<String, List<Channel>> TOPIC_CHANNEL_FACTORY = new ConcurrentHashMap<>();
+    private static final Map<String, Map<Channel, MqttQoS>> TOPIC_CHANNEL_FACTORY = new ConcurrentHashMap<>();
 
     @Override
-    public void add(final List<String> topics, final List<Channel> channels) {
-        CompletableFuture.runAsync(() -> topics.parallelStream().forEach(s -> {
-            List<Channel> list = get(s);
-            list.addAll(channels);
-            TOPIC_CHANNEL_FACTORY.put(s, list);
-        }));
+    public void add(final List<String> topics, final Map<Channel, MqttQoS> channelQos) {
+        CompletableFuture.runAsync(() -> topics.parallelStream().forEach(topic ->
+                channelQos.forEach((channel, qos) -> TOPIC_CHANNEL_FACTORY
+                        .computeIfAbsent(topic, key -> new ConcurrentHashMap<>())
+                        .merge(channel, qos, SubscribeRepository::maxQoS))));
     }
 
     /**
@@ -55,11 +53,11 @@ public void add(final List<String> topics, final List<Channel> channels) {
      * @param mqttTopicSubscription mqtt subscription info
      */
     public void add(final Channel channel, final List<MqttTopicSubscription> mqttTopicSubscription) {
-        CompletableFuture.runAsync(() -> mqttTopicSubscription.parallelStream().forEach(s -> {
-            List<Channel> channels = get(s.topicName());
-            channels.add(channel);
-            TOPIC_CHANNEL_FACTORY.put(s.topicName(), channels);
-        }));
+        CompletableFuture.runAsync(() -> mqttTopicSubscription.parallelStream()
+                .filter(s -> s.qualityOfService() != MqttQoS.FAILURE)
+                .forEach(s -> TOPIC_CHANNEL_FACTORY
+                        .computeIfAbsent(s.topicName(), key -> new ConcurrentHashMap<>())
+                        .merge(channel, s.qualityOfService(), SubscribeRepository::maxQoS)));
     }
 
     @Override
@@ -74,27 +72,41 @@ public void remove(final List<String> topics) {
      */
     public void remove(final List<String> topics, final Channel channel) {
         CompletableFuture.runAsync(() -> topics.parallelStream().forEach(topic -> {
-            List<Channel> channels = TOPIC_CHANNEL_FACTORY.get(topic);
-            if (CollectionUtils.isNotEmpty(channels)) {
-                channels.remove(channel);
+            Map<Channel, MqttQoS> subscribers = TOPIC_CHANNEL_FACTORY.get(topic);
+            if (Objects.nonNull(subscribers)) {
+                subscribers.remove(channel);
             }
         }));
     }
 
+    /**
+     * remove the channel from all topics it subscribed.
+     * @param channel channel
+     */
+    public void remove(final Channel channel) {
+        CompletableFuture.runAsync(() -> TOPIC_CHANNEL_FACTORY.values().parallelStream()
+                .forEach(subscribers -> subscribers.remove(channel)));
+    }
+
     @Override
-    public List<Channel> get(final List<String> topics) {
-        Set<Channel> channels = new CopyOnWriteArraySet<>();
-        topics.parallelStream().forEach(s -> channels.addAll(TOPIC_CHANNEL_FACTORY.get(s)));
-        return new CopyOnWriteArrayList<>(channels);
+    public Map<Channel, MqttQoS> get(final List<String> topics) {
+        Map<Channel, MqttQoS> subscribers = new ConcurrentHashMap<>();
+        topics.parallelStream().forEach(topic -> TOPIC_CHANNEL_FACTORY.getOrDefault(topic, Collections.emptyMap())
+                .forEach((channel, qos) -> subscribers.merge(channel, qos, SubscribeRepository::maxQoS)));
+        return subscribers;
     }
 
     /**
-     * get Channels.
+     * get subscriber channels with their granted qos.
      * @param topic topic
-     * @return Channels
+     * @return map of channel to granted qos
      */
-    public List<Channel> get(final String topic) {
-        return TOPIC_CHANNEL_FACTORY.getOrDefault(topic, new CopyOnWriteArrayList<>());
+    public Map<Channel, MqttQoS> get(final String topic) {
+        return TOPIC_CHANNEL_FACTORY.getOrDefault(topic, Collections.emptyMap());
+    }
+
+    private static MqttQoS maxQoS(final MqttQoS qos1, final MqttQoS qos2) {
+        return qos1.value() >= qos2.value() ? qos1 : qos2;
     }
 
 }