Support ACKNOWLEDGE_MODE for JmsIO#39253
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request adds support for configurable acknowledgement modes in JmsIO, enabling better compatibility with various JMS providers and performance tuning. The changes involve a significant refactoring of how JmsIO manages JMS sessions and consumers, ensuring that resources are correctly handled across different acknowledgement strategies. Additionally, the PR improves the reliability of integration tests and addresses a bug in the direct runner's checkpoint finalization logic. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for multiple acknowledgment modes in JmsIO (CLIENT_ACKNOWLEDGE, CLIENT_ACKNOWLEDGE_UNSAFE, and INDIVIDUAL_ACKNOWLEDGE) to optimize message consumption and prevent message backlogs, along with corresponding updates to checkpointing and reader closing logic. The review feedback highlights two important issues: a potential race condition in JmsIO.java where the consumer and session could be closed concurrently on different threads in CLIENT_ACKNOWLEDGE mode, and a potential hang during reader shutdown if a RuntimeException in JmsCheckpointMark.finalizeCheckpoint() prevents the decrement of activeCheckpoints.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if (watermark.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { | ||
| // If the reader had no elements available, but the shard is not done, reuse it later | ||
| // Might be better to finalize old checkpoint. | ||
| // If the reader had no elements available, but the shard is not done, reuse it later. |
There was a problem hiding this comment.
in a typical run, 5/6 tests observe small amount of messages are not acked due to the known issue noted here:
- testPublishingThenReadingAllClientAcknowledgeUnsafe: 1
- testPublishingThenReadingAllIndividualAcknowledge: 13
- testPublishingThenReadingAllIndividualAcknowledge: 13
- testPublishingThenReadingAll: 34
- testPublishingThenReadingAll: 33
There was a problem hiding this comment.
How about on dataflow runner? Do we have the same problem?
There was a problem hiding this comment.
Will follow up on Dataflow runner (need a remote ActiveMQ deployement)
The original report where this issue was observed was on Direct runner: #30218 (comment)
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
* CheckpointMark behavior in alignment with different ACKNOWLEDGEMENT_MODE * Ref count active checkpoint for quicker onClose that releases session * Fix hanging checkpoint when no incoming data in direct runner. This allows us to do an exact assert * Optimize long running unit test usign a short retry * Re-enable AMQP integration test after stuck unack messages resolved
Abacn
changed the title
There was a problem hiding this comment.
Thanks @Abacn . Looking at this from the streaming IO side you tagged me for. The two safe modes look correct on the data loss axis. Acks fire only in finalizeCheckpoint, and failure paths degrade to redelivery. I have one question on the two new modes when finalization runs on a separate thread, plus a few reader lifecycle and docs items inline. I also noticed one thing in the integration test that was re enabled. It is outside streaming, but I flagged it anyway.
As @shunping already approved, I just added my comments below
| for (Message message : messages) { | ||
| message.acknowledge(); |
There was a problem hiding this comment.
Question on the two new modes when finalization runs on a separate thread:
In INDIVIDUAL_ACKNOWLEDGE and CLIENT_ACKNOWLEDGE_UNSAFE the mark holds Messages bound to the reader's single long lived session. My understanding is that finalizeCheckpoint() may be called on a different thread from the reader on some runners.
You would know the Dataflow specifics far better than I do. If so, does message.acknowledge() here end up running concurrently with the reader's receive loop on the same Session, which JMS (section 4.4.6) says must have a single thread of control? CLIENT_ACKNOWLEDGE looks immune since each mark owns a private session. Could you confirm how ack is serialized against receive for the other two modes, or which providers this has been validated against? (This ties into @shunping 's Dataflow question above).
There was a problem hiding this comment.
My understanding is that finalizeCheckpoint() may be called on a different thread from the reader on some runners.
Correct
does message.acknowledge() here running concurrently with the reader's receive loop on the same Session
No for INDIVIDUAL_ACKNOWLEDGE and CLIENT_ACKNOWLEDGE_UNSAFE. That's why the latter is "unsafe". The former isn't part of Jms spec but some providers' extended feature.
Could you confirm how ack is serialized against receive for the other two modes, or which providers this has been validated against?
The integration tests now covers this question, tested on ActiveMQ and Amqp provider.
CLIENT_ACKNOWLEDGElooks immune since each mark owns a private session
Yes
There was a problem hiding this comment.
Thank you, that mostly clears it up! One part I still want to nail/understand:
The CheckpointMark contract in UnboundedSource.java says finalizeCheckpoint "may be called from any thread, concurrently with calls to the UnboundedReader it was created from", and the reader is single threaded but finalize is exempted. For INDIVIDUAL_ACKNOWLEDGE and CLIENT_ACKNOWLEDGE_UNSAFE getCheckpointMark does not recreate the session, so the mark's messages stay bound to the reader's live session, and finalizeCheckpoint calls message.acknowledge() on it without taking the reader monitor. CLIENT_ACKNOWLEDGE avoids this because it hands each checkpoint its own session.
On Dataflow, where commit finalization runs asynchronously on a pool thread while the cached reader keeps advancing, what stops acknowledge() from racing the reader's receiveNoWait() on the same session? Same question for the legacy Flink UnboundedSourceWrapper, which calls finalizeCheckpoint outside the checkpoint lock. The integration tests run on the direct runner, where finalize is on the reader thread, so they would not exercise this. (Or am I missing something?)
If I am missing the mechanism that serializes them I am happy to be corrected. Otherwise acking under the reader monitor for these two modes, or giving them per checkpoint sessions like CLIENT_ACKNOWLEDGE, would close this.
There was a problem hiding this comment.
@Abacn this is my only remaining point otherwise I can also approve it, just let me know!
There was a problem hiding this comment.
I agree a race is theoretically possible. Jms spec's requirement for session operations running in single thread and Beam's asynchronous checkpointing resulted in this fundamental incompatibility. recreate session in CLIENT_ACKNOWLEDGE tried to workaround it
-
CLIENT_ACKNOWLEDGE_UNSAFE
As noted previously, this mode is intended to restore the same behavior and timing prior to Beam 2.50.0, regardless race exists or not, based on the users report that certain Jms providers work with Beam<2.50.0 but not later. -
INDIVIDUAL_ACKNOWLEDGE is an open question. I did more investigation
public doc emphasizing "session not threadsafe" or "one thread per session" is referring to receive:
The JMS contract is that only 1 session is used by one thread at once - which if you’re using consumers means that only 1 consumer can receive messages at once if using the same session.
(see https://activemq.apache.org/components/classic/documentation/multiple-consumers-on-a-queue)
but none saying acknowledge and receive must be in same thread. I think AI generated answer based on mixing up concepts in the available doc
If prompt is specific, here is an answer I got
Q: In Jms, Active MQ INDIVIDUAL_ACKNOELEDGE mode, is there a race if receive new message and acknowledge old message happens concurrently
A: No, there is no race condition between receiving a new message and acknowledging an old message in INDIVIDUAL_ACKNOWLEDGE mode. ActiveMQ tracks and deletes messages on the broker based on the specific message object you call acknowledge() on.
3 participants
Introduce INDIVIDUAL_ACKNOWLEDGEMENT for supported providers
(Re)introduce CLIENT_ACKNOWLEDGEMENT_UNSAFE that restore the behavior prior to Beam 2.55.0, provide a way to mitigate #30218 (comment)
CheckpointMark behavior in alignment with different ACKNOWLEDGEMENT_MODE
Ref count active checkpoint for quicker onClose that releases session
Fix hanging checkpoint when no incoming data in direct runner. This allows us to do an exact assert
Optimize long running unit test usign a short retry
Re-enable AMQP integration test after stuck unack messages resolved
Please add a meaningful description for your change here
Fix #26203 ; fix #26175
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.