[jbonofre]

JmsCheckpointMark keeps its messages, consumer and session in transient fields. A mark that is serialized and restored, or finalized after the reader was recreated, therefore acknowledges nothing and the messages are redelivered.

This gives each mark an id and keeps live marks in a per-JVM registry, so a restored copy finalizes the live checkpoint with the same id. The entry is removed on finalize, and entries not finalized within an hour are dropped (the broker redelivers them).

Motivated by a Dataflow streaming job (CLIENT_ACKNOWLEDGE, withRequiresDeduping()) where dequeues stayed near zero while readers were recreated every few seconds. This is a hypothesis for that behavior, not yet confirmed. It does not help if finalize is delivered to a different JVM, or if the old connection is already closed.

Adds a test that serializes and restores a mark and checks the original message is acked.

[github-actions]

📋 Review Summary

This PR adds a static per-JVM registry (LIVE_CHECKPOINTS) to JmsCheckpointMark so that checkpoint marks serialized and restored on the same JVM can locate and finalize the original live checkpoint containing the transient JMS messages, consumer, and session. While the overall approach is sound for same-JVM checkpoint restoration, there are concurrency and state management issues in finalizeCheckpoint() as well as a missing id comparison in equals() / hashCode() that should be addressed.

🔍 General Feedback

• Thread safety and idempotency in finalizeCheckpoint(): Because a restored checkpoint mark delegates to the live instance in LIVE_CHECKPOINTS and multiple copies (or the original and a restored copy) could be finalized, finalizeCheckpoint() should atomically remove from the map and guard against double finalization (especially double-decrementing activeCheckpoints).
• Equality contract: Since id is now a core non-transient field distinguishing checkpoint marks, equals() and hashCode() should be updated to include it.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

💬 Inline Findings

• sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java (line 84): 🟠 High: finalizeCheckpoint() is not thread-safe and is vulnerable to double-execution or infinite recursion in edge cases:

1. Infinite recursion hazard: If a live checkpoint were ever registered with messages == null && consumer == null && session == null, LIVE_CHECKPOINTS.get(id) would return this and recurse infinitely before reaching LIVE_CHECKPOINTS.remove(id).
2. Race condition & non-idempotency: Using .get(id) followed later by .remove(id) is non-atomic. Furthermore, messages and activeCheckpoints are never nulled out after finalization, while consumer and session are. If finalizeCheckpoint() is called more than once (e.g., on the live mark and a restored copy, or two restored copies concurrently), activeCheckpoints.decrementAndGet() will be called multiple times—corrupting the reader's active checkpoint count—and messages will be acknowledged again.

Use LIVE_CHECKPOINTS.remove(id) atomically (and check live != this), synchronize finalization on the instance, and clear messages and activeCheckpoints after processing:

  @Override
  public synchronized void finalizeCheckpoint() {
    JmsCheckpointMark live = LIVE_CHECKPOINTS.remove(id);
    if (live != null && live != this) {
      live.finalizeCheckpoint();
      return;
    }

• sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java (line 130): 🟠 High: Null out messages and activeCheckpoints after finalization so that repeated calls to finalizeCheckpoint() on the same instance do not re-acknowledge messages or decrement activeCheckpoints multiple times (causing it to drop below 0).

      activeCheckpoints.decrementAndGet();
      activeCheckpoints = null;

• sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java (line 207): 🟡 Medium: Now that each JmsCheckpointMark has a unique serialized id used to identify the specific live checkpoint, equals and hashCode should include id alongside oldestMessageTimestamp.

    return Objects.equals(id, that.id) && Objects.equals(oldestMessageTimestamp, that.oldestMessageTimestamp);

[github-actions]

Assigning reviewers:

R: @chamikaramj for label java.
Backup reviewer: kennknowles (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[chamikaramj]

Thanks!

[chamikaramj]

Going through all items here during checkpoint creation can be expensive.

Consider replacing with a Guava Cache to put/get/invalidate checkpoints efficiently.

[chamikaramj]

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.

[chamikaramj]

May be return immediately here if messages is null ?

[chamikaramj]

Let's serialize/deserialize using the actual AvroCoder instead of using Java serialization for this test.

[chamikaramj]

waiting on author