Skip to content

Implement Asynchronous wrapper for DoFn in Java SDK#38609

Merged
damccorm merged 17 commits into
apache:masterfrom
tejasiyer-dev:add-async-dofn-wrapper
Jun 11, 2026
Merged

Implement Asynchronous wrapper for DoFn in Java SDK#38609
damccorm merged 17 commits into
apache:masterfrom
tejasiyer-dev:add-async-dofn-wrapper

Conversation

@tejasiyer-dev

@tejasiyer-dev tejasiyer-dev commented May 22, 2026

Copy link
Copy Markdown
Contributor

fixes #38529

R: @AMOOOMA

Overview

AsyncWrapper acts as an execution wrapper around a standard synchronous DoFn, 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.

  1. Ingestion & Local Deduplication (Main Thread)
    • JVM Isolation: Every AsyncWrapper instance generates a unique UUID upon instantiation to keep static JVM registries completely isolated.
    • Deduplication Boundary: Incoming elements pass through an idFn to extract an elementId. If the elementId is already present in the local activeElements Map in JVM memory, scheduling is skipped to enforce exactly-once execution grouping.
  2. Backpressure & Capacity Check (Main Thread)
    • Capacity Management: The main thread checks if the background pool's active task count is below maxItemsToBuffer.
    • Exponential Backoff: If the pool is full, the main thread sleeps using exponential backoff (starting at 10ms, doubling, capped at maxWaitTime / 500ms).
    • Timeout Handling: If capacity doesn't clear within timeout (default 1s), the main thread stops scheduling the task. The element is written directly to persistent storage (BagState) and a Timer is registered to process it later.
  3. Task Creation & Durable State Writing (Main Thread)
    When capacity is available, the main thread performs the following steps sequentially:
    • Task Creation: Wraps the element logic inside a CompletableFuture and submits it to the JVM's task queue.
    • In-Memory Tracking: Registers the elementId and its future in the activeElements Map and increments the itemsInBuffer counter.
    • Durable State Write: Writes the element to the Runner's persistent BagState (ensuring durability if a worker crashes).
    • Timer Scheduling: Schedules/updates a key-scoped Timer callback to manage future reconciliation.
  4. Background Execution (Background Worker Threads)
    • Decoupled Processing: Worker threads independently pull tasks from the JVM queue.
    • Bundle Lifecycle: The thread executes the full synchronous bundle lifecycle of the wrapped DoFn (startBundle $\rightarrow$ processElement $\rightarrow$ finishBundle).
    • Thread-Safe Accumulation: Workers append outputs to a private AccumulatingOutputReceiver held in JVM memory, ensuring background threads do not write downstream unsafely. On completion, the future returns the output list and itemsInBuffer is decremented.
  5. Timer Reconciliation & Cleanup (Main Thread)
    When the Timer fires for Key K, the main thread executes a synchronous reconciliation cycle:
    • Early Exit: If BagState for Key K is empty, it exits immediately to free up CPU.
    • State Reconciliation: Iterates through the elements listed in BagState:
      • Case 1 (Lost Task): The task is not in activeElements. It is immediately rescheduled into the JVM thread pool.
      • Case 2 (Completed Task): The task is in activeElements and completed. The main thread retrieves the outputs from the future, emits them downstream, and removes the element from BagState and JVM memory.
      • Case 3 (In-Flight Task): The task is still running. The main thread leaves it in BagState untouched.
    • Timer Reset: If any elements remain unfinished, a new timer is scheduled for the next check cycle.

Testing & Verification

  • Comprehensive Unit Tests (AsyncWrapperTest.java): Includes 13 unit tests covering all concurrency and logical paths, including:
    • Basic offloaded element execution and output propagation.
    • Concurrency limits, backpressure exponential backoffs, and timeout state fallbacks.
    • Key-scoped exactly-once ingestion deduplication.
    • In-flight task cancellation on runner rollback/state deletion.
    • Lost-task rescheduling and timer reconciliation cycles.
  • Local Verification: All tests compile and pass successfully under the Core Java SDK pre-commits:
    ./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:

  • Mention the appropriate issue in your description (for example: 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, comment fixes #<ISSUE NUMBER> instead.
  • Update CHANGES.md with noteworthy changes.
  • If this contribution is large, please file an Apache Individual Contributor License Agreement.

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)

Build python source distribution and wheels
Python tests
Java tests
Go tests

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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 DoFn in the Apache Beam Java SDK. By decoupling high-latency, I/O-heavy processing from the main event loop, it aims to significantly increase pipeline throughput and prevent synchronous blocking. The implementation includes robust backpressure mechanisms, durable state management, and comprehensive testing to ensure reliable execution in distributed environments.

Highlights

  • Asynchronous Execution Wrapper: Introduced AsyncDoFn to wrap synchronous DoFn instances, offloading element processing to a background thread pool to prevent blocking the main runner thread.
  • Backpressure and Capacity Management: Implemented capacity management with exponential backoff and timeout handling to ensure the background thread pool is not overwhelmed, improving pipeline stability.
  • Durable State and Timer Reconciliation: Integrated BagState and @Timer to ensure exactly-once processing, handle task durability, and manage reconciliation of in-flight tasks.
  • JVM Isolation and Resource Management: Added UUID-based isolation for static registries and thread pools to ensure safe resource reuse across serialized worker instances.
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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@github-actions github-actions Bot added the java label May 22, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
@github-actions

Copy link
Copy Markdown
Contributor

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
Comment thread sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncDoFnTest.java Outdated
@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java Outdated
@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

codecov Bot commented Jun 5, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 74.62687% with 153 lines in your changes missing coverage. Please review.
✅ Project coverage is 57.81%. Comparing base (d9117f4) to head (d495262).
⚠️ Report is 27 commits behind head on master.

Files with missing lines Patch % Lines
...java/org/apache/beam/sdk/transforms/AsyncDoFn.java 75.00% 65 Missing and 13 partials ⚠️
...ache/beam/sdk/transforms/PerKeyConcurrentDoFn.java 74.22% 62 Missing and 13 partials ⚠️
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     
Flag Coverage Δ
java 64.25% <74.62%> (+0.07%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@tejasiyer-dev

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

…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.
… 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)
@tejasiyer-dev tejasiyer-dev force-pushed the add-async-dofn-wrapper branch from 17e66d1 to 7e0ebc5 Compare June 8, 2026 17:04
}
}

if (verboseLogging) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java Outdated
@AMOOOMA

AMOOOMA commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

Thanks! Overall LGTM, once the logging changes are in we will tag Danny for final review.

@AMOOOMA

AMOOOMA commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

I also just realized you have files duplicated and the naming for the class should match AsyncWrapper instead of calling AsyncDoFn

@AMOOOMA

AMOOOMA commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

Thanks! Tagging Danny for the final check R: @damccorm

@github-actions

Copy link
Copy Markdown
Contributor

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

@damccorm damccorm left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@damccorm

Copy link
Copy Markdown
Contributor

Rerunning some workflows to confirm this doesn't break them before merging

@damccorm damccorm merged commit afbbc31 into apache:master Jun 11, 2026
29 of 36 checks passed
@tejasiyer-dev tejasiyer-dev deleted the add-async-dofn-wrapper branch June 22, 2026 20:48
ash6898 pushed a commit to ash6898/beam that referenced this pull request Jun 29, 2026
* 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request]: Implement AsyncWrapper for DoFn in Java

3 participants