ARTEMIS-1118 IO callbacks on AMQP

artemis-commons/src/main/java/org/apache/activemq/artemis/utils/RunnableEx.java

@@ -0,0 +1,22 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.activemq.artemis.utils;
+
+public interface RunnableEx {
+   void run() throws Exception;
+}

artemis-protocols/artemis-amqp-protocol/src/main/java/org/apache/activemq/artemis/protocol/amqp/broker/AMQPSessionCallback.java

@@ -30,6 +30,7 @@ import org.apache.activemq.artemis.api.core.client.ActiveMQClient;
 import org.apache.activemq.artemis.core.io.IOCallback;
 import org.apache.activemq.artemis.core.paging.PagingStore;
 import org.apache.activemq.artemis.core.persistence.OperationContext;
+import org.apache.activemq.artemis.core.persistence.StorageManager;
 import org.apache.activemq.artemis.core.server.AddressQueryResult;
 import org.apache.activemq.artemis.core.server.BindingQueryResult;
 import org.apache.activemq.artemis.core.server.MessageReference;
@@ -53,6 +54,7 @@ import org.apache.activemq.artemis.spi.core.protocol.SessionCallback;
 import org.apache.activemq.artemis.spi.core.remoting.Connection;
 import org.apache.activemq.artemis.spi.core.remoting.ReadyListener;
 import org.apache.activemq.artemis.utils.IDGenerator;
+import org.apache.activemq.artemis.utils.RunnableEx;
 import org.apache.activemq.artemis.utils.SelectorTranslator;
 import org.apache.activemq.artemis.utils.SimpleIDGenerator;
 import org.apache.activemq.artemis.utils.UUIDGenerator;
