[JmsIO] Let restored checkpoint marks acknowledge live messages

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

@@ -26,11 +26,15 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 import org.slf4j.Logger;
@@ -42,8 +46,25 @@
  */
 class JmsCheckpointMark implements UnboundedSource.CheckpointMark, Serializable {
 
+  // Pinned to the value computed for this class before the id field was added, so checkpoints
+  // written by those versions can still be restored.
+  private static final long serialVersionUID = -275230511397134772L;
+
   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. To
+   * bound memory, at most 10,000 checkpoints are kept and those not finalized within an hour are
+   * dropped; the broker redelivers their messages.
+   */
+  private static final Cache<String, JmsCheckpointMark> LIVE_CHECKPOINTS =
+      CacheBuilder.newBuilder().maximumSize(10_000).expireAfterWrite(1, TimeUnit.HOURS).build();
+
+  // May be null when restored from a checkpoint written by a version without this field.
+  private final @Nullable String id = UUID.randomUUID().toString();
   private Instant oldestMessageTimestamp;
   private transient @Nullable List<Message> messages;
   private transient @Nullable MessageConsumer consumer;
@@ -65,12 +86,27 @@ private JmsCheckpointMark(
 
   /** Acknowledge all outstanding message. */
   @Override
-  public void finalizeCheckpoint() {
+  @SuppressWarnings("ReferenceEquality") // live must be a different instance, not an equal one
+  public synchronized void finalizeCheckpoint() {
+    if (id != null) {
+      // Atomic, so only one copy of this checkpoint ever finalizes the live one.
+      JmsCheckpointMark live = LIVE_CHECKPOINTS.asMap().remove(id);
+      if (live != null && live != this) {
+        // This is a restored copy whose state lives in another instance.
+        live.finalizeCheckpoint();
+        return;
+      }
+    }
+    if (messages == null && consumer == null && session == null && activeCheckpoints == null) {
+      // Nothing to acknowledge or release, e.g. a restored copy that was already finalized.
+      return;
+    }
     try {
       if (messages != null) {
         for (Message message : messages) {
           message.acknowledge();
         }
+        messages = null;
       }
     } catch (JMSException e) {
       // The effect of this is message not get acknowledged and thus will be redelivered. It is
@@ -103,9 +139,19 @@ public void finalizeCheckpoint() {
 
     if (activeCheckpoints != null) {
       activeCheckpoints.decrementAndGet();
+      activeCheckpoints = null;
     }
   }
 
+  private void registerLive() {
+    LIVE_CHECKPOINTS.put(Objects.requireNonNull(id), this);
+  }
+
+  @VisibleForTesting
+  static long liveCheckpointCount() {
+    return LIVE_CHECKPOINTS.size();
+  }
+
   @VisibleForTesting
   @Nullable List<Message> getMessages() {
     return messages;
@@ -139,6 +185,7 @@ public boolean equals(@Nullable Object o) {
       return false;
     }
     JmsCheckpointMark that = (JmsCheckpointMark) o;
+    // The id is random and only used to find the live checkpoint, so it is not part of equality.
     return oldestMessageTimestamp.equals(that.oldestMessageTimestamp);
   }
 
@@ -251,6 +298,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

@@ -90,6 +90,7 @@
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.SerializableBiFunction;
+import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
 import org.apache.qpid.jms.JmsAcknowledgeCallback;
@@ -598,6 +599,29 @@ 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 encodes the mark and recreates the reader from it
+    JmsIO.UnboundedJmsSource<?> source = (JmsIO.UnboundedJmsSource<?>) reader.getCurrentSource();
+    JmsCheckpointMark restored = CoderUtils.clone(source.getCheckpointMarkCoder(), mark);
+    reader.close();
+    UnboundedJmsReader recreated = source.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

@@ -35,6 +35,7 @@
 import jakarta.jms.Session;
 import java.io.IOException;
 import java.io.NotSerializableException;
+import java.util.Base64;
 import java.util.List;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -45,6 +46,7 @@
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.testing.CoderProperties;
+import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.util.SerializableUtils;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
 import org.joda.time.Duration;
@@ -170,7 +172,51 @@ public void testJmsCheckpointMarkClientAcknowledgeUnsafeNoSessionRecreation() th
     assertEquals(0, activeCheckpoints.get());
   }
 
-  /** Test the checkpoint mark default coder, which is actually AvroCoder. */
+  @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);
+    long live = JmsCheckpointMark.liveCheckpointCount();
+
+    // encode/decode drops the transient state, as when a runner restores the mark
+    Coder<JmsCheckpointMark> coder =
+        new JmsIO.UnboundedJmsSource<JmsRecord>(null).getCheckpointMarkCoder();
+    JmsCheckpointMark restored = CoderUtils.clone(coder, mark);
+    assertNull(restored.getMessages());
+
+    restored.finalizeCheckpoint();
+
+    Mockito.verify(msg, Mockito.times(1)).acknowledge();
+    assertEquals(0, activeCheckpoints.get());
+    assertEquals(live - 1, JmsCheckpointMark.liveCheckpointCount());
+  }
+
+  /** A checkpoint mark encoded by the checkpoint mark coder before marks had an id. */
+  private static final String CHECKPOINT_MARK_WITHOUT_ID =
+      "rO0ABXNyACxvcmcuYXBhY2hlLmJlYW0uc2RrLmlvLmptcy5KbXNDaGVja3BvaW50TWFya/wuL07g0x5MAgAB"
+          + "TAAWb2xkZXN0TWVzc2FnZVRpbWVzdGFtcHQAF0xvcmcvam9kYS90aW1lL0luc3RhbnQ7eHBzcgAVb3JnLmpv"
+          + "ZGEudGltZS5JbnN0YW50Lci+0MYOnM0CAAFKAAdpTWlsbGlzeHAAAAGLz+VoAA==";
+
+  @Test
+  public void testDecodeCheckpointMarkWithoutId() throws Exception {
+    Coder<JmsCheckpointMark> coder =
+        new JmsIO.UnboundedJmsSource<JmsRecord>(null).getCheckpointMarkCoder();
+    JmsCheckpointMark restored =
+        CoderUtils.decodeFromByteArray(
+            coder, Base64.getDecoder().decode(CHECKPOINT_MARK_WITHOUT_ID));
+    assertNull(restored.getMessages());
+
+    // nothing to acknowledge, and a missing id must not break finalization
+    restored.finalizeCheckpoint();
+  }
+
+  /** Test the checkpoint mark default coder, which is SerializableCoder. */
   @Test
   public void testCheckpointMarkDefaultCoder() throws Exception {
     JmsCheckpointMark jmsCheckpointMark =