From 4d637a349b78497243a51c7f312a8cd184588165 Mon Sep 17 00:00:00 2001 From: andsel Date: Fri, 23 Sep 2016 09:45:47 +0200 Subject: [PATCH] Fixed drop messages in clean session phase, introduced ref counter --- .../java/io/moquette/spi/IMessagesStore.java | 17 ++++++++ .../moquette/spi/impl/ProtocolProcessor.java | 2 +- .../spi/persistence/MapDBMessagesStore.java | 26 ++++++++++- .../spi/persistence/MapDBSessionsStore.java | 3 ++ .../spi/impl/MemoryMessagesStore.java | 14 ++++++ .../persistence/MapDBPersistentStoreTest.java | 43 +++++++++++++++++-- 6 files changed, 99 insertions(+), 6 deletions(-) diff --git a/broker/src/main/java/io/moquette/spi/IMessagesStore.java b/broker/src/main/java/io/moquette/spi/IMessagesStore.java index 6845c6b9..f2ad8ce1 100644 --- a/broker/src/main/java/io/moquette/spi/IMessagesStore.java +++ b/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); } diff --git a/broker/src/main/java/io/moquette/spi/impl/ProtocolProcessor.java b/broker/src/main/java/io/moquette/spi/impl/ProtocolProcessor.java index 3aa32403..53592df0 100644 --- a/broker/src/main/java/io/moquette/spi/impl/ProtocolProcessor.java +++ b/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 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 716ffcaa..67adac89 100644 --- a/broker/src/main/java/io/moquette/spi/persistence/MapDBMessagesStore.java +++ b/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 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); + } } 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 0d7ead81..825fac70 100644 --- a/broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java +++ b/broker/src/main/java/io/moquette/spi/persistence/MapDBSessionsStore.java @@ -203,6 +203,7 @@ class MapDBSessionsStore implements ISessionsStore { List guids = Utils.defaultGet(m_enqueuedStore, clientID, new ArrayList()); guids.add(guid); m_enqueuedStore.put(clientID, guids); + m_messagesStore.incUsageCounter(guid); } @Override @@ -215,6 +216,8 @@ class MapDBSessionsStore implements ISessionsStore { List guids = Utils.defaultGet(m_enqueuedStore, clientID, new ArrayList()); guids.remove(guid); m_enqueuedStore.put(clientID, guids); + m_messagesStore.decUsageCounter(guid); + //TODO if counter gets to 0 then remove from storage } @Override diff --git a/broker/src/test/java/io/moquette/spi/impl/MemoryMessagesStore.java b/broker/src/test/java/io/moquette/spi/impl/MemoryMessagesStore.java index 5c4eed04..685f62eb 100644 --- a/broker/src/test/java/io/moquette/spi/impl/MemoryMessagesStore.java +++ b/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); + } } diff --git a/broker/src/test/java/io/moquette/spi/persistence/MapDBPersistentStoreTest.java b/broker/src/test/java/io/moquette/spi/persistence/MapDBPersistentStoreTest.java index 77d55062..03534395 100644 --- a/broker/src/test/java/io/moquette/spi/persistence/MapDBPersistentStoreTest.java +++ b/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 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); + } } \ No newline at end of file