[broker-30] Rewrite retained messages collecting

Author: crazyrokr

model/src/main/java/javasabr/mqtt/model/topic/AbstractTopic.java

@@ -41,7 +41,7 @@ public int levelsCount() {
     return segments.length;
   }
 
-  String lastSegment() {
+  public String lastSegment() {
     return segments[segments.length - 1];
   }
 

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

@@ -2,30 +2,30 @@
 
 import javasabr.mqtt.model.publishing.Publish;
 import javasabr.mqtt.model.topic.TopicFilter;
-import javasabr.mqtt.model.topic.TopicName;
+import javasabr.rlib.collections.array.Array;
+import javasabr.rlib.collections.array.MutableArray;
 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 ConcurrentRetainedMessageTree implements ThreadSafe {
 
-  TopicMessageNode rootNode;
+  RetainedMessageNode rootNode;
 
   public ConcurrentRetainedMessageTree() {
-    this.rootNode = new TopicMessageNode();
+    this.rootNode = new RetainedMessageNode();
   }
 
   public void retainMessage(Publish message) {
-    rootNode.retainMessage(0, message, message.topicName());
+    if (message.retained()) {
+      rootNode.retainMessage(0, message, message.topicName());
+    }
   }
 
-  public @Nullable Publish getRetainedMessage(TopicName topicName) {
-    return rootNode.getRetainedMessage(0, topicName);
-  }
-
-  public @Nullable Publish getRetainedMessage(TopicFilter topicFilter) {
-    return rootNode.getRetainedMessage(0, topicFilter);
+  public Array<Publish> getRetainedMessage(TopicFilter topicFilter) {
+    var resultArray = MutableArray.ofType(Publish.class);
+    rootNode.collectRetainedMessages(0, topicFilter, topicFilter.levelsCount() - 1, resultArray);
+    return resultArray;
   }
 }

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

@@ -0,0 +1,150 @@
+package javasabr.mqtt.model.topic.tree;
+
+import static javasabr.mqtt.model.topic.TopicFilter.MULTI_LEVEL_WILDCARD;
+import static javasabr.mqtt.model.topic.TopicFilter.SINGLE_LEVEL_WILDCARD;
+
+import java.util.Objects;
+import java.util.PriorityQueue;
+import java.util.Queue;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+import javasabr.mqtt.base.util.DebugUtils;
+import javasabr.mqtt.model.publishing.Publish;
+import javasabr.mqtt.model.topic.TopicFilter;
+import javasabr.mqtt.model.topic.TopicName;
+import javasabr.rlib.collections.array.MutableArray;
+import javasabr.rlib.collections.dictionary.DictionaryFactory;
+import javasabr.rlib.collections.dictionary.LockableRefToRefDictionary;
+import lombok.AccessLevel;
+import lombok.Getter;
+import lombok.experimental.Accessors;
+import lombok.experimental.FieldDefaults;
+import org.jspecify.annotations.Nullable;
+
+@Getter(AccessLevel.PACKAGE)
+@Accessors(fluent = true, chain = false)
+@FieldDefaults(level = AccessLevel.PRIVATE)
+class RetainedMessageNode {
+
+  private final static Supplier<RetainedMessageNode> TOPIC_NODE_FACTORY = RetainedMessageNode::new;
+
+  static {
+    DebugUtils.registerIncludedFields("childNodes", "retainedMessage");
+  }
+
+  @Nullable
+  volatile LockableRefToRefDictionary<String, RetainedMessageNode> childNodes;
+  final AtomicReference<@Nullable Publish> retainedMessage = new AtomicReference<>();
+
+  public void retainMessage(int level, Publish message, TopicName topicName) {
+    if (level + 1 == topicName.levelsCount()) {
+      retainedMessage.set(message.payload().length == 0 ? null : message);
+      return;
+    }
+    RetainedMessageNode childNode = getOrCreateChildNode(topicName.segment(level));
+    childNode.retainMessage(level + 1, message, topicName);
+  }
+
+  public void collectRetainedMessages(int level, TopicFilter topicFilter, int lastLevel, MutableArray<Publish> result) {
+    String segment = topicFilter.segment(level);
+    Publish publish = retainedMessage.get();
+    if (Objects.equals(segment, MULTI_LEVEL_WILDCARD)) {
+      collectAllMessages(this, result);
+    } else if (Objects.equals(segment, SINGLE_LEVEL_WILDCARD)) {
+      var childNodes = childNodes();
+      if (childNodes == null) {
+        return;
+      }
+      long stamp = childNodes.readLock();
+      try {
+        for (RetainedMessageNode n : childNodes) {
+          n.collectRetainedMessages(level + 1, topicFilter, lastLevel, result);
+        }
+      } finally {
+        childNodes.readUnlock(stamp);
+      }
+    } else if (level == lastLevel && publish != null && Objects.equals(segment, publish.topicName().lastSegment())) {
+      result.add(publish);
+    } else {
+      RetainedMessageNode topicFilterNode = childNode(segment);
+      if (topicFilterNode == null) {
+        return;
+      }
+      topicFilterNode.collectRetainedMessages(level + 1, topicFilter, lastLevel, result);
+    }
+  }
+
+  private void collectAllMessages(RetainedMessageNode node, MutableArray<Publish> result) {
+    Queue<RetainedMessageNode> queue = new PriorityQueue<>();
+    queue.add(node);
+    while (!queue.isEmpty()) {
+      RetainedMessageNode poll = queue.poll();
+      Publish message = poll.retainedMessage.get();
+      if (message != null) {
+        result.add(message);
+      }
+      var childNodes = poll.childNodes();
+      if (childNodes == null) {
+        continue;
+      }
+      long stamp = childNodes.readLock();
+      try {
+        for (RetainedMessageNode n : childNodes) {
+          queue.add(n);
+        }
+      } finally {
+        childNodes.readUnlock(stamp);
+      }
+    }
+  }
+
+  @Nullable
+  private RetainedMessageNode childNode(String segment) {
+    LockableRefToRefDictionary<String, RetainedMessageNode> childNodes = childNodes();
+    if (childNodes == null) {
+      return null;
+    }
+    long stamp = childNodes.readLock();
+    try {
+      return childNodes.get(segment);
+    } finally {
+      childNodes.readUnlock(stamp);
+    }
+  }
+
+  private RetainedMessageNode getOrCreateChildNode(String segment) {
+    LockableRefToRefDictionary<String, RetainedMessageNode> childNodes = getOrCreateChildNodes();
+    long stamp = childNodes.readLock();
+    try {
+      RetainedMessageNode topicFilterNode = childNodes.get(segment);
+      if (topicFilterNode != null) {
+        return topicFilterNode;
+      }
+    } finally {
+      childNodes.readUnlock(stamp);
+    }
+    stamp = childNodes.writeLock();
+    try {
+      return childNodes.getOrCompute(segment, TOPIC_NODE_FACTORY);
+    } finally {
+      childNodes.writeUnlock(stamp);
+    }
+  }
+
+  private LockableRefToRefDictionary<String, RetainedMessageNode> getOrCreateChildNodes() {
+    if (childNodes == null) {
+      synchronized (this) {
+        if (childNodes == null) {
+          childNodes = DictionaryFactory.stampedLockBasedRefToRefDictionary();
+        }
+      }
+    }
+    //noinspection ConstantConditions
+    return childNodes;
+  }
+
+  @Override
+  public String toString() {
+    return DebugUtils.toJsonString(this);
+  }
+}

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

@@ -4,10 +4,11 @@
 import javasabr.mqtt.model.subscriber.SingleSubscriber;
 import javasabr.mqtt.model.topic.TopicFilter;
 import javasabr.mqtt.service.publish.handler.PublishHandlingResult;
+import javasabr.rlib.collections.array.Array;
 
 public interface PublishDeliveringService {
 
   PublishHandlingResult startDelivering(Publish publish, SingleSubscriber subscriber);
 
-  PublishHandlingResult deliverRetainedMessages(TopicFilter topicFilter, SingleSubscriber subscriber);
+  Array<PublishHandlingResult> deliverRetainedMessages(TopicFilter topicFilter, SingleSubscriber subscriber);
 }

service/src/main/java/javasabr/mqtt/service/PublishDeliveringService.java

@@ -9,6 +9,8 @@
 import javasabr.mqtt.service.PublishDeliveringService;
 import javasabr.mqtt.service.publish.handler.MqttPublishOutMessageHandler;
 import javasabr.mqtt.service.publish.handler.PublishHandlingResult;
+import javasabr.rlib.collections.array.Array;
+import javasabr.rlib.collections.array.MutableArray;
 import lombok.AccessLevel;
 import lombok.CustomLog;
 import lombok.experimental.FieldDefaults;
@@ -20,7 +22,7 @@ public class DefaultPublishDeliveringService implements PublishDeliveringService
 
   @Nullable
   MqttPublishOutMessageHandler[] publishOutMessageHandlers;
-  ConcurrentRetainedMessageTree topicTree;
+  ConcurrentRetainedMessageTree retainedMessageTree;
 
   public DefaultPublishDeliveringService(
       Collection<? extends MqttPublishOutMessageHandler> knownPublishOutHandlers) {
@@ -42,17 +44,15 @@ public DefaultPublishDeliveringService(
       }
       handlers[qos.level()] = knownPublishOutHandler;
     }
-    this.topicTree = new ConcurrentRetainedMessageTree();
+    this.retainedMessageTree = new ConcurrentRetainedMessageTree();
     this.publishOutMessageHandlers = handlers;
     log.info(publishOutMessageHandlers, DefaultPublishDeliveringService::buildServiceDescription);
   }
 
   @Override
   public PublishHandlingResult startDelivering(Publish publish, SingleSubscriber subscriber) {
     try {
-      if (publish.retained()) {
-        topicTree.retainMessage(publish);
-      }
+      retainedMessageTree.retainMessage(publish);
       //noinspection DataFlowIssue
       return publishOutMessageHandlers[subscriber.qos().level()].handle(publish, subscriber);
     } catch (IndexOutOfBoundsException | NullPointerException ex) {
@@ -62,9 +62,13 @@ public PublishHandlingResult startDelivering(Publish publish, SingleSubscriber s
   }
 
   @Override
-  public PublishHandlingResult deliverRetainedMessages(TopicFilter topicFilter, SingleSubscriber subscriber) {
-    Publish retainedMessage = topicTree.getRetainedMessage(topicFilter);
-    return startDelivering(retainedMessage, subscriber);
+  public Array<PublishHandlingResult> deliverRetainedMessages(TopicFilter topicFilter, SingleSubscriber subscriber) {
+    Array<Publish> retainedMessage = retainedMessageTree.getRetainedMessage(topicFilter);
+    MutableArray<PublishHandlingResult> result = MutableArray.ofType(PublishHandlingResult.class);
+    for (Publish message : retainedMessage) {
+      result.add(startDelivering(message, subscriber));
+    }
+    return result;
   }
 
   private static String buildServiceDescription(

service/src/main/java/javasabr/mqtt/service/impl/DefaultPublishDeliveringService.java

@@ -194,23 +194,27 @@ private void sendRetainedMessages(
       RequestedSubscription requestedSubscription = subscriptions.get(i);
       SubscribeAckReasonCode subscribeAckReasonCode = subscribeResults.get(i);
       Subscription subscription = subs.get(i);
-      if (subscribeAckReasonCode.ordinal() < 3) {
-        TopicFilter topicFilter = TopicFilter.valueOf(requestedSubscription.rawTopicFilter());
-        SingleSubscriber singleSubscriber = new SingleSubscriber(client, subscription);
-        PublishHandlingResult result = publishDeliveringService.deliverRetainedMessages(topicFilter, singleSubscriber);
+      if (subscribeAckReasonCode.ordinal() > 2) {
+        // TODO handle error
+        continue;
+      }
+      TopicFilter topicFilter = TopicFilter.valueOf(requestedSubscription.rawTopicFilter());
+      SingleSubscriber singleSubscriber = new SingleSubscriber(client, subscription);
+      var results = publishDeliveringService.deliverRetainedMessages(topicFilter, singleSubscriber);
+      for (PublishHandlingResult result : results) {
         if (result.error()) {
           errorResult = result;
-        } else if(result == PublishHandlingResult.SUCCESS) {
+        } else if (result == PublishHandlingResult.SUCCESS) {
           count++;
         }
         if (errorResult != null) {
           log.debug(client.clientId(), errorResult,
               "[%s] Found final error:[%s] during sending retained messages"::formatted);
-          // handleError(client, publish, errorResult);
+          // TODO handleError(client, publish, errorResult);
         } else {
           log.debug(client.clientId(), count,
               "[%s] Successfully started delivering retained messages to [%s] subscribers"::formatted);
-          // handleSuccessfulResult(client, publish, count);
+          // TODO handleSuccessfulResult(client, publish, count);
         }
       }
     }