[JmsIO] Let restored checkpoint marks acknowledge live messages

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

@@ -26,6 +26,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import org.apache.beam.sdk.io.UnboundedSource;
@@ -44,6 +46,20 @@ class JmsCheckpointMark implements UnboundedSource.CheckpointMark, Serializable
 
   private static final Logger LOG = LoggerFactory.getLogger(JmsCheckpointMark.class);
 
+  /**
+   * Live checkpoints of this JVM that still hold unacknowledged messages, by id. The messages,
+   * consumer and session of a checkpoint are transient, so a runner that serializes and restores
+   * the mark (or recreates the reader) would otherwise lose them and never acknowledge. A restored
+   * mark uses this registry to find the live checkpoint with the same id and finalize that one.
+   */
+  private static final ConcurrentHashMap<String, JmsCheckpointMark> LIVE_CHECKPOINTS =
+      new ConcurrentHashMap<>();
+
+  /** Live checkpoints not finalized after this long are dropped; the broker redelivers them. */
+  private static final long LIVE_CHECKPOINT_MAX_AGE_MILLIS = 60 * 60 * 1000L;
+
+  private final String id = UUID.randomUUID().toString();
+  private final long creationTimeMillis = System.currentTimeMillis();
   private Instant oldestMessageTimestamp;
   private transient @Nullable List<Message> messages;
   private transient @Nullable MessageConsumer consumer;
@@ -66,6 +82,15 @@ private JmsCheckpointMark(
   /** Acknowledge all outstanding message. */
   @Override
   public void finalizeCheckpoint() {
+    if (messages == null && consumer == null && session == null) {
+      // Possibly a restored copy of a checkpoint whose state lives in another instance.
+      JmsCheckpointMark live = LIVE_CHECKPOINTS.get(id);
+      if (live != null) {
+        live.finalizeCheckpoint();
+        return;
+      }
+    }
+    LIVE_CHECKPOINTS.remove(id);
     try {
       if (messages != null) {
         for (Message message : messages) {
@@ -106,6 +131,19 @@ public void finalizeCheckpoint() {
     }
   }
 
+  private void registerLive() {
+    long now = System.currentTimeMillis();
+    LIVE_CHECKPOINTS
+        .values()
+        .removeIf(m -> now - m.creationTimeMillis > LIVE_CHECKPOINT_MAX_AGE_MILLIS);
+    LIVE_CHECKPOINTS.put(id, this);
+  }
+
+  @VisibleForTesting
+  static int liveCheckpointCount() {
+    return LIVE_CHECKPOINTS.size();
+  }
+
   @VisibleForTesting
   @Nullable List<Message> getMessages() {
     return messages;
@@ -251,6 +289,9 @@ JmsCheckpointMark newCheckpoint(
                   activeCheckpoints);
           messages.clear();
           oldestMessageTimestamp = Instant.now();
+          if (messagesCopy != null) {
+            checkpointMark.registerLive();
+          }
           if (activeCheckpoints != null) {
             activeCheckpoints.incrementAndGet();
           }

sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java

@@ -598,6 +598,39 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception {
     assertEquals(7, count(QUEUE));
   }
 
+  @Test
+  public void testRestoredCheckpointMarkFinalizedAfterReaderRecreation() throws Exception {
+    UnboundedJmsReader reader = setupReaderForTest();
+    assertTrue(startWithRetry(reader));
+    assertTrue(advanceWithRetry(reader));
+    assertTrue(advanceWithRetry(reader));
+
+    JmsCheckpointMark mark = (JmsCheckpointMark) reader.getCheckpointMark();
+    assertEquals(10, count(QUEUE));
+
+    // the runner serializes the mark and recreates the reader from it
+    java.io.ByteArrayOutputStream bytes = new java.io.ByteArrayOutputStream();
+    try (java.io.ObjectOutputStream out = new java.io.ObjectOutputStream(bytes)) {
+      out.writeObject(mark);
+    }
+    JmsCheckpointMark restored;
+    try (java.io.ObjectInputStream in =
+        new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(bytes.toByteArray()))) {
+      restored = (JmsCheckpointMark) in.readObject();
+    }
+    reader.close();
+    UnboundedJmsReader recreated =
+        (UnboundedJmsReader)
+            ((JmsIO.UnboundedJmsSource<?>) reader.getCurrentSource())
+                .createReader(PipelineOptionsFactory.create(), restored);
+    assertTrue(startWithRetry(recreated));
+
+    // finalizing the restored copy acks the messages read by the original reader
+    restored.finalizeCheckpoint();
+
+    assertEquals(7, count(QUEUE));
+  }
+
   @Test
   public void testCheckpointMarkAndFinalizeSeparatelyIndividualAcknowledge() throws Exception {
     UnboundedJmsReader reader = setupReaderForTest(JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE);

sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsLocalTest.java

@@ -170,6 +170,37 @@ public void testJmsCheckpointMarkClientAcknowledgeUnsafeNoSessionRecreation() th
     assertEquals(0, activeCheckpoints.get());
   }
 
+  @Test
+  public void testRestoredCheckpointMarkAcknowledgesLiveMessages() throws Exception {
+    Message msg = Mockito.mock(Message.class);
+    JmsCheckpointMark.Preparer preparer =
+        JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE);
+    preparer.add(msg);
+    AtomicInteger activeCheckpoints = new AtomicInteger(0);
+    JmsCheckpointMark mark =
+        preparer.newCheckpoint(
+            null, null, JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE, activeCheckpoints);
+    int live = JmsCheckpointMark.liveCheckpointCount();
+
+    // serialize/deserialize drops the transient state, as when a runner restores the mark
+    java.io.ByteArrayOutputStream bytes = new java.io.ByteArrayOutputStream();
+    try (java.io.ObjectOutputStream out = new java.io.ObjectOutputStream(bytes)) {
+      out.writeObject(mark);
+    }
+    JmsCheckpointMark restored;
+    try (java.io.ObjectInputStream in =
+        new java.io.ObjectInputStream(new java.io.ByteArrayInputStream(bytes.toByteArray()))) {
+      restored = (JmsCheckpointMark) in.readObject();
+    }
+    assertNull(restored.getMessages());
+
+    restored.finalizeCheckpoint();
+
+    Mockito.verify(msg, Mockito.times(1)).acknowledge();
+    assertEquals(0, activeCheckpoints.get());
+    assertEquals(live - 1, JmsCheckpointMark.liveCheckpointCount());
+  }
+
   /** Test the checkpoint mark default coder, which is actually AvroCoder. */
   @Test
   public void testCheckpointMarkDefaultCoder() throws Exception {