Improve subscription process, part 2 (#51)

* add supporting replacing subscriptions

* improve error handling during subscribing

* update tests

* add validation packet flags

* update validation message flags

* extend tests for subscription process

Author: JavaSaBr

gradle/libs.versions.toml

@@ -1,6 +1,6 @@
 [versions]
 # https://gitlab.com/JavaSaBr/maven-repo/-/packages
-rlib = "10.0.alpha6"
+rlib = "10.0.alpha8"
 # https://mvnrepository.com/artifact/org.projectlombok/lombok
 lombok = "1.18.38"
 # https://mvnrepository.com/artifact/org.jspecify/jspecify

model/src/main/java/javasabr/mqtt/model/MqttClientConnectionConfig.java

@@ -10,4 +10,29 @@ public record MqttClientConnectionConfig(
     int topicAliasMaxValue,
     int keepAlive,
     boolean requestResponseInformation,
-    boolean requestProblemInformation) {}
+    boolean requestProblemInformation) {
+
+  public boolean subscriptionIdAvailable() {
+    return server.subscriptionIdAvailable();
+  }
+
+  public boolean retainAvailable() {
+    return server.retainAvailable();
+  }
+
+  public boolean wildcardSubscriptionAvailable() {
+    return server.wildcardSubscriptionAvailable();
+  }
+
+  public boolean sharedSubscriptionAvailable() {
+    return server.sharedSubscriptionAvailable();
+  }
+
+  public boolean sessionsEnabled() {
+    return server.sessionsEnabled();
+  }
+
+  public int maxTopicLevels() {
+    return server.maxTopicLevels();
+  }
+}

model/src/main/java/javasabr/mqtt/model/MqttProperties.java

@@ -45,9 +45,7 @@ public interface MqttProperties {
   int TOPIC_ALIAS_MAX = 0xFFFF;
   int TOPIC_ALIAS_NOT_SET = 0;
 
-  int SUBSCRIPTION_ID_UNDEFINED = 0;
-
-  int MESSAGE_ID_UNDEFINED = -1;
+  int SUBSCRIPTION_ID_IS_NOT_SET = 0;
   int MESSAGE_ID_IS_NOT_SET = 0;
 
   boolean SESSIONS_ENABLED_DEFAULT = true;

model/src/main/java/javasabr/mqtt/model/MqttServerConnectionConfig.java

@@ -126,4 +126,23 @@ public MqttServerConnectionConfig withSharedSubscriptionAvailable(boolean shared
         subscriptionIdAvailable,
         sharedSubscriptionAvailable);
   }
+
+  public MqttServerConnectionConfig withSubscriptionIdAvailable(boolean subscriptionIdAvailable) {
+    return new MqttServerConnectionConfig(
+        maxQos,
+        maxMessageSize,
+        maxStringLength,
+        maxBinarySize,
+        maxTopicLevels,
+        minKeepAliveTime,
+        receiveMaxPublishes,
+        topicAliasMaxValue,
+        defaultSessionExpiryInterval,
+        keepAliveEnabled,
+        sessionsEnabled,
+        retainAvailable,
+        wildcardSubscriptionAvailable,
+        subscriptionIdAvailable,
+        sharedSubscriptionAvailable);
+  }
 }

model/src/main/java/javasabr/mqtt/model/QoS.java

@@ -1,6 +1,8 @@
 package javasabr.mqtt.model;
 
 import javasabr.mqtt.model.reason.code.SubscribeAckReasonCode;
+import javasabr.rlib.common.util.NumberedEnum;
+import javasabr.rlib.common.util.NumberedEnumMap;
 import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -11,22 +13,28 @@
 @RequiredArgsConstructor
 @Accessors(fluent = true, chain = false)
 @FieldDefaults(level = AccessLevel.PRIVATE, makeFinal = true)
