From 94cf005d80100a03b96a2d8b5a23aa2100755859 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jean-Baptiste=20Onofr=C3=A9?= Date: Thu, 15 Oct 2026 09:22:09 +0200 Subject: [PATCH 1/2] [JmsIO] Let restored checkpoint marks acknowledge live messages JmsCheckpointMark keeps its messages, consumer and session in transient fields, so a mark that is serialized and restored (or finalized after the reader was recreated) acknowledges nothing. Give each mark an id and keep live marks in a per-JVM registry so a restored copy can finalize the live one. Entries not finalized within an hour are dropped. --- .../beam/sdk/io/jms/JmsCheckpointMark.java | 41 +++++++++++++++++++ .../apache/beam/sdk/io/jms/JmsLocalTest.java | 31 ++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java index fdfd2c0c7bce..fe69338779ff 100644 --- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java +++ b/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 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 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 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(); } diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsLocalTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsLocalTest.java index 277c84151703..84386c8cb1f2 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsLocalTest.java +++ b/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 { From 165d4f6ac1c702609309020c6209458740e078ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jean-Baptiste=20Onofr=C3=A9?= Date: Fri, 16 Oct 2026 17:36:43 +0200 Subject: [PATCH 2/2] [JmsIO] Test finalizing a restored checkpoint mark after reader recreation --- .../org/apache/beam/sdk/io/jms/JmsIOTest.java | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java index bdcf27bbfa49..70a2b26cee7d 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java +++ b/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);