[broker-15] base implementation of publishing with QoS 1

JavaSaBr · JavaSaBr · commit 8116d5ca6091 · 2019-11-19T23:21:55.000+03:00
diff --git a/build.gradle b/build.gradle
@@ -32,7 +32,7 @@ allprojects {
 
     ext {
         annotationVersion = "17.0.0"
-        rlibVersion = "9.5.0"
+        rlibVersion = "9.6.0"
         lombokVersion = '1.18.4'
         springbootVersion = '2.2.0.RELEASE'
         springVersion = '5.1.6.RELEASE'
diff --git a/src/main/java/com/ss/mqtt/broker/config/MqttBrokerConfig.java b/src/main/java/com/ss/mqtt/broker/config/MqttBrokerConfig.java
@@ -88,11 +88,6 @@ private interface ChannelFactory extends
         );
     }
 
-    @Bean
-    @NotNull PacketIdGenerator packetIdGenerator() {
-        return new DefaultPacketIdGenerator();
-    }
-
     @Bean
     PacketInHandler @NotNull [] devicePacketHandlers(
         @NotNull AuthenticationService authenticationService,
@@ -164,10 +159,10 @@ private interface ChannelFactory extends
     }
 
     @Bean
-    @NotNull PublishOutHandler[] publishOutHandlers(@NotNull PacketIdGenerator packetIdGenerator) {
+    @NotNull PublishOutHandler[] publishOutHandlers() {
         return new PublishOutHandler[] {
             new Qos0PublishOutHandler(),
-            new Qos1PublishOutHandler(packetIdGenerator),
+            new Qos1PublishOutHandler(),
             new Qos2PublishOutHandler(),
         };
     }
diff --git a/src/main/java/com/ss/mqtt/broker/handler/publish/out/Qos1PublishOutHandler.java b/src/main/java/com/ss/mqtt/broker/handler/publish/out/Qos1PublishOutHandler.java
@@ -4,23 +4,22 @@
 import com.ss.mqtt.broker.model.QoS;
 import com.ss.mqtt.broker.model.Subscriber;
 import com.ss.mqtt.broker.network.packet.in.PublishInPacket;
-import com.ss.mqtt.broker.service.PacketIdGenerator;
 import lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
 
 @RequiredArgsConstructor
 public class Qos1PublishOutHandler extends AbstractPublishOutHandler {
 
-    private final @NotNull PacketIdGenerator packetIdGenerator;
-
     @Override
     public void handle(@NotNull PublishInPacket packet, @NotNull Subscriber subscriber) {
 
         var client = subscriber.getMqttClient();
+        var session = client.getSession();
         var packetOutFactory = client.getPacketOutFactory();
+
         var publish = packetOutFactory.newPublish(
             client,
-            packetIdGenerator.nextPacketId(),
+            session.nextPacketId(),
             QoS.AT_LEAST_ONCE_DELIVERY,
             packet.isRetained(),
             packet.isDuplicate(),
@@ -33,7 +32,6 @@ public void handle(@NotNull PublishInPacket packet, @NotNull Subscriber subscrib
             packet.getUserProperties()
         );
 
-        var session = client.getSession();
         session.registerPendingPublish(publish);
 
         client.send(publish);
diff --git a/src/main/java/com/ss/mqtt/broker/model/MqttPropertyConstants.java b/src/main/java/com/ss/mqtt/broker/model/MqttPropertyConstants.java
@@ -5,6 +5,7 @@ public interface MqttPropertyConstants {
     QoS MAXIMUM_QOS_DEFAULT = QoS.EXACTLY_ONCE_DELIVERY;
 
     int MAXIMUM_PROTOCOL_PACKET_SIZE = 256 * 1024 * 1024;
+    int MAXIMUM_PACKET_ID = 0xFFFF;
 
     long SESSION_EXPIRY_INTERVAL_DISABLED = 0;
     long SESSION_EXPIRY_INTERVAL_DEFAULT = 120;
diff --git a/src/main/java/com/ss/mqtt/broker/model/MqttSession.java b/src/main/java/com/ss/mqtt/broker/model/MqttSession.java
@@ -27,6 +27,8 @@ interface PendingCallback<T extends HasPacketId> {
 
     @NotNull String getClientId();
 
+    int nextPacketId();
+
     /**
      * @return the expiration time in ms or -1 if it should not be expired now.
      */
diff --git a/src/main/java/com/ss/mqtt/broker/model/impl/DefaultMqttSession.java b/src/main/java/com/ss/mqtt/broker/model/impl/DefaultMqttSession.java
@@ -1,5 +1,6 @@
 package com.ss.mqtt.broker.model.impl;
 
+import com.ss.mqtt.broker.model.MqttPropertyConstants;
 import com.ss.mqtt.broker.model.MqttSession.UnsafeMqttSession;
 import com.ss.mqtt.broker.network.client.MqttClient;
 import com.ss.mqtt.broker.network.packet.HasPacketId;
@@ -16,6 +17,7 @@
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Collection;
+import java.util.concurrent.atomic.AtomicInteger;
 
 @Log4j2
 @ToString(of = "clientId")
@@ -30,12 +32,27 @@ private static class PendingPublish<T extends HasPacketId> {
 
     private final @NotNull String clientId;
     private final @NotNull ConcurrentArray<PendingPublish<?>> pendingPublishes;
+    private final @NotNull AtomicInteger packetIdGenerator;
 
     private volatile @Getter @Setter long expirationTime = -1;
 
     public DefaultMqttSession(@NotNull String clientId) {
         this.clientId = clientId;
         this.pendingPublishes = ConcurrentArray.ofType(PendingPublish.class);
+        this.packetIdGenerator = new AtomicInteger(0);
+    }
+
+    @Override
+    public int nextPacketId() {
+
+        var nextId = packetIdGenerator.incrementAndGet();
+
+        if (nextId >= MqttPropertyConstants.MAXIMUM_PACKET_ID) {
+            packetIdGenerator.compareAndSet(nextId, 0);
+            return nextPacketId();
+        }
+
+        return nextId;
     }
 
     @Override
@@ -54,13 +71,10 @@ public <T extends MqttReadablePacket & HasPacketId> void registerPendingPublish(
         @NotNull PublishOutPacket publish,
         @NotNull PendingCallback<T> callback
     ) {
-        // FIXME add new method to array
-        var stamp = pendingPublishes.writeLock();
-        try {
-            pendingPublishes.add(new PendingPublish<>(publish, callback, System.currentTimeMillis()));
-        } finally {
-            pendingPublishes.writeUnlock(stamp);
-        }
+        pendingPublishes.runInWriteLock(
+            new PendingPublish<>(publish, callback, System.currentTimeMillis()),
+            Array::add
+        );
     }
 
     @Override
@@ -70,11 +84,10 @@ public <T extends MqttReadablePacket & HasPacketId> void unregisterPendingPacket
     ) {
 
         var packetId = feedback.getPacketId();
-
-        // FIXME add new method to array
-        var pendingPublish = pendingPublishes.findAnyInReadLock(
+        var pendingPublish = pendingPublishes.findAnyConvertedToIntInReadLock(
             packetId,
-            (id, pending) -> id == pending.publish.getPacketId()
+            pending -> pending.publish.getPacketId(),
+            (id, targetId) -> id == targetId
         );
 
         if (pendingPublish == null) {
diff --git a/src/main/java/com/ss/mqtt/broker/network/packet/in/PublishAckInPacket.java b/src/main/java/com/ss/mqtt/broker/network/packet/in/PublishAckInPacket.java
@@ -68,11 +68,16 @@ protected void readVariableHeader(@NotNull MqttConnection connection, @NotNull B
         packetId = readUnsignedShort(buffer);
 
         // https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901123
-        if (connection.isSupported(MqttVersion.MQTT_5)) {
+        if (connection.isSupported(MqttVersion.MQTT_5) && buffer.hasRemaining()) {
             reasonCode = PublishAckReasonCode.of(readUnsignedByte(buffer));
         }
     }
 
+    @Override
+    protected boolean isPropertiesSupported(@NotNull MqttConnection connection, @NotNull ByteBuffer buffer) {
+        return connection.isSupported(MqttVersion.MQTT_5) && buffer.hasRemaining();
+    }
+
     @Override
     protected @NotNull Set<PacketProperty> getAvailableProperties() {
         return AVAILABLE_PROPERTIES;
diff --git a/src/main/java/com/ss/mqtt/broker/service/PacketIdGenerator.java b/src/main/java/com/ss/mqtt/broker/service/PacketIdGenerator.java
diff --git a/src/main/java/com/ss/mqtt/broker/service/impl/DefaultPacketIdGenerator.java b/src/main/java/com/ss/mqtt/broker/service/impl/DefaultPacketIdGenerator.java
diff --git a/src/main/java/com/ss/mqtt/broker/service/impl/InMemoryMqttSessionService.java b/src/main/java/com/ss/mqtt/broker/service/impl/InMemoryMqttSessionService.java
@@ -1,6 +1,7 @@
 package com.ss.mqtt.broker.service.impl;
 
 import com.ss.mqtt.broker.model.MqttSession;
+import com.ss.mqtt.broker.model.MqttSession.UnsafeMqttSession;
 import com.ss.mqtt.broker.model.impl.DefaultMqttSession;
 import com.ss.mqtt.broker.service.MqttSessionService;
 import com.ss.rlib.common.concurrent.util.ThreadUtils;
@@ -18,7 +19,7 @@
 @Log4j2
 public class InMemoryMqttSessionService implements MqttSessionService, Closeable {
 
-    private final @NotNull ConcurrentObjectDictionary<String, MqttSession> storedSession;
+    private final @NotNull ConcurrentObjectDictionary<String, UnsafeMqttSession> storedSession;
     private final @NotNull Thread cleanThread;
 
     private final int cleanInterval;
@@ -65,10 +66,10 @@ public InMemoryMqttSessionService(int cleanInterval) {
     @Override
     public @NotNull Mono<Boolean> store(@NotNull String clientId, @NotNull MqttSession session, long expiryInterval) {
 
-        var unsafe = (MqttSession.UnsafeMqttSession) session;
+        var unsafe = (UnsafeMqttSession) session;
         unsafe.setExpirationTime(System.currentTimeMillis() + (expiryInterval * 1000));
 
-        storedSession.runInWriteLock(clientId, session, ObjectDictionary::put);
+        storedSession.runInWriteLock(clientId, unsafe, ObjectDictionary::put);
 
         log.debug("Stored session for client {}", clientId);
 
@@ -77,8 +78,8 @@ public InMemoryMqttSessionService(int cleanInterval) {
 
     private void cleanup() {
 
-        var toCheck = ArrayFactory.newArray(MqttSession.class);
-        var toRemove = ArrayFactory.newArray(MqttSession.class);
+        var toCheck = ArrayFactory.newArray(UnsafeMqttSession.class);
+        var toRemove = ArrayFactory.newArray(UnsafeMqttSession.class);
 
         while (!closed) {
 
@@ -110,13 +111,18 @@ private void cleanup() {
                     // if we already have new session under the same client id
                     if (removed != null && removed != session) {
                         dictionary.put(session.getClientId(), removed);
+                    } else if (removed != null) {
+                        removed.clear();
                     }
                 }
             });
         }
     }
 
-    private boolean findToRemove(@NotNull Array<MqttSession> toCheck, @NotNull Array<MqttSession> toRemove) {
+    private boolean findToRemove(
+        @NotNull Array<UnsafeMqttSession> toCheck,
+        @NotNull Array<UnsafeMqttSession> toRemove
+    ) {
 
         var currentTime = System.currentTimeMillis();
 
diff --git a/src/test/groovy/com/ss/mqtt/broker/test/integration/ConnectSubscribePublishTest.groovy b/src/test/groovy/com/ss/mqtt/broker/test/integration/ConnectSubscribePublishTest.groovy
@@ -6,11 +6,13 @@ import com.hivemq.client.mqtt.mqtt5.message.publish.Mqtt5PayloadFormatIndicator
 import com.hivemq.client.mqtt.mqtt5.message.publish.Mqtt5Publish
 import com.hivemq.client.mqtt.mqtt5.message.subscribe.suback.Mqtt5SubAckReasonCode
 
+import java.util.concurrent.atomic.AtomicReference
+
 class ConnectSubscribePublishTest extends MqttBrokerTest {
     
-    def "publisher should publish message to broker"() {
+    def "publisher should publish message QoS 0"() {
         given:
-            Mqtt5Publish receivedMessage = null
+            def received = new AtomicReference<Mqtt5Publish>()
             def subscriber = buildClient()
             def publisher = buildClient()
         when:
@@ -20,7 +22,7 @@ class ConnectSubscribePublishTest extends MqttBrokerTest {
             def subscribeResult = subscriber.subscribeWith()
                 .topicFilter(topicFilter)
                 .qos(MqttQos.AT_MOST_ONCE)
-                .callback({ publish -> receivedMessage = publish })
+                .callback({ publish -> received.set(publish) })
                 .send()
                 .join()
         
@@ -45,10 +47,57 @@ class ConnectSubscribePublishTest extends MqttBrokerTest {
             publishResult.publish.type == Mqtt5MessageType.PUBLISH
             publishResult.publish.topic.levels.join("/") == topicFilter
         
-            receivedMessage != null
-            receivedMessage.qos == MqttQos.AT_MOST_ONCE
-            receivedMessage.type == Mqtt5MessageType.PUBLISH
-            receivedMessage.topic.levels.join("/") == topicFilter
+            received.get() != null
+            received.get().qos == MqttQos.AT_MOST_ONCE
+            received.get().type == Mqtt5MessageType.PUBLISH
+            received.get().topic.levels.join("/") == topicFilter
+        cleanup:
+            subscriber.disconnect()
+            publisher.disconnect()
+    }
+    
+    def "publisher should publish message QoS 1"() {
+        given:
+            def received = new AtomicReference<Mqtt5Publish>()
+            def subscriber = buildClient()
+            def publisher = buildClient()
+        when:
+           
+            subscriber.connect().join()
+            publisher.connect().join()
+            
+            def subscribeResult = subscriber.subscribeWith()
+                .topicFilter(topicFilter)
+                .qos(MqttQos.AT_LEAST_ONCE)
+                .callback({ publish -> received.set(publish) })
+                .send()
+                .join()
+            
+            def publishResult = publisher.publishWith()
+                .topic(topicFilter)
+                .qos(MqttQos.AT_LEAST_ONCE)
+                .payload(publishPayload)
+                .payloadFormatIndicator(Mqtt5PayloadFormatIndicator.UTF_8)
+                .send()
+                .join()
+            
+            Thread.sleep(500)
+        then:
+            noExceptionThrown()
+            
+            subscribeResult != null
+            subscribeResult.reasonCodes.contains(Mqtt5SubAckReasonCode.GRANTED_QOS_1)
+            subscribeResult.type == Mqtt5MessageType.SUBACK
+            
+            publishResult != null
+            publishResult.publish.qos == MqttQos.AT_LEAST_ONCE
+            publishResult.publish.type == Mqtt5MessageType.PUBLISH
+            publishResult.publish.topic.levels.join("/") == topicFilter
+    
+            received.get() != null
+            received.get().qos == MqttQos.AT_LEAST_ONCE
+            received.get().type == Mqtt5MessageType.PUBLISH
+            received.get().topic.levels.join("/") == topicFilter
         cleanup:
             subscriber.disconnect()
             publisher.disconnect()
