[broker-11] finish implement mqtt session service, add correct disconnect flow

Author: JavaSaBr

src/main/java/com/ss/mqtt/broker/config/MqttBrokerConfig.java

@@ -3,8 +3,10 @@
 import com.ss.mqtt.broker.model.MqttPropertyConstants;
 import com.ss.mqtt.broker.model.QoS;
 import com.ss.mqtt.broker.network.MqttConnection;
+import com.ss.mqtt.broker.network.client.MqttClientReleaseHandler;
 import com.ss.mqtt.broker.network.client.UnsafeMqttClient;
 import com.ss.mqtt.broker.network.client.impl.DeviceMqttClient;
+import com.ss.mqtt.broker.network.client.impl.DeviceMqttClientReleaseHandler;
 import com.ss.mqtt.broker.network.packet.PacketType;
 import com.ss.mqtt.broker.network.packet.in.handler.*;
 import com.ss.mqtt.broker.service.*;
@@ -96,25 +98,34 @@ private interface ChannelFactory extends
         handlers[PacketType.SUBSCRIBE.ordinal()] = new SubscribeInPacketHandler(subscriptionService);
         handlers[PacketType.UNSUBSCRIBE.ordinal()] = new UnsubscribeInPacketHandler(subscriptionService);
         handlers[PacketType.PUBLISH.ordinal()] = new PublishInPacketHandler(publishingService);
+        handlers[PacketType.DISCONNECT.ordinal()] = new DisconnetInPacketHandler();
 
         return handlers;
     }
 
+    @Bean
+    @NotNull MqttClientReleaseHandler deviceMqttClientReleaseHandler(
+        @NotNull ClientIdRegistry clientIdRegistry,
+        @NotNull MqttSessionService mqttSessionService
+    ) {
+        return new DeviceMqttClientReleaseHandler(clientIdRegistry, mqttSessionService);
+    }
+
     @Bean
     @NotNull ServerNetwork<@NotNull MqttConnection> deviceNetwork(
         @NotNull ServerNetworkConfig networkConfig,
         @NotNull BufferAllocator bufferAllocator,
         @NotNull MqttConnectionConfig deviceConnectionConfig,
         PacketInHandler @NotNull [] devicePacketHandlers,
-        @NotNull MqttSessionService mqttSessionService
+        @NotNull MqttClientReleaseHandler deviceMqttClientReleaseHandler
     ) {
         return NetworkFactory.newServerNetwork(
             networkConfig,
             deviceConnectionFactory(
                 bufferAllocator,
                 deviceConnectionConfig,
                 devicePacketHandlers,
-                mqttSessionService
+                deviceMqttClientReleaseHandler
             )
         );
     }
@@ -188,7 +199,7 @@ private interface ChannelFactory extends
         @NotNull BufferAllocator bufferAllocator,
         @NotNull MqttConnectionConfig connectionConfig,
         PacketInHandler @NotNull [] packetHandlers,
-        @NotNull MqttSessionService mqttSessionService
+        @NotNull MqttClientReleaseHandler deviceMqttClientReleaseHandler
     ) {
         return (network, channel) -> new MqttConnection(
             network,
@@ -197,8 +208,7 @@ private interface ChannelFactory extends
             100,
             packetHandlers,
             connectionConfig,
-            DeviceMqttClient::new,
-            mqttSessionService
+            mqttConnection -> new DeviceMqttClient(mqttConnection, deviceMqttClientReleaseHandler)
         );
     }
 }

src/main/java/com/ss/mqtt/broker/network/MqttConnection.java

@@ -3,14 +3,12 @@
 import com.ss.mqtt.broker.config.MqttConnectionConfig;
 import com.ss.mqtt.broker.model.MqttSession;
 import com.ss.mqtt.broker.model.MqttVersion;
