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..6c19a554db8e 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,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 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 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 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(); } 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..f01000db60ee 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 @@ -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); 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..39c0b0ae90ff 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 @@ -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 coder = + new JmsIO.UnboundedJmsSource(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 coder = + new JmsIO.UnboundedJmsSource(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 =