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
| MetricsReader metricsReader = new MetricsReader(readResult, NAMESPACE); | ||
| long startTime = System.currentTimeMillis(); | ||
| long timeoutMillis = OPTIONS.getReadTimeout() * 1000L; | ||
| PipelineResult.State readState = readResult.getState(); |
There was a problem hiding this comment.
Outside the streaming IO scope I was tagged for, but flagging it since I noticed it. readState comes from readResult.getState(), which on the direct runner returns RUNNING and never null. So when readTimeout expires with the pipeline still running (the amqp slow on CI case from #26175 that got this test disabled), cancelIfTimeouted, which only cancels when readState == null, is a no op. The asserts then fail while the pipeline keeps consuming from the class scoped shared BROKERS, leaking it into the next parameterized test. Suggestion:
if (readState == null || !readState.isTerminal()) {
readResult.cancel();
}There was a problem hiding this comment.
Thanks for pointing that. Before #28014 (edit: probably another earlier fix) direct runner returns null on querying state when running. This part was outdated and now fixed.
| } | ||
| } | ||
| return checkpointMarkPreparer.newCheckpoint(consumerToClose, sessionTofinalize); | ||
| activeCheckpoints.incrementAndGet(); |
There was a problem hiding this comment.
activeCheckpoints.incrementAndGet() runs outside the synchronized(this) block and unconditionally after the isEmpty() check.
If doClose() or discard() races in between, newCheckpoint takes the discarded branch and returns emptyCheckpoint() whose activeCheckpoints is null (JmsCheckpointMark.java:265), so the increment is never balanced and the executor loop waits the full closeTimeout. In CLIENT_ACKNOWLEDGE, the old consumerToClose and sessionTofinalize captured just above (after recreateSession already swapped in a new session) are then silently dropped, leaking the old session and its prefetched messages. Suggest moving the increment inside newCheckpoint under the write lock on the non discarded branch, and closing the passed in consumer and session in the discarded branch.
There was a problem hiding this comment.
moving the increment inside newCheckpoint under the write lock on the non discarded branch
Done
closing the passed in consumer and session in the discarded branch
I'm not sure if we should do this. On discard branch we don't mutate consumer and session, and it constructively it may be used later
| executorService.submit( | ||
| () -> { | ||
| long startTime = System.currentTimeMillis(); | ||
| long timeoutMillis = source.spec.getCloseTimeout().getMillis(); | ||
| while (activeCheckpoints.get() > 0 | ||
| && System.currentTimeMillis() - startTime < timeoutMillis) { | ||
| try { | ||
| Thread.sleep(1_000); // poll in 1 sec interval | ||
| } catch (InterruptedException ignored) { | ||
| break; | ||
| } | ||
| } | ||
| LOG.debug( | ||
| "Closing connection after checkpoints finalized or timeout: {}", | ||
| source.spec.getCloseTimeout()); | ||
| closeConsumer(); | ||
| closeSession(); | ||
| closeConnection(); | ||
| }); |
There was a problem hiding this comment.
The close path submit()s a while (...) Thread.sleep(1000) loop onto the shared ExecutorOptions scheduled pool. On any runner where checkpoint finalization is best effort or delayed, the count may never reach 0, so the loop runs up to the full closeTimeout (default 60s).
When many readers close at once, for example on scale down, that parks a lot of pool threads and can starve other scheduled work. I am not sure how often Dataflow drops finalization in practice, you would know better, but even setting that aside the busy wait on a shared pool seems worth avoiding.
Suggest rescheduling via schedule() with a short period, or having the last finalizeCheckpoint decrement trigger the close when a closed flag is set (last one out), with a single schedule(closeTimeout) fallback.
There was a problem hiding this comment.
Agree this would theoretically happen, though in my observation this is rare.
The ScheduledExecutorService only started when activeCheckpoints > 0 (either there is active checkpoint or there is checkpoint never get finalized). JmsCheckpointMark.finalizeCheckpoint does not throw, so when it's executed, the counter is decremented.
Suggest rescheduling via schedule() with a short period
Done
| * AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}. | ||
| * | ||
| * <p>Different JMS providers use different proprietary integer constants for individual | ||
| * acknowledgment (e.g., ActiveMQ uses 4, Qpid JMS / ActiveMQ Artemis / IBM MQ use 101). |
There was a problem hiding this comment.
The withIndividualAcknowledgeModeCode javadoc says "Qpid JMS / ActiveMQ Artemis / IBM MQ use 101". IBM MQ classes for JMS only support the standard acknowledgment modes (AUTO=1, CLIENT=2, DUPS_OK=3) and have no individual acknowledge constant, so 101 is not a valid mode there and createSession(false, 101) on IBM MQ would throw a JMSException rather than give individual ack.
Suggest dropping IBM MQ from that list. For per message behavior on IBM MQ the guidance is CLIENT_ACKNOWLEDGE, which acks cumulatively. The "ActiveMQ uses 4" part checks out against ActiveMQSession.INDIVIDUAL_ACKNOWLEDGE. Minor, since getAckModeCode correctly throws for unrecognized providers, it might be worth a one line note that INDIVIDUAL mode on other providers requires an explicit withIndividualAcknowledgeModeCode.
There was a problem hiding this comment.
Removed IBM MQ. checked that "IBM MQ use 101" was wrong and fixed getAckModeCode manually before pushing, but forgot to clean up the doc here.
Added description in javadoc
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.