Implement Asynchronous wrapper for DoFn in Java SDK#38609
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 introduces an asynchronous wrapper for 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
|
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper designed to convert synchronous DoFn operations into asynchronous ones within Apache Beam pipelines. It includes a core implementation using a thread pool and persistent state for buffering, along with a comprehensive test suite. The review identified several critical issues: a performance bottleneck caused by global map scans during timer firing, potential data loss due to ignored timestamps in the output receiver, and logic bugs in the reconciliation loop that could lead to redundant task rescheduling and quadratic complexity. Additionally, the use of a single static lock for all instances may cause thread contention.
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class designed to process elements asynchronously using a background thread pool, along with a comprehensive test suite AsyncDoFnTest. The code review highlights several critical concurrency and correctness issues in the implementation. Specifically, the wrapped syncFn is shared across background threads without being thread-safe, exceptions during async execution are caught and result in silent data loss, and tearing down cloned instances can prematurely shut down shared thread pools. Additionally, the reviewer noted a potential division-by-zero bug if timerFrequency is non-positive, and recommended documenting the class's limitations in the Javadoc.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class that converts synchronous Beam DoFns into asynchronous ones using thread pools and state-backed timers, along with a comprehensive test suite. The review feedback highlights several critical areas for improvement: avoiding blocking the runner's timer thread during rescheduling, resolving a performance bottleneck caused by a single JVM-wide static lock, removing redundant locking around atomic counter decrements, ensuring thread safety for the wrapped DoFn via thread-local isolation, preventing potential memory leaks in the static tracking map using an expiring cache, leveraging Beam's relative timer API for deterministic testing, and cleaning up an unused test rule.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class designed to convert synchronous Apache Beam DoFn operations into asynchronous ones, along with a comprehensive test suite. The reviewer identified several critical and high-severity issues in the implementation: event-time timestamp corruption in AccumulatingOutputReceiver, severe JVM-wide lock contention from a static ReentrantLock, timer starvation in processElement under continuous element arrival, redundant rescheduling of completed tasks on retry when sibling tasks fail, unnecessary thread pool creation when useThreadPool is false, and potential state corruption in commitFinishedItems if InputT lacks proper equals and hashCode implementations.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class that converts synchronous DoFn operations into asynchronous ones, along with a comprehensive test suite to verify its behavior under various scenarios. The code review identified several critical issues: high complexity in commitFinishedItems due to flat map iteration over all active elements across all keys, JVM-wide lock contention caused by a static ReentrantLock, thread-safety issues when iterating over a synchronized list in AccumulatingOutputReceiver without manual synchronization, and the ineffectiveness of CompletableFuture.cancel(true) in interrupting background threads.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class designed to convert synchronous DoFn operations into asynchronous ones, along with a comprehensive suite of unit tests in AsyncDoFnTest. The review feedback highlights several critical issues: severe thread contention caused by a JVM-wide static lock, performance degradation during timer firing due to scanning all active elements across all keys, loss of custom windowing and watermark tracking from hardcoding GlobalWindow.INSTANCE and shifting timestamps when rescheduling, and unnecessary synchronization overhead in AccumulatingOutputReceiver from using ConcurrentLinkedQueue instead of ArrayList.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class designed to convert synchronous DoFn operations into asynchronous ones within Apache Beam pipelines, accompanied by a comprehensive suite of unit tests. The review feedback highlights several critical performance and design improvements: replacing the global static ReentrantLock with per-UUID locks to prevent severe thread contention, partitioning the processingElements map by key to eliminate expensive O(N) scans during timer processing, simplifying the timer firing calculations with pure integer math, and unwrapping ExecutionException to improve error trace readability.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #38609 +/- ##
============================================
+ Coverage 57.77% 57.81% +0.04%
- Complexity 12969 13086 +117
============================================
Files 2509 2512 +3
Lines 260525 261222 +697
Branches 10658 10760 +102
============================================
+ Hits 150516 151038 +522
- Misses 104318 104456 +138
- Partials 5691 5728 +37
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces AsyncDoFn, a wrapper class designed to convert synchronous Beam DoFns into asynchronous ones, along with a comprehensive test suite in AsyncDoFnTest. The code review highlights several critical concurrency, performance, and operational issues. Specifically, there is a race condition in teardown() where the wrapped DoFn is destroyed before the thread pool is shut down, and a risk of blocking the runner's main thread during timer execution by calling scheduleItem instead of a non-blocking alternative. Additionally, the reviewer noted performance bottlenecks from O(N) lookups in processingElements and JVM-wide lock contention on a single static lock, as well as operational concerns regarding thread dump readability and the use of ForkJoinPool.commonPool() for high-latency operations.
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.
…he Apache Beam Java SDK (apache#38529)
…emoved O(N) global activeElements scan. Fixed logic bug where duplicate elements were incorrectly marked for rescheduling. Optimized lookups by converting finishedItems from a list to a HashSet.
…ng sleep. Added two more tests to match Python SDK. Fixed formatting issues.
…ixes. Spot Bugs potential fixes.
… that preserves deterministic jitter behavior and avoids pressure on garbage collector (apache#38529)
…mini-code-assist - Propagates asynchronous task exceptions as RuntimeExceptions to prevent silent data loss and enable runner-level retries. - Implements a static refCounts registry to safely tear down the shared executor service only when the last cloned instance is destroyed. - Validates timerFrequency in the constructor to prevent zero/negative values from entering infinite loops. - Documents multi-threading requirements, multi-output limitations, and bundle lifecycle behaviors in a class-level comment.
Stores the partition key inside InFlightElement and cancels/purges orphaned futures inside commitFinishedItems. This prevents silent memory leaks on bundle rollbacks (apache#38529)
…r to preserve event-time downstream. Initializes the ExecutorService thread pool only when useThreadPool is true. Refactors state filtering to use finishedElementIds instead of finishedItems, preventing duplicate processing (apache#38529)
…parameter without Hardcoding GlobalWindow.Instance (apache#38529)
…s involving inputTimestamp(apache#38529)
17e66d1 to
7e0ebc5
Compare
…Might also need to rename python sdk too.
| } | ||
| } | ||
|
|
||
| if (verboseLogging) { |
There was a problem hiding this comment.
If we do verboselogging checks we should do it for all places, let's add the check for other logs and also make a util function for logging maybe so we don't have to do this if check everytime.
There was a problem hiding this comment.
Added verbose logging checks everywhere and made private util functions for logging in latest commit. Let me know if I should change it.
…ArgumentProvider code block into util functions. Refactored verbose logging into util functions. Replaced hardcoded numbers with static constants. Updated test helper methods with @VisibleForTesting (apache#38529)
|
Thanks! Overall LGTM, once the logging changes are in we will tag Danny for final review. |
|
I also just realized you have files duplicated and the naming for the class should match AsyncWrapper instead of calling AsyncDoFn |
|
Thanks! Tagging Danny for the final check R: @damccorm |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
|
Rerunning some workflows to confirm this doesn't break them before merging |
* Created an Asynchronous Wrapper for DoFn as well as JUnit tests for the Apache Beam Java SDK (apache#38529) * Optimize State reconciliation loop and eliminate O(N^2) complexity. Removed O(N) global activeElements scan. Fixed logic bug where duplicate elements were incorrectly marked for rescheduling. Optimized lookups by converting finishedItems from a list to a HashSet. * Added check for long overflow possibility when exponentially increasing sleep. Added two more tests to match Python SDK. Fixed formatting issues. * Fix Timestamp propagation and add relevant test too. Spotless Apply fixes. Spot Bugs potential fixes. * Resolve SpotBugs DMI_RANDOM_USED_ONLY_ONCE replacing new Random(seed) that preserves deterministic jitter behavior and avoids pressure on garbage collector (apache#38529) * Improve AsyncDoFn robustness and fix critical warnings provided by gemini-code-assist - Propagates asynchronous task exceptions as RuntimeExceptions to prevent silent data loss and enable runner-level retries. - Implements a static refCounts registry to safely tear down the shared executor service only when the last cloned instance is destroyed. - Validates timerFrequency in the constructor to prevent zero/negative values from entering infinite loops. - Documents multi-threading requirements, multi-output limitations, and bundle lifecycle behaviors in a class-level comment. * Fixed formatting issue * Implement cross-key task cancellation in AsyncDoFn. Stores the partition key inside InFlightElement and cancels/purges orphaned futures inside commitFinishedItems. This prevents silent memory leaks on bundle rollbacks (apache#38529) * Passes original element's inputTimestamp to AccumulatingOutputReceiver to preserve event-time downstream. Initializes the ExecutorService thread pool only when useThreadPool is true. Refactors state filtering to use finishedElementIds instead of finishedItems, preventing duplicate processing (apache#38529) * Changed outputs from List type to ConcurrentLinkedQueue to prevent lock contention (apache#38529) * Reverted outputs back to list type. Takes original Bounded Window as parameter without Hardcoding GlobalWindow.Instance (apache#38529) * Refactor AsyncDoFn to simplify output receiver and fix type mismatches involving inputTimestamp(apache#38529) * Rename AsyncDoFn to PerKeyConcurrentDoFn to address naming feedback. Might also need to rename python sdk too. * Final Formatting fixes. (apache#38529) * Revert class, file, test naming back to AsyncDoFn. Extracted verbose ArgumentProvider code block into util functions. Refactored verbose logging into util functions. Replaced hardcoded numbers with static constants. Updated test helper methods with @VisibleForTesting (apache#38529) * Rename class and files to AsyncWrapper and clean up old duplicates (apache#38529) * Refactor verbose logging to a single logInfo utility function (apache#38529)
fixes #38529
R: @AMOOOMA
Overview
AsyncWrapperacts as an execution wrapper around a standard synchronousDoFn, offloading element processing to a background thread pool. Decoupling the runner's event loop (main thread) from high-latency, I/O-heavy element processing (background threads) prevents synchronous blocking, implements backpressure, and significantly increases pipeline throughput.AsyncWrapperinstance generates a unique UUID upon instantiation to keep static JVM registries completely isolated.idFnto extract anelementId. If theelementIdis already present in the localactiveElementsMap in JVM memory, scheduling is skipped to enforce exactly-once execution grouping.maxItemsToBuffer.maxWaitTime/ 500ms).timeout(default 1s), the main thread stops scheduling the task. The element is written directly to persistent storage (BagState) and aTimeris registered to process it later.When capacity is available, the main thread performs the following steps sequentially:
CompletableFutureand submits it to the JVM's task queue.elementIdand its future in theactiveElementsMap and increments theitemsInBuffercounter.BagState(ensuring durability if a worker crashes).Timercallback to manage future reconciliation.DoFn(startBundleprocessElementfinishBundle).AccumulatingOutputReceiverheld in JVM memory, ensuring background threads do not write downstream unsafely. On completion, the future returns the output list anditemsInBufferis decremented.When the
Timerfires for Key K, the main thread executes a synchronous reconciliation cycle:BagStatefor Key K is empty, it exits immediately to free up CPU.BagState:activeElements. It is immediately rescheduled into the JVM thread pool.activeElementsand completed. The main thread retrieves the outputs from the future, emits them downstream, and removes the element fromBagStateand JVM memory.BagStateuntouched.Testing & Verification
AsyncWrapperTest.java): Includes 13 unit tests covering all concurrency and logical paths, including:./gradlew :sdks:java:core:test --tests "org.apache.beam.sdk.transforms.AsyncWrapperTest"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.