-import com.ss.mqtt.broker.network.client.MqttClient;
 import com.ss.mqtt.broker.network.client.UnsafeMqttClient;
 import com.ss.mqtt.broker.network.packet.MqttPacketReader;
 import com.ss.mqtt.broker.network.packet.MqttPacketWriter;
 import com.ss.mqtt.broker.network.packet.in.MqttReadablePacket;
 import com.ss.mqtt.broker.network.packet.in.handler.PacketInHandler;
 import com.ss.mqtt.broker.network.packet.out.MqttWritablePacket;
-import com.ss.mqtt.broker.service.MqttSessionService;
 import com.ss.rlib.network.BufferAllocator;
 import com.ss.rlib.network.Connection;
 import com.ss.rlib.network.Network;
@@ -36,11 +34,10 @@ public class MqttConnection extends AbstractConnection<MqttReadablePacket, MqttW
 
     @Getter(AccessLevel.PROTECTED)
     private final @NotNull PacketWriter packetWriter;
-    private final @NotNull MqttSessionService sessionService;
 
     private final @Getter PacketInHandler @NotNull [] packetHandlers;
 
-    private final @Getter @NotNull MqttClient client;
+    private final @Getter @NotNull UnsafeMqttClient client;
     private final @Getter @NotNull MqttConnectionConfig config;
 
     private volatile @Getter @Setter @NotNull MqttVersion mqttVersion;
