[broker-11] finish fixing issues with mqtt sessions

Author: JavaSaBr

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

@@ -41,7 +41,7 @@ public interface MqttPropertyConstants {
 
     int SUBSCRIPTION_ID_NOT_DEFINED = 0;
 
-    boolean SESSIONS_ENABLED_DEFAULT = false;
+    boolean SESSIONS_ENABLED_DEFAULT = true;
     boolean KEEP_ALIVE_ENABLED_DEFAULT = false;
     boolean RETAIN_AVAILABLE_DEFAULT = false;
     boolean WILDCARD_SUBSCRIPTION_AVAILABLE_DEFAULT = false;

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

@@ -4,8 +4,10 @@
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.Setter;
+import lombok.ToString;
 import org.jetbrains.annotations.NotNull;
 
+@ToString
 @RequiredArgsConstructor
 public class DefaultMqttSession implements UnsafeMqttSession {
 

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

@@ -38,8 +38,8 @@ public abstract class AbstractMqttClientReleaseHandler<T extends AbstractMqttCli
 
         Mono<?> asyncActions = null;
 
-        if (session != null) {
-            asyncActions = sessionService.store(clientId, session);
+        if (session != null && client.getConnectionConfig().isSessionsEnabled()) {
+            asyncActions = sessionService.store(clientId, session, client.getSessionExpiryInterval());
             client.setSession(null);
         }
 

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

@@ -93,6 +93,11 @@ protected void readVariableHeader(@NotNull MqttConnection connection, @NotNull B
         }
     }
 
+    @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;

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

@@ -82,15 +82,14 @@ private Mono<Boolean> onConnected(
     ) {
 
         var connection = client.getConnection();
+        var config = connection.getConfig();
 
         // if it was closed in parallel
-        if (connection.isClosed()) {
+        if (connection.isClosed() && config.isSessionsEnabled()) {
             // store the session again
-            return mqttSessionService.store(client.getClientId(), session);
+            return mqttSessionService.store(client.getClientId(), session, config.getDefaultSessionExpiryInterval());
         }
 
-        var config = connection.getConfig();
-
         // select result keep alive time
         var minimalKeepAliveTime = Math.max(config.getMinKeepAliveTime(), packet.getKeepAlive());
         var keepAlive = config.isKeepAliveEnabled() ? minimalKeepAliveTime : SERVER_KEEP_ALIVE_DISABLED;
@@ -115,6 +114,7 @@ private Mono<Boolean> onConnected(
         var topicAliasMaximum = packet.getTopicAliasMaximum() == TOPIC_ALIAS_MAXIMUM_UNDEFINED ?
             TOPIC_ALIAS_MAXIMUM_DISABLED : Math.min(packet.getTopicAliasMaximum(), config.getTopicAliasMaximum());
 
+        client.setSession(session);
         client.configure(
             sessionExpiryInterval,
             receiveMax,

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

@@ -10,5 +10,5 @@ public interface MqttSessionService {
 
     @NotNull Mono<MqttSession> create(@NotNull String clientId);
 
-    @NotNull Mono<Boolean> store(@NotNull String clientId, @NotNull MqttSession session);
+    @NotNull Mono<Boolean> store(@NotNull String clientId, @NotNull MqttSession session, long expiryInterval);
 }

src/main/java/com/ss/mqtt/broker/service/impl/InMemoryMqttSessionService.java

@@ -9,11 +9,13 @@
 import com.ss.rlib.common.util.dictionary.ConcurrentObjectDictionary;
 import com.ss.rlib.common.util.dictionary.DictionaryFactory;
 import com.ss.rlib.common.util.dictionary.ObjectDictionary;
+import lombok.extern.log4j.Log4j2;
 import org.jetbrains.annotations.NotNull;
 import reactor.core.publisher.Mono;
 
 import java.io.Closeable;
 
+@Log4j2
 public class InMemoryMqttSessionService implements MqttSessionService, Closeable {
 
     private final @NotNull ConcurrentObjectDictionary<String, MqttSession> storedSession;
@@ -34,18 +36,42 @@ public InMemoryMqttSessionService(int cleanInterval) {
 
     @Override
     public @NotNull Mono<MqttSession> restore(@NotNull String clientId) {
+
         var session = storedSession.getInWriteLock(clientId, ObjectDictionary::remove);
+
+        if (session != null) {
+            log.debug("Restored session for client {}", clientId);
+        } else {
+            log.debug("No stored session for client {}", clientId);
+        }
+
         return Mono.justOrEmpty(session);
     }
 
     @Override
     public @NotNull Mono<MqttSession> create(@NotNull String clientId) {
+
+        var session = storedSession.getInWriteLock(clientId, ObjectDictionary::remove);
+
+        if (session != null) {
+            log.debug("Removed old session for client {}", clientId);
+        }
+
+        log.debug("Created new session for client {}", clientId);
+
         return Mono.just(new DefaultMqttSession(clientId));
     }
 
     @Override
-    public @NotNull Mono<Boolean> store(@NotNull String clientId, @NotNull MqttSession session) {
+    public @NotNull Mono<Boolean> store(@NotNull String clientId, @NotNull MqttSession session, long expiryInterval) {
+
+        var unsafe = (MqttSession.UnsafeMqttSession) session;
+        unsafe.setExpirationTime(System.currentTimeMillis() + (expiryInterval * 1000));
+
         storedSession.runInWriteLock(clientId, session, ObjectDictionary::put);
+
+        log.debug("Stored session for client {}", clientId);
+
         return Mono.just(Boolean.TRUE);
     }
 
@@ -79,6 +105,8 @@ private void cleanup() {
 
                     var removed = dictionary.remove(session.getClientId());
 
+                    log.debug("Removed expired session for client {}", session.getClientId());
+
                     // if we already have new session under the same client id
                     if (removed != null && removed != session) {
                         dictionary.put(session.getClientId(), removed);

src/test/groovy/com/ss/mqtt/broker/test/integration/ConnectionTest.groovy

@@ -21,7 +21,7 @@ class ConnectionTest extends MqttBrokerTest {
         then:
             result.reasonCode == Mqtt5ConnAckReasonCode.SUCCESS
             result.sessionExpiryInterval.present
-            result.sessionExpiryInterval.getAsLong() == MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DISABLED
+            result.sessionExpiryInterval.getAsLong() == MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DEFAULT
             result.serverKeepAlive.present
             result.serverKeepAlive.getAsInt() == MqttPropertyConstants.SERVER_KEEP_ALIVE_DISABLED
             !result.serverReference.present
@@ -40,7 +40,7 @@ class ConnectionTest extends MqttBrokerTest {
         then:
             result.reasonCode == Mqtt5ConnAckReasonCode.SUCCESS
             result.sessionExpiryInterval.present
-            result.sessionExpiryInterval.getAsLong() == MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DISABLED
+            result.sessionExpiryInterval.getAsLong() == MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DEFAULT
             result.serverKeepAlive.present
             result.serverKeepAlive.getAsInt() == MqttPropertyConstants.SERVER_KEEP_ALIVE_DISABLED
             !result.serverReference.present

src/test/groovy/com/ss/mqtt/broker/test/integration/service/ClientIdRegistryTest.groovy

@@ -117,7 +117,7 @@ class ClientIdRegistryTest extends MqttBrokerTest {
             !clientIdRegistry.register(clientId).block()
         when:
             client.disconnect().join()
-            Thread.sleep(50)
+            Thread.sleep(100)
         then:
             clientIdRegistry.register(clientId).block()
             clientIdRegistry.unregister(clientId).block()

src/test/groovy/com/ss/mqtt/broker/test/integration/service/MqttSessionServiceTest.groovy

@@ -1,6 +1,7 @@
 package com.ss.mqtt.broker.test.integration.service
 
 import com.hivemq.client.mqtt.mqtt5.message.connect.connack.Mqtt5ConnAckReasonCode
+import com.ss.mqtt.broker.config.MqttConnectionConfig
 import com.ss.mqtt.broker.service.ClientIdRegistry
 import com.ss.mqtt.broker.service.MqttSessionService
 import com.ss.mqtt.broker.test.integration.MqttBrokerTest
@@ -14,6 +15,9 @@ class MqttSessionServiceTest extends MqttBrokerTest {
     @Autowired
     MqttSessionService mqttSessionService
 
+    @Autowired
+    MqttConnectionConfig connectionConfig
+    
     def "subscriber should create and re-use mqtt session"() {
         given:
             def clientId = clientIdRegistry.generate().block()
@@ -27,21 +31,21 @@ class MqttSessionServiceTest extends MqttBrokerTest {
             mqttSessionService.restore(clientId).block() == null
         when:
             client.disconnect().join()
-            Thread.sleep(50)
+            Thread.sleep(100)
             def restored = mqttSessionService.restore(clientId).block()
-            mqttSessionService.store(clientId, restored).block()
         then:
             restored != null
         when:
+            mqttSessionService.store(clientId, restored, connectionConfig.getDefaultSessionExpiryInterval()).block()
             client.connect().join()
             shouldNoSession = mqttSessionService.restore(clientId).block()
         then:
             shouldNoSession == null
         when:
             client.disconnect().join()
-            Thread.sleep(50)
-            def restored2 = mqttSessionService.restore(clientId).block()
+            Thread.sleep(100)
+            restored = mqttSessionService.restore(clientId).block()
         then:
-            restored == restored2
+            restored != null
     }
 }