-public enum QoS {
+public enum QoS implements NumberedEnum<QoS> {
   AT_MOST_ONCE(0, SubscribeAckReasonCode.GRANTED_QOS_0),
   AT_LEAST_ONCE(1, SubscribeAckReasonCode.GRANTED_QOS_1),
   EXACTLY_ONCE(2, SubscribeAckReasonCode.GRANTED_QOS_2),
   INVALID(3, SubscribeAckReasonCode.IMPLEMENTATION_SPECIFIC_ERROR);
 
-  private static final QoS[] VALUES = values();
+  private static final NumberedEnumMap<QoS> NUMBERED_MAP =
+      new NumberedEnumMap<>(QoS.class);
 
   public static QoS ofCode(int level) {
-    if (level < 0 || level > EXACTLY_ONCE.ordinal()) {
-      return INVALID;
-    } else {
-      return VALUES[level];
-    }
+    return NUMBERED_MAP.resolve(level, QoS.INVALID);
   }
 
   int level;
   SubscribeAckReasonCode subscribeAckReasonCode;
+
+  @Override
+  public int number() {
+    return level;
+  }
+
+  public QoS lower(QoS alternative) {
+    return level > alternative.level ? alternative : this;
+  }
 }

model/src/main/java/javasabr/mqtt/model/subscribtion/Subscription.java

@@ -41,7 +41,7 @@ public record Subscription(
   public static Subscription minimal(TopicFilter topicFilter, QoS qos) {
     return new Subscription(
         topicFilter,
-        MqttProperties.SUBSCRIPTION_ID_UNDEFINED,
+        MqttProperties.SUBSCRIPTION_ID_IS_NOT_SET,
         qos,
         SubscribeRetainHandling.SEND,
         true,

model/src/main/java/javasabr/mqtt/model/topic/tree/ConcurrentTopicTree.java

@@ -10,6 +10,7 @@
 import javasabr.rlib.common.ThreadSafe;
 import lombok.AccessLevel;
 import lombok.experimental.FieldDefaults;
+import org.jspecify.annotations.Nullable;
 
 @FieldDefaults(level = AccessLevel.PRIVATE, makeFinal = true)
 public class ConcurrentTopicTree implements ThreadSafe {
@@ -20,8 +21,9 @@ public ConcurrentTopicTree() {
     this.rootNode = new TopicNode();
   }
 
-  public void subscribe(SubscriptionOwner owner, Subscription subscription) {
-    rootNode.subscribe(0, owner, subscription, subscription.topicFilter());
+  @Nullable
+  public SingleSubscriber subscribe(SubscriptionOwner owner, Subscription subscription) {
+    return rootNode.subscribe(0, owner, subscription, subscription.topicFilter());
   }
 
   public boolean unsubscribe(SubscriptionOwner owner, TopicFilter topicFilter) {

model/src/main/java/javasabr/mqtt/model/topic/tree/TopicNode.java

@@ -35,13 +35,16 @@ class TopicNode extends TopicTreeBase {
   @Nullable
   volatile LockableArray<Subscriber> subscribers;
 
-  public void subscribe(int level, SubscriptionOwner owner, Subscription subscription, TopicFilter topicFilter) {
+  /**
+   * @return the previous subscription from the same owner
+   */
+  @Nullable
+  public SingleSubscriber subscribe(int level, SubscriptionOwner owner, Subscription subscription, TopicFilter topicFilter) {
     if (level == topicFilter.levelsCount()) {
-      addSubscriber(getOrCreateSubscribers(), owner, subscription, topicFilter);
-      return;
+      return addSubscriber(getOrCreateSubscribers(), owner, subscription, topicFilter);
     }
     TopicNode childNode = getOrCreateChildNode(topicFilter.segment(level));
-    childNode.subscribe(level + 1, owner, subscription, topicFilter);
+    return childNode.subscribe(level + 1, owner, subscription, topicFilter);
   }
 
   public boolean unsubscribe(int level, SubscriptionOwner owner, TopicFilter topicFilter) {

model/src/main/java/javasabr/mqtt/model/topic/tree/TopicTreeBase.java

@@ -20,7 +20,11 @@
 @FieldDefaults(level = AccessLevel.PROTECTED, makeFinal = true)
 abstract class TopicTreeBase {
 
-  protected static void addSubscriber(
+  /**
+   * @return previous subscriber with the same owner
+   */
+  @Nullable
+  protected static SingleSubscriber addSubscriber(
       LockableArray<Subscriber> subscribers,
       SubscriptionOwner owner,
       Subscription subscription,
@@ -29,14 +33,30 @@ protected static void addSubscriber(
     try {
       if (topicFilter instanceof SharedTopicFilter stf) {
         addSharedSubscriber(subscribers, owner, subscription, stf);
+        return null;
       } else {
+        SingleSubscriber previous = removePreviousIfExist(subscribers, owner);
         subscribers.add(new SingleSubscriber(owner, subscription));
+        return previous;
       }
     } finally {
       subscribers.writeUnlock(stamp);
     }
   }
 
+  @Nullable
+  private static SingleSubscriber removePreviousIfExist(
+      LockableArray<Subscriber> subscribers,
+      SubscriptionOwner owner) {
+    int index = subscribers.indexOf(Subscriber::resolveOwner, owner);
+    if (index < 0) {
+      return null;
+    }
+    return subscribers
+        .remove(index)
+        .resolveSingle();
+  }
+
   private static void addSharedSubscriber(
       LockableArray<Subscriber> subscribers,
       SubscriptionOwner owner,

model/src/test/groovy/javasabr/mqtt/model/topic/tree/TopicTreeTest.groovy

@@ -539,14 +539,94 @@ class TopicTreeTest extends UnitSpecification {
         !id3WasUnsubscribed
   }
 
+  def "should replace the same subscriptions"() {
+    given:
+        ConcurrentTopicTree topicTree = new ConcurrentTopicTree()
+        def owner1 = makeOwner("id1")
+        def originalSub = makeSubscription('topic/name1')
+        def replacementSub = makeSubscription('topic/name1')
+        topicTree.subscribe(makeOwner("id2"), makeSubscription('topic/name1'))
+        topicTree.subscribe(makeOwner("id3"), makeSubscription('topic/name1'))
+    when:
+        def previous = topicTree.subscribe(owner1, originalSub)
+        def matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 3
+        previous == null;
+    when:
+        previous = topicTree.subscribe(owner1, replacementSub)
+        matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 3
+        matched.first().subscription() == replacementSub
+        previous != null
+        previous.subscription() == originalSub
+  }
+
+  def "should extend shared subscription group on multiply subscribing by the same topic"() {
+    given:
+        ConcurrentTopicTree topicTree = new ConcurrentTopicTree()
+        def owner1 = makeOwner("id1")
+        def owner2 = makeOwner("id2")
+        topicTree.subscribe(owner1, makeSharedSubscription('$share/group1/topic/name1'))
+        topicTree.subscribe(owner2, makeSharedSubscription('$share/group1/topic/name1'))
+    when:
+        def matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 1
+        matched.first().owner() == owner2
+    when:
+        matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 1
+        matched.first().owner() == owner1
+    when:
+        matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 1
+        matched.first().owner() == owner2
+    when:
+        topicTree.subscribe(owner1, makeSharedSubscription('$share/group1/topic/name1'))
+        matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 1
+        matched.first().owner() == owner2
+    when:
+        matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 1
+        matched.first().owner() == owner1
+    when:
+        matched = topicTree
+            .matches(TopicName.valueOf("topic/name1"))
+            .toSet()
+    then:
+        matched.size() == 1
+        matched.first().owner() == owner1
+  }
+
   static def makeOwner(String id) {
     return new TestSubscriptionOwner(id)
   }
 
   static def makeSubscription(String topicFilter) {
     return new Subscription(
         TopicFilter.valueOf(topicFilter),
-        MqttProperties.SUBSCRIPTION_ID_UNDEFINED,
+        MqttProperties.SUBSCRIPTION_ID_IS_NOT_SET,
         QoS.AT_LEAST_ONCE,
         SubscribeRetainHandling.SEND,
         true,
@@ -556,7 +636,7 @@ class TopicTreeTest extends UnitSpecification {
   static def makeSharedSubscription(String topicFilter) {
     return new Subscription(
         SharedTopicFilter.valueOf(topicFilter),
-        MqttProperties.SUBSCRIPTION_ID_UNDEFINED,
+        MqttProperties.SUBSCRIPTION_ID_IS_NOT_SET,
         QoS.AT_LEAST_ONCE,
         SubscribeRetainHandling.SEND,
         true,
@@ -566,7 +646,7 @@ class TopicTreeTest extends UnitSpecification {
   static def makeSubscription(String topicFilter, int qos) {
     return new Subscription(
         TopicFilter.valueOf(topicFilter),
-        MqttProperties.SUBSCRIPTION_ID_UNDEFINED,
+        MqttProperties.SUBSCRIPTION_ID_IS_NOT_SET,
         QoS.ofCode(qos),
         SubscribeRetainHandling.SEND,
         true,