chamikaramj commented on code in PR #40605:
URL: https://github.com/apache/beam/pull/40605#discussion_r4355037331
##########
sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java:
##########
@@ -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.
Review Comment:
This uses `AvroCoder` for encoding right ? `id` field could be null during
certain scenarios (for example, restoring from an older version during Dataflow
streaming update). You should guard against that here.
##########
sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java:
##########
@@ -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 =
Review Comment:
Going through all items here during checkpoint creation can be expensive.
Consider replacing with a Guava Cache to put/get/invalidate checkpoints
efficiently.
##########
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
Review Comment:
Let's serialize/deserialize using the actual `AvroCoder` instead of using
Java serialization for this test.
##########
sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java:
##########
@@ -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 {
Review Comment:
May be return immediately here if `messages` is null ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
