ARTEMIS-2641 Openwire client runs out of credits after reconnection

Clear the messagePullHandler on setting prefetchSize.

artemis-protocols/artemis-openwire-protocol/pom.xml

@@ -105,6 +105,16 @@
          <groupId>org.osgi</groupId>
          <artifactId>osgi.cmpn</artifactId>
       </dependency>
+      <dependency>
+         <groupId>junit</groupId>
+         <artifactId>junit</artifactId>
+         <scope>test</scope>
+      </dependency>
+      <dependency>
+         <groupId>org.mockito</groupId>
+         <artifactId>mockito-core</artifactId>
+         <scope>test</scope>
+      </dependency>
    </dependencies>
 
    <build>

artemis-protocols/artemis-openwire-protocol/src/main/java/org/apache/activemq/artemis/core/protocol/openwire/amq/AMQConsumer.java

@@ -27,6 +27,7 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.activemq.advisory.AdvisorySupport;
 import org.apache.activemq.artemis.api.core.ActiveMQQueueExistsException;
@@ -68,7 +69,7 @@ public class AMQConsumer {
    private int prefetchSize;
    private final AtomicInteger currentWindow;
    private long messagePullSequence = 0;
-   private MessagePullHandler messagePullHandler;
+   private final AtomicReference<MessagePullHandler> messagePullHandler = new AtomicReference<>(null);
    //internal means we don't expose
    //it's address/queue to management service
    private boolean internalAddress = false;
@@ -87,7 +88,7 @@ public class AMQConsumer {
       this.prefetchSize = info.getPrefetchSize();
       this.currentWindow = new AtomicInteger(prefetchSize);
       if (prefetchSize == 0) {
-         messagePullHandler = new MessagePullHandler();
+         messagePullHandler.set(new MessagePullHandler());
       }
       this.internalAddress = internalAddress;
       this.rolledbackMessageRefs = null;
@@ -228,7 +229,7 @@ public class AMQConsumer {
    }
 
    public void acquireCredit(int n) throws Exception {
-      if (messagePullHandler != null) {
+      if (messagePullHandler.get() != null) {
          //don't acquire any credits when the pull handler controls it!!
          return;
       }
@@ -245,7 +246,8 @@ public class AMQConsumer {
    public int handleDeliver(MessageReference reference, ICoreMessage message, int deliveryCount) {
       MessageDispatch dispatch;
       try {
-         if (messagePullHandler != null && !messagePullHandler.checkForcedConsumer(message)) {
+         MessagePullHandler pullHandler = messagePullHandler.get();
+         if (pullHandler != null && !pullHandler.checkForcedConsumer(message)) {
             return 0;
          }
 
@@ -359,8 +361,9 @@ public class AMQConsumer {
 
    public void processMessagePull(MessagePull messagePull) throws Exception {
       currentWindow.incrementAndGet();
-      if (messagePullHandler != null) {
-         messagePullHandler.nextSequence(messagePullSequence++, messagePull.getTimeout());
+      MessagePullHandler pullHandler = messagePullHandler.get();
+      if (pullHandler != null) {
+         pullHandler.nextSequence(messagePullSequence++, messagePull.getTimeout());
       }
    }
 
@@ -380,6 +383,11 @@ public class AMQConsumer {
       this.prefetchSize = prefetchSize;
       this.currentWindow.set(prefetchSize);
       this.info.setPrefetchSize(prefetchSize);
+      if (this.prefetchSize == 0) {
+         messagePullHandler.compareAndSet(null, new MessagePullHandler());
+      } else {
+         messagePullHandler.set(null);
+      }
       if (this.prefetchSize > 0) {
          serverConsumer.promptDelivery();
       }

artemis-protocols/artemis-openwire-protocol/src/test/java/org/apache/activemq/artemis/core/protocol/openwire/amq/AMQConsumerTest.java

@@ -0,0 +1,91 @@
+/**
+ * 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.core.protocol.openwire.amq;
+
+import java.util.concurrent.ScheduledExecutorService;
+
+import org.apache.activemq.artemis.api.core.ICoreMessage;
+import org.apache.activemq.artemis.api.core.SimpleString;
+import org.apache.activemq.artemis.core.message.impl.CoreMessage;
+import org.apache.activemq.artemis.core.protocol.openwire.OpenWireConnection;
+import org.apache.activemq.artemis.core.server.ActiveMQServer;
+import org.apache.activemq.artemis.core.server.MessageReference;
+import org.apache.activemq.artemis.core.server.ServerSession;
+import org.apache.activemq.artemis.core.server.SlowConsumerDetectionListener;
+import org.apache.activemq.artemis.core.server.impl.ServerConsumerImpl;
+import org.apache.activemq.command.ActiveMQTopic;
+import org.apache.activemq.command.ConsumerInfo;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
+
+public class AMQConsumerTest {
+
+   @Test
+   public void testCreditsWithPrefetch() throws Exception {
+      AMQConsumer consumer = getConsumer(1);
+
+      testCredits(consumer);
+   }
+
+   @Test
+   public void testCreditsUpdatingPrefetchSize() throws Exception {
+      AMQConsumer consumer = getConsumer(0);
+
+      consumer.setPrefetchSize(1);
+
+      testCredits(consumer);
+   }
+
+   private AMQConsumer getConsumer(int prefetchSize) throws Exception {
+      ServerSession coreSession = Mockito.mock(ServerSession.class);
+      Mockito.when(coreSession.createConsumer(ArgumentMatchers.anyLong(), ArgumentMatchers.nullable(SimpleString.class),
+                                              ArgumentMatchers.nullable(SimpleString.class), ArgumentMatchers.anyInt(),
+                                              ArgumentMatchers.anyBoolean(), ArgumentMatchers.anyBoolean(),
+                                              ArgumentMatchers.nullable(Integer.class))).thenReturn(Mockito.mock(ServerConsumerImpl.class));
+      AMQSession session = Mockito.mock(AMQSession.class);
+      Mockito.when(session.getConnection()).thenReturn(Mockito.mock(OpenWireConnection.class));
+      Mockito.when(session.getCoreServer()).thenReturn(Mockito.mock(ActiveMQServer.class));
+      Mockito.when(session.getCoreSession()).thenReturn(coreSession);
+      Mockito.when(session.convertWildcard(ArgumentMatchers.any(String.class))).thenReturn("");
+
+      ConsumerInfo info = new ConsumerInfo();
+      info.setPrefetchSize(prefetchSize);
+
+      AMQConsumer consumer = new AMQConsumer(session, new ActiveMQTopic("TEST"), info, Mockito.mock(ScheduledExecutorService.class), false);
+      consumer.init(Mockito.mock(SlowConsumerDetectionListener.class), 0);
+
+      return consumer;
+   }
+
+   private void testCredits(AMQConsumer consumer) throws  Exception {
+      ICoreMessage message = new CoreMessage(1, 0);
+      MessageReference reference = Mockito.mock(MessageReference.class);
+
+      Assert.assertTrue(consumer.hasCredits());
+
+      consumer.handleDeliver(reference, message, 0);
+
+      Assert.assertFalse(consumer.hasCredits());
+
+      consumer.acquireCredit(1);
+
+      Assert.assertTrue(consumer.hasCredits());
+   }
+}