[broker-11] implement in-memory mqtt session service

Author: JavaSaBr

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

@@ -9,7 +9,10 @@
 import com.ss.mqtt.broker.network.packet.in.handler.*;
 import com.ss.mqtt.broker.service.*;
 import com.ss.mqtt.broker.service.impl.*;
-import com.ss.rlib.network.*;
+import com.ss.rlib.network.BufferAllocator;
+import com.ss.rlib.network.Network;
+import com.ss.rlib.network.NetworkFactory;
+import com.ss.rlib.network.ServerNetworkConfig;
 import com.ss.rlib.network.impl.DefaultBufferAllocator;
 import com.ss.rlib.network.server.ServerNetwork;
 import lombok.RequiredArgsConstructor;
@@ -44,11 +47,6 @@ private interface ChannelFactory extends
         return new DefaultBufferAllocator(networkConfig);
     }
 
-    @Bean
-    @NotNull ClientService clientService() {
-        return new DefaultClientService();
-    }
-
     @Bean
     @NotNull ClientIdRegistry clientIdRegistry() {
         return new InMemoryClientIdRegistry(
@@ -85,11 +83,16 @@ private interface ChannelFactory extends
         @NotNull AuthenticationService authenticationService,
         @NotNull ClientIdRegistry clientIdRegistry,
         @NotNull SubscriptionService subscriptionService,
-        @NotNull PublishingService publishingService
+        @NotNull PublishingService publishingService,
+        @NotNull MqttSessionService mqttSessionService
     ) {
 
         var handlers = new PacketInHandler[PacketType.INVALID.ordinal()];
-        handlers[PacketType.CONNECT.ordinal()] = new ConnectInPacketHandler(clientIdRegistry, authenticationService);
+        handlers[PacketType.CONNECT.ordinal()] = new ConnectInPacketHandler(
+            clientIdRegistry,
+            authenticationService,
+            mqttSessionService
+        );
         handlers[PacketType.SUBSCRIBE.ordinal()] = new SubscribeInPacketHandler(subscriptionService);
         handlers[PacketType.UNSUBSCRIBE.ordinal()] = new UnsubscribeInPacketHandler(subscriptionService);
         handlers[PacketType.PUBLISH.ordinal()] = new PublishInPacketHandler(publishingService);

src/main/java/com/ss/mqtt/broker/model/MqttSession.java

@@ -4,6 +4,11 @@
 
 public interface MqttSession {
 
+    interface UnsafeMqttSession extends MqttSession {
+
+        void setExpirationTime(long expirationTime);
+    }
+
     @NotNull String getClientId();
 
     /**

src/main/java/com/ss/mqtt/broker/model/impl/DefaultMqttSession.java

@@ -1,5 +1,20 @@
 package com.ss.mqtt.broker.model.impl;
 
-import com.ss.mqtt.broker.model.MqttSession;
+import com.ss.mqtt.broker.model.MqttSession.UnsafeMqttSession;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.Setter;
+import org.jetbrains.annotations.NotNull;
 
-public class DefaultMqttSession implements MqttSession {}
+@RequiredArgsConstructor
+public class DefaultMqttSession implements UnsafeMqttSession {
+
+    private final @NotNull String clientId;
+
+    private volatile @Getter @Setter long expirationTime = -1;
+
+    @Override
+    public @NotNull String getClientId() {
+        return clientId;
+    }
+}

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

@@ -104,7 +104,7 @@ protected void doClose() {
         var session = getSession();
 
         if (session != null) {
-            sessionService.store(getClient().getClientId(), session)
+            sessionService.store(client.getClientId(), session)
                 .subscribe(result -> setSession(null));
         }
 

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

@@ -1,6 +1,7 @@
 package com.ss.mqtt.broker.network.client;
 
 import com.ss.mqtt.broker.model.ConnectAckReasonCode;
+import com.ss.mqtt.broker.model.MqttSession;
 import com.ss.mqtt.broker.network.MqttConnection;
 import com.ss.mqtt.broker.network.packet.in.MqttReadablePacket;
 import org.jetbrains.annotations.NotNull;
@@ -21,5 +22,7 @@ void configure(
 
     void setClientId(@NotNull String clientId);
 
+    void setSession(@NotNull MqttSession session);
+
     void reject(@NotNull ConnectAckReasonCode reasonCode);
 }

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

@@ -3,6 +3,7 @@
 import com.ss.mqtt.broker.config.MqttConnectionConfig;
 import com.ss.mqtt.broker.model.ConnectAckReasonCode;
 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.UnsafeMqttClient;
 import com.ss.mqtt.broker.network.packet.factory.MqttPacketOutFactory;
@@ -14,6 +15,7 @@
 import lombok.Setter;
 import lombok.extern.log4j.Log4j2;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 @Getter
 @Log4j2
@@ -23,6 +25,7 @@ public abstract class AbstractMqttClient implements UnsafeMqttClient {
     protected final @NotNull MqttConnection connection;
 
     private volatile @Setter @NotNull String clientId;
+    private volatile @Setter @Getter @Nullable MqttSession session;
 
     private volatile long sessionExpiryInterval = MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DEFAULT;
     private volatile int receiveMax = MqttPropertyConstants.RECEIVE_MAXIMUM_DEFAULT;

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

@@ -1,14 +1,17 @@
 package com.ss.mqtt.broker.network.packet.in.handler;
 
-import static reactor.core.publisher.Mono.fromCallable;
-import static reactor.core.publisher.Mono.fromRunnable;
+import static com.ss.mqtt.broker.model.ConnectAckReasonCode.BAD_USER_NAME_OR_PASSWORD;
+import static com.ss.mqtt.broker.model.ConnectAckReasonCode.CLIENT_IDENTIFIER_NOT_VALID;
+import static com.ss.mqtt.broker.util.ReactorUtils.ifTrue;
 import com.ss.mqtt.broker.exception.ConnectionRejectException;
 import com.ss.mqtt.broker.exception.MalformedPacketMqttException;
 import com.ss.mqtt.broker.model.ConnectAckReasonCode;
+import com.ss.mqtt.broker.model.MqttSession;
 import com.ss.mqtt.broker.network.client.UnsafeMqttClient;
 import com.ss.mqtt.broker.network.packet.in.ConnectInPacket;
 import com.ss.mqtt.broker.service.AuthenticationService;
 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 org.jetbrains.annotations.NotNull;
@@ -19,6 +22,7 @@ public class ConnectInPacketHandler extends AbstractPacketHandler<UnsafeMqttClie
 
     private final ClientIdRegistry clientIdRegistry;
     private final AuthenticationService authenticationService;
+    private final MqttSessionService mqttSessionService;
 
     @Override
     protected void handleImpl(@NotNull UnsafeMqttClient client, @NotNull ConnectInPacket packet) {
@@ -31,61 +35,60 @@ protected void handleImpl(@NotNull UnsafeMqttClient client, @NotNull ConnectInPa
         }
 
         authenticationService.auth(packet.getUsername(), packet.getPassword())
-            .filter(Boolean::booleanValue)
-            .flatMap(authenticated -> registerClient(client, packet))
-            .switchIfEmpty(fromRunnable(() -> client.reject(ConnectAckReasonCode.BAD_USER_NAME_OR_PASSWORD)))
+            .flatMap(ifTrue(client, packet, this::registerClient, BAD_USER_NAME_OR_PASSWORD, client::reject))
+            .flatMap(ifTrue(client, packet, this::restoreSession, CLIENT_IDENTIFIER_NOT_VALID, client::reject))
             .subscribe();
     }
 
     private @NotNull Mono<Boolean> registerClient(
         @NotNull UnsafeMqttClient client,
         @NotNull ConnectInPacket packet
     ) {
+
         var requestedClientId = packet.getClientId();
-        if (StringUtils.isEmpty(requestedClientId)) {
-            return processWithoutClientId(client, packet, requestedClientId);
+
+        if (StringUtils.isNotEmpty(requestedClientId)) {
+            return clientIdRegistry.register(requestedClientId)
+                .map(ifTrue(requestedClientId, client::setClientId));
         } else {
-            return processWithClientId(client, packet, requestedClientId);
+            return clientIdRegistry.generate()
+                .flatMap(newClientId -> clientIdRegistry.register(newClientId)
+                    .map(ifTrue(newClientId, client::setClientId)));
         }
     }
 
-    private @NotNull Mono<Boolean> processWithClientId(
+    private @NotNull Mono<Boolean> restoreSession(
         @NotNull UnsafeMqttClient client,
-        @NotNull ConnectInPacket packet,
-        @NotNull String requestedClientId
+        @NotNull ConnectInPacket packet
     ) {
-        return clientIdRegistry.register(requestedClientId)
-            .filter(Boolean::booleanValue)
-            .flatMap(registered -> fromCallable(() -> {
-                client.setClientId(requestedClientId);
-                onConnected(client, packet, requestedClientId);
-                return true;
-            }))
-            .switchIfEmpty(fromRunnable(() -> client.reject(ConnectAckReasonCode.CLIENT_IDENTIFIER_NOT_VALID)));
-    }
 
-    private @NotNull Mono<Boolean> processWithoutClientId(
-        @NotNull UnsafeMqttClient client,
-        @NotNull ConnectInPacket packet,
-        @NotNull String requestedClientId
-    ) {
-        return clientIdRegistry.generate()
-            .doOnNext(client::setClientId)
-            .flatMap(clientIdRegistry::register)
-            .filter(Boolean::booleanValue)
-            .flatMap(registered -> fromCallable(() -> {
-                onConnected(client, packet, requestedClientId);
-                return true;
-            }))
-            .switchIfEmpty(fromRunnable(() -> client.reject(ConnectAckReasonCode.CLIENT_IDENTIFIER_NOT_VALID)));
+        if (packet.isCleanStart()) {
+            return mqttSessionService.create(client.getClientId())
+                .flatMap(session -> onConnected(client, packet, session, false));
+        } else {
+            return mqttSessionService.restore(client.getClientId())
+                .flatMap(session -> onConnected(client, packet, session, true))
+                .switchIfEmpty(Mono.defer(() -> mqttSessionService.create(client.getClientId())
+                    .flatMap(session -> onConnected(client, packet, session,false))));
+        }
     }
 
-    private void onConnected(
+    private Mono<Boolean> onConnected(
         @NotNull UnsafeMqttClient client,
         @NotNull ConnectInPacket packet,
-        @NotNull String requestedClientId
+        @NotNull MqttSession session,
+        boolean sessionRestored
     ) {
 
+        var connection = client.getConnection();
+
+        // if it was closed in parallel
+        if (connection.isClosed()) {
+            // store the session again
+            return mqttSessionService.store(client.getClientId(), session);
+        }
+
+        client.setSession(session);
         client.configure(
             packet.getSessionExpiryInterval(),
             packet.getReceiveMax(),
@@ -97,13 +100,14 @@ private void onConnected(
         client.send(client.getPacketOutFactory().newConnectAck(
             client,
             ConnectAckReasonCode.SUCCESS,
-            false,
-            requestedClientId,
+            sessionRestored,
+            packet.getClientId(),
             packet.getSessionExpiryInterval(),
             packet.getKeepAlive()
         ));
-    }
 
+        return Mono.just(Boolean.TRUE);
+    }
 
     private boolean checkPacketException(@NotNull UnsafeMqttClient client, @NotNull ConnectInPacket packet) {
 

src/main/java/com/ss/mqtt/broker/service/MqttSessionService.java

@@ -6,9 +6,9 @@
 
 public interface MqttSessionService {
 
-    @NotNull Mono<MqttSession> getOrCreate(@NotNull String clientId);
+    @NotNull Mono<MqttSession> restore(@NotNull String clientId);
 
-    @NotNull Mono<MqttSession> createNew(@NotNull String clientId);
+    @NotNull Mono<MqttSession> create(@NotNull String clientId);
 
     @NotNull Mono<Boolean> store(@NotNull String clientId, @NotNull MqttSession session);
 }