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