@@ -78,6 +80,8 @@ public class AMQPSessionCallback implements SessionCallback {
 
    private final ProtonProtocolManager manager;
 
+   private final StorageManager storageManager;
+
    private final AMQPConnectionContext connection;
 
    private final Connection transportConnection;
@@ -100,6 +104,7 @@ public class AMQPSessionCallback implements SessionCallback {
                               OperationContext operationContext) {
       this.protonSPI = protonSPI;
       this.manager = manager;
+      this.storageManager = manager.getServer().getStorageManager();
       this.connection = connection;
       this.transportConnection = transportConnection;
       this.closeExecutor = executor;
@@ -134,6 +139,24 @@ public class AMQPSessionCallback implements SessionCallback {
       }
    }
 
+   public void withinContext(RunnableEx run) throws Exception {
+      OperationContext context = recoverContext();
+      try {
+         run.run();
+      } finally {
+         resetContext(context);
+      }
+   }
+
+   public void afterIO(IOCallback ioCallback) {
+      OperationContext context = recoverContext();
+      try {
+         manager.getServer().getStorageManager().afterCompleteOperations(ioCallback);
+      } finally {
+         resetContext(context);
+      }
+   }
+
    @Override
    public void browserFinished(ServerConsumer consumer) {
 
@@ -315,11 +338,11 @@ public class AMQPSessionCallback implements SessionCallback {
    public void close() throws Exception {
       //need to check here as this can be called if init fails
       if (serverSession != null) {
-         recoverContext();
+         OperationContext context = recoverContext();
          try {
             serverSession.close(false);
          } finally {
-            resetContext();
+            resetContext(context);
          }
       }
    }
@@ -328,30 +351,30 @@ public class AMQPSessionCallback implements SessionCallback {
       if (transaction == null) {
          transaction = serverSession.getCurrentTransaction();
       }
-      recoverContext();
+      OperationContext oldContext = recoverContext();
       try {
          ((ServerConsumer) brokerConsumer).individualAcknowledge(transaction, message.getMessageID());
       } finally {
-         resetContext();
+         resetContext(oldContext);
       }
    }
 
    public void cancel(Object brokerConsumer, Message message, boolean updateCounts) throws Exception {
-      recoverContext();
+      OperationContext oldContext = recoverContext();
       try {
          ((ServerConsumer) brokerConsumer).individualCancel(message.getMessageID(), updateCounts);
          ((ServerConsumer) brokerConsumer).getQueue().forceDelivery();
       } finally {
-         resetContext();
+         resetContext(oldContext);
       }
    }
 
    public void reject(Object brokerConsumer, Message message) throws Exception {
-      recoverContext();
+      OperationContext oldContext = recoverContext();
       try {
          ((ServerConsumer) brokerConsumer).reject(message.getMessageID());
       } finally {
-         resetContext();
+         resetContext(oldContext);
       }
    }
 
@@ -380,8 +403,9 @@ public class AMQPSessionCallback implements SessionCallback {
          }
       }
 
-      recoverContext();
+      OperationContext oldcontext = recoverContext();
 
+      try {
          PagingStore store = manager.getServer().getPagingManager().getPageStore(message.getAddressSimpleString());
          if (store.isRejectingMessages()) {
             // We drop pre-settled messages (and abort any associated Tx)
@@ -397,6 +421,9 @@ public class AMQPSessionCallback implements SessionCallback {
          } else {
             serverSend(transaction, message, delivery, receiver);
          }
+      } finally {
+         resetContext(oldcontext);
+      }
    }
 
    private void rejectMessage(Delivery delivery, Symbol errorCondition, String errorMessage) {
@@ -406,6 +433,9 @@ public class AMQPSessionCallback implements SessionCallback {
       Rejected rejected = new Rejected();
       rejected.setError(condition);
 
+      afterIO(new IOCallback() {
+         @Override
+         public void done() {
             connection.lock();
             try {
                delivery.disposition(rejected);
@@ -416,17 +446,23 @@ public class AMQPSessionCallback implements SessionCallback {
             connection.flush();
          }
 
+         @Override
+         public void onError(int errorCode, String errorMessage) {
+
+         }
+      });
+
+   }
+
    private void serverSend(final Transaction transaction,
                            final Message message,
                            final Delivery delivery,
                            final Receiver receiver) throws Exception {
-      try {
-
       message.setConnectionID(receiver.getSession().getConnection().getRemoteContainer());
 
       serverSession.send(transaction, message, false, false);
 
-         manager.getServer().getStorageManager().afterCompleteOperations(new IOCallback() {
+      afterIO(new IOCallback() {
          @Override
          public void done() {
             connection.lock();
@@ -458,9 +494,6 @@ public class AMQPSessionCallback implements SessionCallback {
             }
          }
       });
-      } finally {
-         resetContext();
-      }
    }
 
    public void offerProducerCredit(final String address,
@@ -502,12 +535,15 @@ public class AMQPSessionCallback implements SessionCallback {
       manager.getServer().destroyQueue(new SimpleString(queueName));
    }
 
-   private void resetContext() {
+   public void resetContext(OperationContext oldContext) {
-      manager.getServer().getStorageManager().setContext(null);
+      storageManager.setContext(oldContext);
    }
 
-   private void recoverContext() {
+   public OperationContext recoverContext() {
+
+      OperationContext oldContext = storageManager.getContext();
       manager.getServer().getStorageManager().setContext(serverSession.getSessionContext());
+      return oldContext;
    }
 
    @Override

artemis-protocols/artemis-amqp-protocol/src/main/java/org/apache/activemq/artemis/protocol/amqp/client/AMQPClientConnectionFactory.java

@@ -52,7 +52,7 @@ public class AMQPClientConnectionFactory {
 
       Executor executor = server.getExecutorFactory().getExecutor();
 
-      AMQPConnectionContext amqpConnection = new AMQPConnectionContext(null, connectionCallback, containerId, ttl, protocolManager.getMaxFrameSize(), AMQPConstants.Connection.DEFAULT_CHANNEL_MAX, server.getScheduledPool());
+      AMQPConnectionContext amqpConnection = new AMQPConnectionContext(protocolManager, connectionCallback, containerId, ttl, protocolManager.getMaxFrameSize(), AMQPConstants.Connection.DEFAULT_CHANNEL_MAX, server.getScheduledPool());
       eventHandler.ifPresent(amqpConnection::addEventHandler);
 
       ActiveMQProtonRemotingConnection delegate = new ActiveMQProtonRemotingConnection(protocolManager, amqpConnection, connection, executor);

artemis-protocols/artemis-amqp-protocol/src/main/java/org/apache/activemq/artemis/protocol/amqp/proton/ProtonServerSenderContext.java

@@ -29,6 +29,8 @@ import org.apache.activemq.artemis.api.core.ActiveMQQueueExistsException;
 import org.apache.activemq.artemis.api.core.Message;
 import org.apache.activemq.artemis.api.core.RoutingType;
 import org.apache.activemq.artemis.api.core.SimpleString;
+import org.apache.activemq.artemis.core.io.IOCallback;
+import org.apache.activemq.artemis.core.persistence.OperationContext;
 import org.apache.activemq.artemis.core.postoffice.impl.CompositeAddress;
 import org.apache.activemq.artemis.core.server.AddressQueryResult;
 import org.apache.activemq.artemis.core.server.Consumer;
@@ -486,6 +488,8 @@ public class ProtonServerSenderContext extends ProtonInitializable implements Pr
          return;
       }
 
+      OperationContext oldContext = sessionSPI.recoverContext();
+
       try {
          Message message = ((MessageReference) delivery.getContext()).getMessage();
 
@@ -590,8 +594,20 @@ public class ProtonServerSenderContext extends ProtonInitializable implements Pr
             // todo not sure if we need to do anything here
          }
       } finally {
+         sessionSPI.afterIO(new IOCallback() {
+            @Override
+            public void done() {
                connection.flush();
             }
+
+            @Override
+            public void onError(int errorCode, String errorMessage) {
+               connection.flush();
+            }
+         });
+
+         sessionSPI.resetContext(oldContext);
+      }
    }
 
    public void settle(Delivery delivery) {

artemis-protocols/artemis-amqp-protocol/src/main/java/org/apache/activemq/artemis/protocol/amqp/proton/transaction/ProtonTransactionHandler.java

@@ -18,6 +18,7 @@ package org.apache.activemq.artemis.protocol.amqp.proton.transaction;
 
 import java.nio.ByteBuffer;
 
+import org.apache.activemq.artemis.core.io.IOCallback;
 import org.apache.activemq.artemis.protocol.amqp.broker.AMQPSessionCallback;
 import org.apache.activemq.artemis.protocol.amqp.exceptions.ActiveMQAMQPException;
 import org.apache.activemq.artemis.protocol.amqp.proton.AMQPConnectionContext;
@@ -118,24 +119,29 @@ public class ProtonTransactionHandler implements ProtonDeliveryHandler {
             ProtonTransactionImpl tx = (ProtonTransactionImpl) sessionSPI.getTransaction(txID, true);
             tx.discharge();
 
+            IOCallback ioAction = new IOCallback() {
+               @Override
+               public void done() {
+                  connection.lock();
+                  try {
+                     delivery.disposition(new Accepted());
+                  } finally {
+                     connection.unlock();
+                  }
+               }
+
+               @Override
+               public void onError(int errorCode, String errorMessage) {
+
+               }
+            };
+
             if (discharge.getFail()) {
-               tx.rollback();
+               sessionSPI.withinContext(() -> tx.rollback());
-               connection.lock();
+               sessionSPI.afterIO(ioAction);
-               try {
-                  delivery.disposition(new Accepted());
-               } finally {
-                  connection.unlock();
-               }
-               connection.flush();
             } else {
-               tx.commit();
+               sessionSPI.withinContext(() -> tx.commit());
-               connection.lock();
+               sessionSPI.afterIO(ioAction);
-               try {
-                  delivery.disposition(new Accepted());
-               } finally {
-                  connection.unlock();
-               }
-               connection.flush();
             }
          }
       } catch (ActiveMQAMQPException amqpE) {
@@ -157,6 +163,9 @@ public class ProtonTransactionHandler implements ProtonDeliveryHandler {
          }
          connection.flush();
       } finally {
+         sessionSPI.afterIO(new IOCallback() {
+            @Override
+            public void done() {
                connection.lock();
                try {
                   delivery.settle();
@@ -165,6 +174,13 @@ public class ProtonTransactionHandler implements ProtonDeliveryHandler {
                }
                connection.flush();
             }
+
+            @Override
+            public void onError(int errorCode, String errorMessage) {
+
+            }
+         });
+      }
    }
 
    @Override