Fixed drop messages in clean session phase, introduced ref counter

broker/src/main/java/io/moquette/spi/IMessagesStore.java

@@ -36,6 +36,7 @@ public interface IMessagesStore {
         //Optional attribute, available only fo QoS 1 and 2
         private Integer m_msgID;
         private String m_guid;
+        private int referenceCounter = 0;
 
         public StoredMessage(byte[] message, AbstractMessage.QOSType qos, String topic) {
             m_qos = qos;
@@ -91,6 +92,18 @@ public interface IMessagesStore {
             return m_retained;
         }
 
+        public void incReferenceCounter() {
+            this.referenceCounter++;
+        }
+
+        public void decReferenceCounter() {
+            this.referenceCounter--;
+        }
+
+        public int getReferenceCounter() {
+            return this.referenceCounter;
+        }
+
         @Override
         public String toString() {
             return "PublishEvent{" +
@@ -137,4 +150,8 @@ public interface IMessagesStore {
     StoredMessage getMessageByGuid(String guid);
 
     void cleanRetained(String topic);
+
+    void incUsageCounter(String guid);
+
+    void decUsageCounter(String guid);
 }

broker/src/main/java/io/moquette/spi/impl/ProtocolProcessor.java

@@ -400,7 +400,7 @@ public class ProtocolProcessor {
             LOG.info("server {} replying with PubAck to MSG ID {}", m_server_port, messageID);
         } else if (qos == AbstractMessage.QOSType.EXACTLY_ONCE) { //QoS2
             guid = m_messagesStore.storePublishForFuture(toStoreMsg);
-            if(msg.isLocal()) {
+            if (msg.isLocal()) {
                 sendPubRec(clientID, messageID);
             }
             //Next the client will send us a pub rel

broker/src/main/java/io/moquette/spi/persistence/MapDBMessagesStore.java

@@ -99,8 +99,16 @@ class MapDBMessagesStore implements IMessagesStore {
 
     @Override
     public void dropMessagesInSession(String clientID) {
-        m_db.getHashMap(MapDBSessionsStore.messageId2GuidsMapName(clientID)).clear();
-        m_persistentMessageStore.remove(clientID);
+        ConcurrentMap<Integer, String> messageIdToGuid = m_db.getHashMap(MapDBSessionsStore.messageId2GuidsMapName(clientID));
+        for (String guid : messageIdToGuid.values()) {
+            //remove only the not retained and no more referenced
+            IMessagesStore.StoredMessage storedMessage = m_persistentMessageStore.get(guid);
+            if (!storedMessage.isRetained() && storedMessage.getReferenceCounter() == 0) {
+                LOG.debug("Cleaning not retained message guid {}", guid);
+                m_persistentMessageStore.remove(guid);
+            }
+        }
+        messageIdToGuid.clear();
     }
 
     @Override
@@ -112,4 +120,18 @@ class MapDBMessagesStore implements IMessagesStore {
     public void cleanRetained(String topic) {
         m_retainedStore.remove(topic);
     }
+
+    @Override
+    public void incUsageCounter(String guid) {
+        IMessagesStore.StoredMessage storedMessage = m_persistentMessageStore.get(guid);
+        storedMessage.incReferenceCounter();
+        m_persistentMessageStore.replace(guid, storedMessage);
+    }
+
+    @Override
+    public void decUsageCounter(String guid) {
+        IMessagesStore.StoredMessage storedMessage = m_persistentMessageStore.get(guid);
+        storedMessage.decReferenceCounter();
+        m_persistentMessageStore.replace(guid, storedMessage);
+    }
 }

broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java

@@ -203,6 +203,7 @@ class MapDBSessionsStore implements ISessionsStore {
         List<String> guids = Utils.defaultGet(m_enqueuedStore, clientID, new ArrayList<String>());
         guids.add(guid);
         m_enqueuedStore.put(clientID, guids);
+        m_messagesStore.incUsageCounter(guid);
     }
 
     @Override
@@ -215,6 +216,8 @@ class MapDBSessionsStore implements ISessionsStore {
         List<String> guids = Utils.defaultGet(m_enqueuedStore, clientID, new ArrayList<String>());
         guids.remove(guid);
         m_enqueuedStore.put(clientID, guids);
+        m_messagesStore.decUsageCounter(guid);
+        //TODO if counter gets to 0 then remove from storage
     }
 
     @Override

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

@@ -100,4 +100,18 @@ public class MemoryMessagesStore implements IMessagesStore {
     public void cleanRetained(String topic) {
         m_retainedStore.remove(topic);
     }
+
+    @Override
+    public void incUsageCounter(String guid) {
+        IMessagesStore.StoredMessage storedMessage = m_persistentMessageStore.get(guid);
+        storedMessage.incReferenceCounter();
+        m_persistentMessageStore.replace(guid, storedMessage);
+    }
+
+    @Override
+    public void decUsageCounter(String guid) {
+        IMessagesStore.StoredMessage storedMessage = m_persistentMessageStore.get(guid);
+        storedMessage.decReferenceCounter();
+        m_persistentMessageStore.replace(guid, storedMessage);
+    }
 }

broker/src/test/java/io/moquette/spi/persistence/MapDBPersistentStoreTest.java

@@ -41,6 +41,7 @@ import static org.junit.Assert.*;
  */
 public class MapDBPersistentStoreTest {
 
+    public static final String TEST_CLIENT = "TestClient";
     MapDBPersistentStore m_storageService;
     ISessionsStore m_sessionsStore;
     IMessagesStore m_messagesStore;
@@ -75,14 +76,14 @@ public class MapDBPersistentStoreTest {
         session1.subscribe(oldSubscription);
 
         // Subscribe on /topic again that overrides the previous subscription.
-        Subscription overrindingSubscription = new Subscription(session1.clientID, "/topic", AbstractMessage.QOSType.EXACTLY_ONCE);
-        session1.subscribe(overrindingSubscription);
+        Subscription overridingSubscription = new Subscription(session1.clientID, "/topic", AbstractMessage.QOSType.EXACTLY_ONCE);
+        session1.subscribe(overridingSubscription);
 
         //Verify
         List<ClientTopicCouple> subscriptions = m_sessionsStore.listAllSubscriptions();
         assertEquals(1, subscriptions.size());
         Subscription sub = m_sessionsStore.getSubscription(subscriptions.get(0));
-        assertEquals(overrindingSubscription.getRequestedQos(), sub.getRequestedQos());
+        assertEquals(overridingSubscription.getRequestedQos(), sub.getRequestedQos());
     }
 
     @Test
@@ -127,4 +128,40 @@ public class MapDBPersistentStoreTest {
                 m_storageService.m_scheduler.awaitTermination(3, TimeUnit.SECONDS));
         assertTrue(m_storageService.m_scheduler.isTerminated());
     }
+
+    @Test
+    public void testDropMessagesInSessionCleanAllNotRetainedStoredMessages() {
+        m_sessionsStore.createNewSession("TestClient", true);
+        IMessagesStore.StoredMessage publishToStore = new IMessagesStore.StoredMessage("Hello".getBytes(),
+                AbstractMessage.QOSType.EXACTLY_ONCE, "/topic");
+        publishToStore.setClientID(TEST_CLIENT);
+        publishToStore.setMessageID(1);
+        publishToStore.setRetained(false);
+        String guid = m_messagesStore.storePublishForFuture(publishToStore);
+
+        //Exercise
+        m_messagesStore.dropMessagesInSession("TestClient");
+
+        //Verify the message store for session is empty.
+        IMessagesStore.StoredMessage storedPublish = m_messagesStore.getMessageByGuid(guid);
+        assertNull("The stored message must'n be present anymore", storedPublish);
+    }
+
+    @Test
+    public void testDropMessagesInSessionDoesntCleanAnyRetainedStoredMessages() {
+        m_sessionsStore.createNewSession("TestClient", true);
+        IMessagesStore.StoredMessage publishToStore = new IMessagesStore.StoredMessage("Hello".getBytes(),
+                AbstractMessage.QOSType.EXACTLY_ONCE, "/topic");
+        publishToStore.setClientID(TEST_CLIENT);
+        publishToStore.setMessageID(1);
+        publishToStore.setRetained(true);
+        String guid = m_messagesStore.storePublishForFuture(publishToStore);
+
+        //Exercise
+        m_messagesStore.dropMessagesInSession("TestClient");
+
+        //Verify the message store for session is empty.
+        IMessagesStore.StoredMessage storedPublish = m_messagesStore.getMessageByGuid(guid);
+        assertNotNull("The stored retained message must be present after client's session drop", storedPublish);
+    }
 }