@@ -53,13 +50,11 @@ public MqttConnection(
         int maxPacketsByRead,
         PacketInHandler @NotNull [] packetHandlers,
         @NotNull MqttConnectionConfig config,
-        @NotNull Function<MqttConnection, UnsafeMqttClient> clientFactory,
-        @NotNull MqttSessionService mqttSessionService
+        @NotNull Function<MqttConnection, UnsafeMqttClient> clientFactory
     ) {
         super(network, channel, NetworkCryptor.NULL, bufferAllocator, maxPacketsByRead);
         this.packetHandlers = packetHandlers;
         this.config = config;
-        this.sessionService = mqttSessionService;
         this.mqttVersion = MqttVersion.MQTT_3_1_1;
         this.packetReader = createPacketReader();
         this.packetWriter = createPacketWriter();
@@ -100,14 +95,7 @@ public boolean isSupported(@NotNull MqttVersion mqttVersion) {
 
     @Override
     protected void doClose() {
-
-        var session = getSession();
-
-        if (session != null) {
-            sessionService.store(client.getClientId(), session)
-                .subscribe(result -> setSession(null));
-        }
-
+        client.release().subscribe();
         super.doClose();
     }
 }

src/main/java/com/ss/mqtt/broker/network/client/MqttClient.java

@@ -1,15 +1,19 @@
 package com.ss.mqtt.broker.network.client;
 
 import com.ss.mqtt.broker.config.MqttConnectionConfig;
+import com.ss.mqtt.broker.model.MqttSession;
 import com.ss.mqtt.broker.network.packet.factory.MqttPacketOutFactory;
 import com.ss.mqtt.broker.network.packet.out.MqttWritablePacket;
 import org.jetbrains.annotations.NotNull;
+import reactor.core.publisher.Mono;
 
 public interface MqttClient {
 
     @NotNull MqttPacketOutFactory getPacketOutFactory();
     @NotNull MqttConnectionConfig getConnectionConfig();
+
     @NotNull String getClientId();
+    @NotNull MqttSession getSession();
 
     int getKeepAlive();
     int getMaximumPacketSize();

src/main/java/com/ss/mqtt/broker/network/client/MqttClientReleaseHandler.java

@@ -0,0 +1,9 @@
+package com.ss.mqtt.broker.network.client;
+
+import org.jetbrains.annotations.NotNull;
+import reactor.core.publisher.Mono;
+
+public interface MqttClientReleaseHandler {
+
+    @NotNull Mono<?> release(@NotNull UnsafeMqttClient client);
+}

src/main/java/com/ss/mqtt/broker/network/client/UnsafeMqttClient.java

@@ -5,6 +5,7 @@
 import com.ss.mqtt.broker.network.MqttConnection;
 import com.ss.mqtt.broker.network.packet.in.MqttReadablePacket;
 import org.jetbrains.annotations.NotNull;
+import reactor.core.publisher.Mono;
 
 public interface UnsafeMqttClient extends MqttClient {
 
@@ -25,4 +26,6 @@ void configure(
     void setSession(@NotNull MqttSession session);
 
     void reject(@NotNull ConnectAckReasonCode reasonCode);
+
+    @NotNull Mono<?> release();
 }

src/main/java/com/ss/mqtt/broker/network/client/impl/AbstractMqttClient.java

@@ -5,6 +5,7 @@
 import com.ss.mqtt.broker.model.MqttPropertyConstants;
 import com.ss.mqtt.broker.model.MqttSession;
 import com.ss.mqtt.broker.network.MqttConnection;
+import com.ss.mqtt.broker.network.client.MqttClientReleaseHandler;
 import com.ss.mqtt.broker.network.client.UnsafeMqttClient;
 import com.ss.mqtt.broker.network.packet.factory.MqttPacketOutFactory;
 import com.ss.mqtt.broker.network.packet.in.MqttReadablePacket;
@@ -13,16 +14,23 @@
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.Setter;
+import lombok.ToString;
 import lombok.extern.log4j.Log4j2;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import reactor.core.publisher.Mono;
+
+import java.util.concurrent.atomic.AtomicBoolean;
 
 @Getter
 @Log4j2
+@ToString(of = "clientId")
 @EqualsAndHashCode(of = "clientId")
 public abstract class AbstractMqttClient implements UnsafeMqttClient {
 
     protected final @NotNull MqttConnection connection;
+    protected final MqttClientReleaseHandler releaseHandler;
+    protected final AtomicBoolean released;
 
     private volatile @Setter @NotNull String clientId;
     private volatile @Setter @Getter @Nullable MqttSession session;
@@ -33,8 +41,10 @@ public abstract class AbstractMqttClient implements UnsafeMqttClient {
     private volatile int topicAliasMaximum = MqttPropertyConstants.TOPIC_ALIAS_MAXIMUM_DEFAULT;
     private volatile int keepAlive = MqttPropertyConstants.SERVER_KEEP_ALIVE_MAX;
 
-    public AbstractMqttClient(@NotNull MqttConnection connection) {
+    public AbstractMqttClient(@NotNull MqttConnection connection, @NotNull MqttClientReleaseHandler releaseHandler) {
         this.connection = connection;
+        this.releaseHandler = releaseHandler;
+        this.released = new AtomicBoolean(false);
         this.clientId = StringUtils.EMPTY;
     }
 
@@ -86,4 +96,13 @@ public void reject(@NotNull ConnectAckReasonCode reasonCode) {
     public @NotNull MqttConnectionConfig getConnectionConfig() {
         return connection.getConfig();
     }
+
+    @Override
+    public @NotNull Mono<?> release() {
+        if (released.compareAndSet(false, true)) {
+            return releaseHandler.release(this);
+        } else {
+            return Mono.empty();
+        }
+    }
 }

src/main/java/com/ss/mqtt/broker/network/client/impl/AbstractMqttClientReleaseHandler.java

@@ -0,0 +1,58 @@
+package com.ss.mqtt.broker.network.client.impl;
+
+import com.ss.mqtt.broker.network.client.MqttClientReleaseHandler;
+import com.ss.mqtt.broker.network.client.UnsafeMqttClient;
+import com.ss.mqtt.broker.service.ClientIdRegistry;
+import com.ss.mqtt.broker.service.MqttSessionService;
+import com.ss.rlib.common.util.StringUtils;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.log4j.Log4j2;
+import org.jetbrains.annotations.NotNull;
+import reactor.core.publisher.Mono;
+
+@Log4j2
+@RequiredArgsConstructor
+public abstract class AbstractMqttClientReleaseHandler<T extends AbstractMqttClient> implements
+    MqttClientReleaseHandler {
+
+    private final @NotNull ClientIdRegistry clientIdRegistry;
+    private final @NotNull MqttSessionService sessionService;
+
+    @Override
+    public @NotNull Mono<?> release(@NotNull UnsafeMqttClient client) {
+        //noinspection unchecked
+        return releaseImpl((T) client);
+    }
+
+    protected @NotNull Mono<?> releaseImpl(@NotNull T client) {
+
+        var clientId = client.getClientId();
+        client.setClientId(StringUtils.EMPTY);
+
+        if (StringUtils.isEmpty(clientId)) {
+            log.warn("This client {} is already released.", client);
+            return Mono.empty();
+        }
+
+        var session = client.getSession();
+
+        Mono<?> asyncActions = null;
+
+        if (session != null) {
+            asyncActions = sessionService.store(clientId, session);
+            client.setSession(null);
+        }
+
+        if (asyncActions != null) {
+            asyncActions = asyncActions.flatMap(any -> clientIdRegistry.unregister(clientId));
+        } else {
+            asyncActions = clientIdRegistry.unregister(clientId);
+        }
+
+        if (asyncActions != null) {
+            return asyncActions;
+        }
+
+        return Mono.empty();
+    }
+}

src/main/java/com/ss/mqtt/broker/network/client/impl/DeviceMqttClient.java

@@ -1,11 +1,12 @@
 package com.ss.mqtt.broker.network.client.impl;
 
 import com.ss.mqtt.broker.network.MqttConnection;
+import com.ss.mqtt.broker.network.client.MqttClientReleaseHandler;
 import org.jetbrains.annotations.NotNull;
 
 public class DeviceMqttClient extends AbstractMqttClient {
 
-    public DeviceMqttClient(@NotNull MqttConnection connection) {
-        super(connection);
+    public DeviceMqttClient(@NotNull MqttConnection connection, @NotNull MqttClientReleaseHandler releaseHandler) {
+        super(connection, releaseHandler);
     }
 }

src/main/java/com/ss/mqtt/broker/network/client/impl/DeviceMqttClientReleaseHandler.java

@@ -0,0 +1,15 @@
+package com.ss.mqtt.broker.network.client.impl;
+
+import com.ss.mqtt.broker.service.ClientIdRegistry;
+import com.ss.mqtt.broker.service.MqttSessionService;
+import org.jetbrains.annotations.NotNull;
+
+public class DeviceMqttClientReleaseHandler extends AbstractMqttClientReleaseHandler<DeviceMqttClient> {
+
+    public DeviceMqttClientReleaseHandler(
+        @NotNull ClientIdRegistry clientIdRegistry,
+        @NotNull MqttSessionService sessionService
+    ) {
+        super(clientIdRegistry, sessionService);
+    }
+}

src/main/java/com/ss/mqtt/broker/network/packet/in/DisconnectInPacket.java

@@ -8,6 +8,7 @@
 import com.ss.mqtt.broker.network.packet.PacketType;
 import com.ss.rlib.common.util.StringUtils;
 import lombok.Getter;
+import lombok.ToString;
 import org.jetbrains.annotations.NotNull;
 
 import java.nio.ByteBuffer;
@@ -18,6 +19,7 @@
  * Disconnect notification.
  */
 @Getter
+@ToString
 public class DisconnectInPacket extends MqttReadablePacket {
 
     public static final byte PACKET_TYPE = (byte) PacketType.DISCONNECT.ordinal();
@@ -86,7 +88,7 @@ protected void readImpl(@NotNull MqttConnection connection, @NotNull ByteBuffer
     protected void readVariableHeader(@NotNull MqttConnection connection, @NotNull ByteBuffer buffer) {
 
         // https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901207
-        if (connection.isSupported(MqttVersion.MQTT_5)) {
+        if (connection.isSupported(MqttVersion.MQTT_5) && buffer.hasRemaining()) {
             reasonCode = DisconnectReasonCode.of(readUnsignedByte(buffer));
         }
     }