Support ACKNOWLEDGEMENT_MODE for JmsIO

* CheckpointMark behavior in alignment with different ACKNOWLEDGEMENT_MODE

* Ref count active checkpoint for quicker onClose that releases session

* Fix hanging checkpoint when no incoming data in direct runner. This
  allows us to do an exact assert

* Optimize long running unit test usign a short retry

* Re-enable AMQP integration test after stuck unack messages resolved

Author: Abacn

runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java

@@ -173,16 +173,17 @@ public void processElement(
         } else {
           Instant watermark = reader.getWatermark();
           if (watermark.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
-            // If the reader had no elements available, but the shard is not done, reuse it later
-            // Might be better to finalize old checkpoint.
+            // If the reader had no elements available, but the shard is not done, reuse it later.
+            // Finalize old checkpoint now.
+            final CheckpointMarkT checkpoint = shard.getCheckpoint();
+            if (checkpoint != null) {
+              checkpoint.finalizeCheckpoint();
+            }
             resultBuilder.addUnprocessedElements(
                 Collections.<WindowedValue<?>>singleton(
                     WindowedValues.timestampedValueInGlobalWindow(
                         UnboundedSourceShard.of(
-                            shard.getSource(),
-                            shard.getDeduplicator(),
-                            reader,
-                            shard.getCheckpoint()),
+                            shard.getSource(), shard.getDeduplicator(), reader, null),
                         watermark)));
           } else {
             // End of input. Close the reader after finalizing old checkpoint.

sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java

@@ -19,13 +19,17 @@
 
 import java.io.IOException;
 import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import javax.jms.JMSException;
 import javax.jms.Message;
 import javax.jms.MessageConsumer;
 import javax.jms.Session;
 import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.io.jms.JmsIO.AcknowledgeMode;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
@@ -41,29 +45,32 @@ class JmsCheckpointMark implements UnboundedSource.CheckpointMark, Serializable
   private static final Logger LOG = LoggerFactory.getLogger(JmsCheckpointMark.class);
 
   private Instant oldestMessageTimestamp;
-  private transient @Nullable Message lastMessage;
+  private transient @Nullable List<Message> messages;
   private transient @Nullable MessageConsumer consumer;
   private transient @Nullable Session session;
+  private transient @Nullable AtomicInteger activeCheckpoints;
 
   private JmsCheckpointMark(
       Instant oldestMessageTimestamp,
-      @Nullable Message lastMessage,
+      @Nullable List<Message> messages,
       @Nullable MessageConsumer consumer,
-      @Nullable Session session) {
+      @Nullable Session session,
+      @Nullable AtomicInteger activeCheckpoints) {
     this.oldestMessageTimestamp = oldestMessageTimestamp;
-    this.lastMessage = lastMessage;
+    this.messages = messages;
     this.consumer = consumer;
     this.session = session;
+    this.activeCheckpoints = activeCheckpoints;
   }
 
   /** Acknowledge all outstanding message. */
   @Override
   public void finalizeCheckpoint() {
     try {
-      // Jms spec will implicitly acknowledge _all_ messaged already received by the same
-      // session if one message in this session is being acknowledged.
-      if (lastMessage != null) {
-        lastMessage.acknowledge();
+      if (messages != null) {
+        for (Message message : messages) {
+          message.acknowledge();
+        }
       }
     } catch (JMSException e) {
       // The effect of this is message not get acknowledged and thus will be redelivered. It is
@@ -93,14 +100,37 @@ public void finalizeCheckpoint() {
         LOG.info("Error closing JMS session. It may have already been closed.");
       }
     }
+
+    if (activeCheckpoints != null) {
+      activeCheckpoints.decrementAndGet();
+    }
+  }
+
+  @VisibleForTesting
+  @Nullable
+  List<Message> getMessages() {
+    return messages;
+  }
+
+  @VisibleForTesting
+  @Nullable
+  Session getSession() {
+    return session;
+  }
+
+  @VisibleForTesting
+  @Nullable
+  MessageConsumer getConsumer() {
+    return consumer;
   }
 
   // set an empty list to messages when deserialize
   private void readObject(java.io.ObjectInputStream stream)
       throws IOException, ClassNotFoundException {
     stream.defaultReadObject();
-    lastMessage = null;
+    messages = null;
     session = null;
+    consumer = null;
   }
 
   @Override
@@ -120,24 +150,27 @@ public int hashCode() {
     return Objects.hash(oldestMessageTimestamp);
   }
 
-  static Preparer newPreparer() {
-    return new Preparer();
+  static Preparer newPreparer(AcknowledgeMode acknowledgeMode) {
+    return new Preparer(acknowledgeMode);
   }
 
   /**
    * A class preparing the immutable checkpoint. It is mutable so that new messages can be added.
    */
   static class Preparer {
     private Instant oldestMessageTimestamp = Instant.now();
-    private transient @Nullable Message lastMessage = null;
+    private transient List<Message> messages = new ArrayList<>();
+    private final AcknowledgeMode acknowledgeMode;
 
     @VisibleForTesting transient boolean discarded = false;
 
     @VisibleForTesting final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
 
-    private Preparer() {}
+    private Preparer(AcknowledgeMode acknowledgeMode) {
+      this.acknowledgeMode = acknowledgeMode;
+    }
 
-    void add(Message message) throws Exception {
+    void add(Message message) throws JMSException {
       lock.writeLock().lock();
       try {
         if (discarded) {
@@ -149,7 +182,15 @@ void add(Message message) throws Exception {
         if (currentMessageTimestamp.isBefore(oldestMessageTimestamp)) {
           oldestMessageTimestamp = currentMessageTimestamp;
         }
-        lastMessage = message;
+        if (acknowledgeMode == AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE) {
+          messages.add(message);
+        } else {
+          if (messages.isEmpty()) {
+            messages.add(message);
+          } else {
+            messages.set(0, message);
+          }
+        }
       } finally {
         lock.writeLock().unlock();
       }
@@ -167,6 +208,7 @@ Instant getOldestMessageTimestamp() {
     void discard() {
       lock.writeLock().lock();
       try {
+        messages.clear();
         this.discarded = true;
       } finally {
         lock.writeLock().unlock();
@@ -175,20 +217,39 @@ void discard() {
 
     /**
      * Create a new checkpoint mark based on the current preparer. This will reset the messages held
-     * by the preparer, and the owner of the preparer is responsible to create a new Jms session
-     * after this call.
+     * by the preparer. If AcknowledgeMode is CLIENT_ACKNOWLEDGE, the owner of the preparer is
+     * responsible to create a new Jms session after this call.
      */
-    JmsCheckpointMark newCheckpoint(@Nullable MessageConsumer consumer, @Nullable Session session) {
+    JmsCheckpointMark newCheckpoint(
+        @Nullable MessageConsumer consumer,
+        @Nullable Session session,
+        @Nullable AcknowledgeMode acknowledgeMode,
+        @Nullable AtomicInteger activeCheckpoints) {
       JmsCheckpointMark checkpointMark;
       lock.writeLock().lock();
       try {
         if (discarded) {
-          lastMessage = null;
+          messages.clear();
           checkpointMark = this.emptyCheckpoint();
         } else {
+          List<Message> messagesCopy = null;
+          MessageConsumer consumerToPass = null;
+          Session sessionToPass = null;
+          if (!messages.isEmpty()) {
+            messagesCopy = new ArrayList<>(messages);
+          }
+          if (acknowledgeMode == AcknowledgeMode.CLIENT_ACKNOWLEDGE) {
+            consumerToPass = consumer;
+            sessionToPass = session;
+          }
           checkpointMark =
-              new JmsCheckpointMark(oldestMessageTimestamp, lastMessage, consumer, session);
-          lastMessage = null;
+              new JmsCheckpointMark(
+                  oldestMessageTimestamp,
+                  messagesCopy,
+                  consumerToPass,
+                  sessionToPass,
+                  activeCheckpoints);
+          messages.clear();
           oldestMessageTimestamp = Instant.now();
         }
       } finally {
@@ -198,11 +259,11 @@ JmsCheckpointMark newCheckpoint(@Nullable MessageConsumer consumer, @Nullable Se
     }
 
     JmsCheckpointMark emptyCheckpoint() {
-      return new JmsCheckpointMark(oldestMessageTimestamp, null, null, null);
+      return new JmsCheckpointMark(oldestMessageTimestamp, null, null, null, null);
     }
 
     boolean isEmpty() {
-      return lastMessage == null;
+      return messages.isEmpty();
     }
   }
 }