From ac419adbd68ffd4edd8a39ea42a5df08f11fc672 Mon Sep 17 00:00:00 2001
From: andsel <selva.andre@gmail.com>
Date: Fri, 23 Sep 2016 10:16:24 +0200
Subject: [PATCH] Drop the message when the ref counter reaches 0

---
 .../spi/persistence/MapDBMessagesStore.java      | 16 ++++++++++------
 .../spi/persistence/MapDBSessionsStore.java      |  4 ++--
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/broker/src/main/java/io/moquette/spi/persistence/MapDBMessagesStore.java b/broker/src/main/java/io/moquette/spi/persistence/MapDBMessagesStore.java
index 67adac89..9affdc65 100644
--- a/broker/src/main/java/io/moquette/spi/persistence/MapDBMessagesStore.java
+++ b/broker/src/main/java/io/moquette/spi/persistence/MapDBMessagesStore.java
@@ -101,16 +101,20 @@ class MapDBMessagesStore implements IMessagesStore {
     public void dropMessagesInSession(String 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);
-            }
+            removeStoredMessage(guid);
         }
         messageIdToGuid.clear();
     }
 
+    void removeStoredMessage(String guid) {
+        //remove only the not retained and no more referenced
+        StoredMessage storedMessage = m_persistentMessageStore.get(guid);
+        if (!storedMessage.isRetained() && storedMessage.getReferenceCounter() == 0) {
+            LOG.debug("Cleaning not retained message guid {}", guid);
+            m_persistentMessageStore.remove(guid);
+        }
+    }
+
     @Override
     public StoredMessage getMessageByGuid(String guid) {
         return m_persistentMessageStore.get(guid);
diff --git a/broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java b/broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java
index ad6f8bd3..11c4180b 100644
--- a/broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java
+++ b/broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java
@@ -16,7 +16,6 @@
 package io.moquette.spi.persistence;
 
 import io.moquette.spi.ClientSession;
-import io.moquette.spi.IMessagesStore;
 import io.moquette.spi.IMessagesStore.StoredMessage;
 import io.moquette.spi.ISessionsStore;
 import io.moquette.spi.impl.Utils;
@@ -217,7 +216,8 @@ class MapDBSessionsStore implements ISessionsStore {
         guids.remove(guid);
         m_enqueuedStore.put(clientID, guids);
         m_messagesStore.decUsageCounter(guid);
-        //TODO if counter gets to 0 then remove from storage
+        //if counter gets to 0 then remove from storage
+        m_messagesStore.removeStoredMessage(guid);
     }
 
     @Override