ARTEMIS-6284 make MQTT cache removal transactional

artemis-protocols/artemis-mqtt-protocol/src/main/java/org/apache/activemq/artemis/core/protocol/mqtt/MQTTPublishManager.java

@@ -325,17 +325,28 @@ private void releaseFlowControl(Long consumerId) {
    }
 
    synchronized void handlePubComp(int packetId) throws Exception {
-      session.getState().getPubRecCache().remove(packetId);
+      removePacketIdFromCache(session.getState().getPubRecCache(), packetId, MqttMessageType.PUBCOMP, MqttMessageType.PUBREC);
    }
 
    synchronized void handlePubRel(int packetId) throws Exception {
-      boolean deleted = session.getState().getPublishCache().remove(packetId);
-      if (!deleted) {
-         logger.debug("MQTT client {} sent PUBREL for packet {} but no corresponding PUBLISH was found in the cache", session.getState().getClientId(), packetId);
-      }
+      removePacketIdFromCache(session.getState().getPublishCache(), packetId, MqttMessageType.PUBREL, MqttMessageType.PUBLISH);
       session.getProtocolHandler().sendPubComp(packetId);
    }
 
+   private void removePacketIdFromCache(PacketIdCache cache, int packetId, MqttMessageType receivedType, MqttMessageType cachedType) throws Exception {
+      Transaction tx = session.getServerSession().newTransaction();
+      try {
+         boolean deleted = cache.remove(packetId, tx);
+         if (!deleted) {
+            logger.debug("MQTT client {} sent {} for packet {} but no corresponding {} was found in the cache", session.getState().getClientId(), receivedType, packetId, cachedType);
+         }
+         tx.commit();
+      } catch (Throwable t) {
+         tx.rollback();
+         throw t;
+      }
+   }
+
    synchronized void handlePubAck(int packetId) throws Exception {
       acknowledgeDelivery(packetId, false);
    }
@@ -377,6 +388,7 @@ private void acknowledgeDelivery(int packetId, boolean needsPubRel) throws Excep
                state.getPubRecCache().add(packetId, tx);
             }
             session.getStateManager().removePacketIdCorrelation(state.getClientId(), delivery.getPacketIdCorrelationKey(), tx.getID());
+            tx.setContainsPersistent();
             consumer.individualAcknowledge(tx, delivery.getCoreMessageId());
             tx.commit();
             state.removeCoreDeliveryInfo(packetId);

artemis-protocols/artemis-mqtt-protocol/src/main/java/org/apache/activemq/artemis/core/protocol/mqtt/PacketIdCache.java

@@ -72,9 +72,9 @@ public boolean contains(int packetId) {
       return cache != null && cache.contains(ByteUtil.intToBytes(packetId));
    }
 
-   public boolean remove(int packetId) throws Exception {
+   public boolean remove(int packetId, Transaction tx) throws Exception {
       check();
-      return cache != null && cache.deleteFromCache(ByteUtil.intToBytes(packetId));
+      return cache != null && cache.deleteFromCache(ByteUtil.intToBytes(packetId), tx);
    }
 
    public int size() {

artemis-server/src/main/java/org/apache/activemq/artemis/core/postoffice/DuplicateIDCache.java

@@ -41,6 +41,8 @@ public interface DuplicateIDCache {
 
    boolean deleteFromCache(byte[] duplicateID) throws Exception;
 
+   boolean deleteFromCache(byte[] duplicateID, Transaction tx) throws Exception;
+
    void load(List<Pair<byte[], Long>> ids) throws Exception;
 
    void load(Transaction tx, byte[] duplID);

artemis-server/src/main/java/org/apache/activemq/artemis/core/postoffice/impl/InMemoryDuplicateIDCache.java

@@ -81,6 +81,11 @@ public boolean deleteFromCache(byte[] duplicateID) {
       return deleteFromCache(new ByteArray(duplicateID));
    }
 
+   @Override
+   public boolean deleteFromCache(byte[] duplicateID, Transaction tx) {
+      return deleteFromCache(new ByteArray(duplicateID));
+   }
+
    private boolean deleteFromCache(final ByteArray duplicateID) {
       if (logger.isTraceEnabled()) {
          logger.trace("deleting id = {}", describeID(duplicateID.bytes));

artemis-server/src/main/java/org/apache/activemq/artemis/core/postoffice/impl/NoOpDuplicateIDCache.java

@@ -55,6 +55,11 @@ public boolean deleteFromCache(byte[] duplicateID) throws Exception {
       return true;
    }
 
+   @Override
+   public boolean deleteFromCache(byte[] duplicateID, Transaction tx) throws Exception {
+      return true;
+   }
+
    @Override
    public void load(List<Pair<byte[], Long>> ids) throws Exception {
 

artemis-server/src/main/java/org/apache/activemq/artemis/core/postoffice/impl/PersistentDuplicateIDCache.java

@@ -146,10 +146,16 @@ public synchronized void load(final List<Pair<byte[], Long>> ids) throws Excepti
 
    @Override
    public boolean deleteFromCache(byte[] duplicateID) throws Exception {
-      return deleteFromCache(new ByteArray(duplicateID));
+      return deleteFromCache(new ByteArray(duplicateID), null);
    }
 
-   private boolean deleteFromCache(final ByteArray duplicateID) throws Exception {
+   @Override
+   public boolean deleteFromCache(byte[] duplicateID, Transaction tx) throws Exception {
+      Objects.requireNonNull(tx, "tx must be not null");
+      return deleteFromCache(new ByteArray(duplicateID), tx);
+   }
+
+   private boolean deleteFromCache(final ByteArray duplicateID, final Transaction tx) throws Exception {
       if (logger.isTraceEnabled()) {
          logger.trace("deleting id = {}", describeID(duplicateID.bytes));
       }
@@ -167,7 +173,12 @@ private boolean deleteFromCache(final ByteArray duplicateID) throws Exception {
                if (logger.isTraceEnabled()) {
                   logger.trace("address = {} deleting id = {}", address, describeID(duplicateID.bytes, id.getB()));
                }
-               storageManager.deleteDuplicateID(recordID);
+               if (tx == null) {
+                  storageManager.deleteDuplicateID(recordID);
+               } else {
+                  storageManager.deleteDuplicateIDTransactional(tx.getID(), recordID);
+                  tx.setContainsPersistent();
+               }
             }
          }
          return true;
@@ -393,7 +404,7 @@ public void afterCommit(final Transaction tx) {
       @Override
       public void beforeRollback(Transaction tx) throws Exception {
          if (!afterCommit) {
-            deleteFromCache(holder);
+            deleteFromCache(holder, null);
          }
       }