Purge nodes without subscriptions as described by PR #274

broker/src/main/java/io/moquette/spi/impl/subscriptions/SubscriptionsStore.java

@@ -122,9 +122,10 @@ public class SubscriptionsStore {
         do {
             oldRoot = subscriptions.get();
             couple = recreatePath(newSubscription.topicFilter, oldRoot);
-            couple.createdNode.addSubscription(newSubscription); // createdNode could be null?
+            couple.createdNode.addSubscription(newSubscription); //createdNode could be null?
-            // spin lock repeating till we can, swap root, if can't swap just re-do the operation
+            couple.root.recalculateSubscriptionsSize();
-        } while (!subscriptions.compareAndSet(oldRoot, couple.root));
+            //spin lock repeating till we can, swap root, if can't swap just re-do the operation
+        } while(!subscriptions.compareAndSet(oldRoot, couple.root));
         LOG.debug("A subscription has been added. Root = {}, oldRoot = {}.", couple.root, oldRoot);
     }
 
@@ -165,8 +166,9 @@ public class SubscriptionsStore {
             couple = recreatePath(topic, oldRoot);
 
             couple.createdNode.remove(new ClientTopicCouple(clientID, topic));
-            // spin lock repeating till we can, swap root, if can't swap just re-do the operation
+            couple.root.recalculateSubscriptionsSize();
-        } while (!subscriptions.compareAndSet(oldRoot, couple.root));
+            //spin lock repeating till we can, swap root, if can't swap just re-do the operation
+        } while(!subscriptions.compareAndSet(oldRoot, couple.root));
     }
 
     /**

broker/src/main/java/io/moquette/spi/impl/subscriptions/Topic.java

@@ -40,11 +40,6 @@ public class Topic implements Serializable {
         this.topic = topic;
     }
 
-    @Override
-    public String toString() {
-        return topic;
-    }
-
     public List<Token> getTokens() {
         if (tokens == null) {
             try {
@@ -115,8 +110,6 @@ public class Topic implements Serializable {
     /**
      * Verify if the 2 topics matching respecting the rules of MQTT Appendix A
      *
-     * @param msgTopic
-     *            the topic to match from the message
      * @param subscriptionTopic
      *            the topic filter of the subscription
      * @return true if the two topics match.
@@ -152,6 +145,11 @@ public class Topic implements Serializable {
         return i == msgTokens.size();
     }
 
+    @Override
+    public String toString() {
+        return topic;
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (obj == null) {

broker/src/main/java/io/moquette/spi/impl/subscriptions/TreeNode.java

@@ -24,9 +24,10 @@ class TreeNode {
 
     Token m_token;
     List<TreeNode> m_children = new ArrayList<>();
-    // TODO move to set of ClientIDthe set of clientIDs that has subscriptions to this topic
     Set<ClientTopicCouple> m_subscriptions = new HashSet<>();
 
+    private int subtreeSubscriptions = 0;
+
     TreeNode() {
     }
 
@@ -40,10 +41,12 @@ class TreeNode {
 
     void addSubscription(ClientTopicCouple s) {
         m_subscriptions.add(s);
+        this.subtreeSubscriptions++;
     }
 
     void addChild(TreeNode child) {
         m_children.add(child);
+        this.subtreeSubscriptions += child.subtreeSubscriptions;
     }
 
     /**
@@ -54,6 +57,7 @@ class TreeNode {
         copy.m_children = new ArrayList<>(m_children);
         copy.m_subscriptions = new HashSet<>(m_subscriptions);
         copy.m_token = m_token;
+        copy.subtreeSubscriptions = this.subtreeSubscriptions;
         return copy;
     }
 
@@ -73,6 +77,7 @@ class TreeNode {
     void updateChild(TreeNode oldChild, TreeNode newChild) {
         m_children.remove(oldChild);
         m_children.add(newChild);
+        this.subtreeSubscriptions += newChild.subtreeSubscriptions - oldChild.subtreeSubscriptions;
     }
 
     Collection<ClientTopicCouple> subscriptions() {
@@ -81,6 +86,7 @@ class TreeNode {
 
     public void remove(ClientTopicCouple clientTopicCouple) {
         m_subscriptions.remove(clientTopicCouple);
+        this.subtreeSubscriptions--;
     }
 
     // TODO smell a query method that return the result modifing the parameter (matchingSubs)
@@ -120,11 +126,7 @@ class TreeNode {
      * Return the number of registered subscriptions
      */
     int size() {
-        int res = m_subscriptions.size();
+        return this.subtreeSubscriptions;
-        for (TreeNode child : m_children) {
-            res += child.size();
-        }
-        return res;
     }
 
     /**
@@ -133,6 +135,27 @@ class TreeNode {
     TreeNode removeClientSubscriptions(String clientID) {
         // collect what to delete and then delete to avoid ConcurrentModification
         TreeNode newSubRoot = this.copy();
+        remoteSubscriptions(clientID, newSubRoot);
+        removeSubscriptionFromChildren(clientID, newSubRoot);
+        return newSubRoot;
+    }
+
+    private void removeSubscriptionFromChildren(String clientID, TreeNode newSubRoot) {
+        int newSubtreeSubscriptions = 0;
+        //go deep
+        List<TreeNode> newChildren = new ArrayList<>(newSubRoot.m_children.size());
+        for (TreeNode child : newSubRoot.m_children) {
+            final TreeNode purgedSubtree = child.removeClientSubscriptions(clientID);
+            if (purgedSubtree.size() != 0) {
+                newSubtreeSubscriptions += purgedSubtree.size();
+                newChildren.add(purgedSubtree);
+            }
+        }
+        newSubRoot.m_children = newChildren;
+        newSubRoot.subtreeSubscriptions += newSubtreeSubscriptions;
+    }
+
+    private void remoteSubscriptions(String clientID, TreeNode newSubRoot) {
         List<ClientTopicCouple> subsToRemove = new ArrayList<>();
         for (ClientTopicCouple s : newSubRoot.m_subscriptions) {
             if (s.clientID.equals(clientID)) {
@@ -143,13 +166,15 @@ class TreeNode {
         for (ClientTopicCouple s : subsToRemove) {
             newSubRoot.m_subscriptions.remove(s);
         }
+        newSubRoot.subtreeSubscriptions = newSubRoot.m_subscriptions.size();
+    }
 
-        // go deep
+    int recalculateSubscriptionsSize() {
-        List<TreeNode> newChildren = new ArrayList<>(newSubRoot.m_children.size());
+        int res = m_subscriptions.size();
-        for (TreeNode child : newSubRoot.m_children) {
+        for (TreeNode child : m_children) {
-            newChildren.add(child.removeClientSubscriptions(clientID));
+            res += child.recalculateSubscriptionsSize();
         }
-        newSubRoot.m_children = newChildren;
+        this.subtreeSubscriptions = res;
-        return newSubRoot;
+        return res;
     }
 }

broker/src/test/java/io/moquette/spi/impl/ProtocolProcessorTest.java

@@ -284,8 +284,9 @@ public class ProtocolProcessorTest {
         assertEquals(0, subscriptions.size());
 
         m_processor.processSubscribe(m_channel, msg);
-
+        assertEquals(1, subscriptions.size());
-        // Exercise
+                
+        //Exercise
         m_processor.processSubscribe(m_channel, msg);
 
         // Verify

broker/src/test/java/io/moquette/spi/impl/subscriptions/TreeNodeTest.java

@@ -0,0 +1,60 @@
+package io.moquette.spi.impl.subscriptions;
+
+import io.moquette.spi.ISessionsStore;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+/**
+ * Created by andrea on 03/03/17.
+ */
+public class TreeNodeTest {
+
+    @Test
+    public void testRemoveTreeNodesAfterEmptySubscriptions() {
+        TreeNode root = new TreeNode();
+        TreeNode subNode1 = new TreeNode();
+
+        subNode1.addSubscription(new ISessionsStore.ClientTopicCouple("Client1", new Topic("/topic1")));
+        root.addChild(subNode1);
+
+        TreeNode updatedRoot = root.removeClientSubscriptions("Client1");
+
+        assertTrue("Empty subscriptions node should be removed", updatedRoot.m_children.isEmpty());
+    }
+
+    @Test
+    public void testRemoveTreeNodesAfterEmptySubscriptions_multiLayer() {
+        TreeNode root = new TreeNode();
+        TreeNode intermediateEmptyNode = new TreeNode();
+        TreeNode subNode1 = new TreeNode();
+
+        subNode1.addSubscription(new ISessionsStore.ClientTopicCouple("Client1", new Topic("/sub/topic1")));
+        intermediateEmptyNode.addChild(subNode1);
+        root.addChild(intermediateEmptyNode);
+
+        TreeNode updatedRoot = root.removeClientSubscriptions("Client1");
+
+        assertTrue("Empty subscriptions node should be removed", updatedRoot.m_children.isEmpty());
+    }
+
+    @Test
+    public void testRemoveTreeNodesAfterEmptySubscriptions_multiLayer_doesntRemoveIntermediateNodes() {
+        TreeNode root = new TreeNode();
+        TreeNode intermediateEmptyNode = new TreeNode();
+
+        TreeNode subNode1 = new TreeNode();
+        subNode1.addSubscription(new ISessionsStore.ClientTopicCouple("Client1", new Topic("/sub/topic1")));
+        intermediateEmptyNode.addChild(subNode1);
+
+        TreeNode subNode2 = new TreeNode();
+        subNode2.addSubscription(new ISessionsStore.ClientTopicCouple("Client2", new Topic("/sub/topic2")));
+        intermediateEmptyNode.addChild(subNode2);
+        root.addChild(intermediateEmptyNode);
+
+        TreeNode updatedRoot = root.removeClientSubscriptions("Client1");
+
+        assertFalse("Remove a subscription shouldn't remove intermediate empty nodes", updatedRoot.m_children.isEmpty());
+    }
+
+}