You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "tejasiyer-dev (via GitHub)" <gi...@apache.org> on 2026/05/22 22:23:11 UTC
[PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
tejasiyer-dev opened a new pull request, #38609:
URL: https://github.com/apache/beam/pull/38609
fixes #38529
R: @AMOOOMA
This PR introduces AsyncDoFn and AsyncDoFnTest to the Apache Beam Java SDK.
AsyncDoFn 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 AsyncDoFn 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:
1. Task Creation: Wraps the element logic inside a CompletableFuture and submits it to the JVM's task queue.
2. In-Memory Tracking: Registers the elementId and its future in the activeElements Map and increments the itemsInBuffer counter.
3. Durable State Write: Writes the element to the Runner's persistent BagState (ensuring durability if a worker crashes).
4. 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 $\to$ processElement $\to$ 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.
- Orphan Cleanup: Scans the activeElements Map and cancels/deletes any running JVM task belonging to Key K whose ID is not in BagState (e.g., orphans from failed bundle retries).
- State Reconciliation: Iterates through the elements listed in BagState:
1. Case 1 (Lost Task): The task is not in activeElements. It is immediately rescheduled into the JVM thread pool.
2. 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.
3. 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.
------------------------
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](https://www.apache.org/licenses/icla.pdf).
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://github.com/apache/beam/blob/master/CONTRIBUTING.md#make-the-reviewers-job-easier).
To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "AMOOOMA (via GitHub)" <gi...@apache.org>.
AMOOOMA commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4673219295
Thanks! Overall LGTM, once the logging changes are in we will tag Danny for final review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3326269830
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
Review Comment:
The `AccumulatingOutputReceiver` currently uses `Instant.now()` for output timestamps in `builder(T value)`, and defaults to `null` (which becomes the timer's processing-time firing timestamp) when `syncFn` calls `output(T output)`. This corrupts the event-time timestamps of the elements, shifting them to the processing time of the background thread or timer firing, which will break downstream windowing and watermarks.
To preserve event-time correctness, we should pass the input element's original timestamp to `AccumulatingOutputReceiver` and use it as the default timestamp for any outputs.
```java
private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> { private final List<TimestampedOutput<T>> outputs = Collections.synchronizedList(new ArrayList<>()); private final Instant inputTimestamp; AccumulatingOutputReceiver(Instant inputTimestamp) { this.inputTimestamp = inputTimestamp; } @Override public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) { return org.apache.beam.sdk.values.WindowedValues.<T>builder() .setValue(value) .setTimestamp(inputTimestamp) .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE)) .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING) .setReceiver( windowedValue -> outputs.add( new TimestampedOutput<>( windowedValue.getValue(), windowedValue.getTimestamp()))); } @Override public void output(T output) { outp
uts.add(new TimestampedOutput<>(output, inputTimestamp)); } @Override public void outputWithTimestamp(T output, Instant timestamp) { outputs.add(new TimestampedOutput<>(output, timestamp)); } public List<T> getOutputs() { List<T> rawOutputs = new ArrayList<>(); for (TimestampedOutput<T> out : outputs) { rawOutputs.add(out.value); } return rawOutputs; } public List<TimestampedOutput<T>> getTimestampedOutputs() { return outputs; } }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
The `lock` is a static `ReentrantLock` shared across all instances of `AsyncDoFn` in the entire JVM. Since `scheduleIfRoom` (called for every element in `processElement`) and `commitFinishedItems` (called on every timer firing) both acquire this lock, this will cause severe lock contention and serialize processing across all threads and all `AsyncDoFn` steps in the pipeline. This completely defeats the purpose of asynchronous processing and high throughput.
We should use per-uuid locks (e.g., stored in a static `ConcurrentHashMap<String, ReentrantLock>`) to avoid JVM-wide contention.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
Review Comment:
In `processElement`, the timer is set to `nextTimeToFire(element.getKey())` for every incoming element. If elements for a key arrive continuously at a rate faster than the timer frequency, the timer will be repeatedly overwritten and pushed further into the future. This causes starvation, as the timer may never fire, and elements in `toProcessState` will be delayed indefinitely.
To prevent this, we should use a `ValueState<Boolean>` to track if a timer is already active for the key, and only set the timer if one is not already pending.
```java
@ProcessElement public void processElement( ProcessContext c, BoundedWindow window, @StateId("to_process") BagState<KV<K, InputT>> toProcessState, @StateId("timer_set") ValueState<Boolean> timerSetState, @TimerId("timer") Timer timer) { KV<K, InputT> element = c.element(); scheduleItem(element, window, c.timestamp()); toProcessState.add(element); Boolean isSet = timerSetState.read(); if (isSet == null || !isSet) { Instant timeToFire = nextTimeToFire(element.getKey()); timer.set(timeToFire); timerSetState.write(true); } }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
Review Comment:
If `useThreadPool` is false, we still create a fixed thread pool of size `parallelism` and store it in `pool`. This is wasteful and can lead to resource leaks. We should only create the thread pool if `useThreadPool` is true.
```java
try { if (useThreadPool) { pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism)); } processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>()); itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0)); refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet(); } finally {
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
Review Comment:
In `commitFinishedItems`, when a background task fails, `inFlight.future.get()` throws an exception, and we immediately throw a `RuntimeException`. However, any successfully completed sibling tasks in the same batch have already been removed from `activeElements` (via `activeElements.remove(elementId)`).
Since the runner retries the entire timer on failure, and `toProcessState` is transactional (so the removals are rolled back in the runner's state), those successfully completed elements will still be in `toProcessState` on retry. But because they were already removed from the non-transactional in-memory `activeElements` map, they will be treated as "Lost Tasks" and rescheduled/re-run!
To avoid re-running successfully completed tasks on retry, we should keep them in `activeElements` until the entire batch succeeds, and only remove them from `activeElements` at the very end of `commitFinishedItems`. If a task fails, we should remove only the failed task from `activeElements` (so it can be retried) and throw the exception.
```java
if (activeElements.containsKey(elementId)) { InFlightElement<OutputT> inFlight = activeElements.get(elementId); if (inFlight.future.isDone()) { try { if (!inFlight.future.isCancelled()) { toReturn.add(inFlight.future.get()); } finishedItems.add(element); finishedElementIds.add(elementId); itemsFinished++; } catch (Exception e) { activeElements.remove(elementId); LOG.error("Error executing async task for element {}", element, e); throw new RuntimeException("Error executing async task for element " + element, e); } } else { inFlightElementIds.add(elementId); itemsNotYetFinished++; } }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ if (!finishedItems.contains(element)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
Review Comment:
In `commitFinishedItems`, `toProcessState` is updated by checking `!finishedItems.contains(element)`. Since `finishedItems` is a `Set<KV<K, InputT>>`, this relies on `InputT` correctly implementing `equals` and `hashCode`. If `InputT` is a byte array or a complex object that does not override `equals`, this check will fail, leading to duplicate processing or state corruption.
We should instead use `finishedElementIds.contains(elementId)` (using the user-provided `idFn` which is guaranteed to uniquely identify elements) to filter out finished elements.
```java
for (KV<K, InputT> element : stateList) { Object elementId = idFn.apply(element.getValue()); if (!finishedElementIds.contains(elementId)) { toProcessState.add(element); itemsInProcessingState++; } }
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3327177364
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
Using a single global static `ReentrantLock` across all instances of `AsyncDoFn` in the JVM will cause severe thread contention and serialize execution across all pipelines, keys, and steps using this DoFn. This completely defeats the purpose of asynchronous execution. We should use a per-UUID lock instead to isolate different instances of `AsyncDoFn`.
```java
private static final ConcurrentHashMap<String, ReentrantLock> locks =
new ConcurrentHashMap<>();
private ReentrantLock getLock() {
return locks.computeIfAbsent(uuid, k -> new ReentrantLock());
}
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
Review Comment:
The `processingElements` map is currently flat per `uuid` (`ConcurrentHashMap<Object, InFlightElement<?>>`). In `commitFinishedItems`, the code iterates over *all* active elements across *all* keys for this `uuid` to find elements to cancel or reconcile. Under high load with many keys, this $O(N_{\text{total}})$ scan will cause extremely high CPU utilization and degrade performance. Partitioning `processingElements` by key avoids this scan entirely.
```suggestion
private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, ConcurrentHashMap<Object, InFlightElement<?>>>>
processingElements = new ConcurrentHashMap<>();
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs = new ArrayList<>();
+ private final Instant inputTimestamp;
+ private final BoundedWindow window;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp, BoundedWindow window) {
+ this.inputTimestamp = inputTimestamp;
+ this.window = window;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(window))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp, window);
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
Review Comment:
With `activeElements` now partitioned by key, we no longer need to scan all active elements across all keys to find elements to cancel. We can directly iterate over the elements for the current key, reducing the complexity from $O(N_{\text{total}})$ to $O(N_{\text{key}})$.
```java
List<Object> toCancelIds = new ArrayList<>();
for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
if (!stateElementIds.contains(entry.getKey())) {
toCancelIds.add(entry.getKey());
}
}
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs = new ArrayList<>();
+ private final Instant inputTimestamp;
+ private final BoundedWindow window;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp, BoundedWindow window) {
+ this.inputTimestamp = inputTimestamp;
+ this.window = window;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(window))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
Review Comment:
Update `getProcessingElements` to accept a `key` parameter and return the partitioned map for that key. This works in tandem with partitioning `processingElements` by key to avoid $O(N)$ scans.
```java
@SuppressWarnings("unchecked")
private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements(Object key) {
ConcurrentHashMap<Object, ConcurrentHashMap<Object, InFlightElement<?>>> uuidElements =
(ConcurrentHashMap<Object, ConcurrentHashMap<Object, InFlightElement<?>>>) (ConcurrentHashMap<?, ?>) processingElements.get(uuid);
if (uuidElements == null) {
throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
}
Object safeKey = (key == null) ? "NULL_KEY" : key;
return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>)
uuidElements.computeIfAbsent(safeKey, k -> new ConcurrentHashMap<>());
}
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs = new ArrayList<>();
+ private final Instant inputTimestamp;
+ private final BoundedWindow window;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp, BoundedWindow window) {
+ this.inputTimestamp = inputTimestamp;
+ this.window = window;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(window))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp, window);
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
Review Comment:
The `nextTimeToFire` method uses floating-point math (`double`) and divisions/multiplications by `1000.0` to calculate the next firing time. This can be simplified to pure integer/long math, which is more efficient, less prone to rounding errors, and much easier to read.
```java
private Instant nextTimeToFire(@Nullable K key) {
long seed = (key == null) ? 0 : key.hashCode();
long frequencyMs = timerFrequency.getMillis();
long now = System.currentTimeMillis();
long base = ((now + frequencyMs) / frequencyMs) * frequencyMs;
long offset = Math.abs(seed % 1000000) * frequencyMs / 1000000;
return new Instant(base + offset);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321304173
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
Review Comment:
@gemini-code-assist The manual absolute time calculation in nextTimeToFire is an intentional design choice to maintain parity with the Python SDK implementation in async_dofn.py. Python calculates the absolute timestamp using system time (time()) and sets it absolutely (timer.set()). By using System.currentTimeMillis() in Java, you are matching Python's behavior 1-1.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3319377748
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
@gemini-code-assist is this issue still there in the most recent commit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "AMOOOMA (via GitHub)" <gi...@apache.org>.
AMOOOMA commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3390623266
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,796 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+ private static final double HASH_MODULO_LIMIT = 1000000.0;
+ private static final double MS_PER_SEC = 1000.0;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ // Setup is called by the runner exactly once on each worker node when this DoFn is initialized.
+ // It is responsible for setting up the wrapped synchronous DoFn
+ // and initializing the shared JVM-wide thread pool and registries.
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ logDuplicateElement(element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ bundleArgProvider(receiver);
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ processArgProvider(element, window, timestamp, receiver);
+
+ invoker.invokeStartBundle(bundleArgProvider);
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ // Add element to active elements map and increment buffer counter
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider(
+ AccumulatingOutputReceiver<OutputT> receiver) {
+ return new BundleArgProvider(receiver);
+ }
+
+ private DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ OutputReceiver<OutputT> receiver) {
+ return new ProcessArgProvider(element, window, timestamp, receiver);
+ }
+
+ // Named BaseArgumentProvider supplying bundle-level lifecycle context to the invoker.
+ private class BundleArgProvider extends DoFnInvoker.BaseArgumentProvider<InputT, OutputT> {
+ private final AccumulatingOutputReceiver<OutputT> receiver;
+
+ BundleArgProvider(AccumulatingOutputReceiver<OutputT> receiver) {
+ this.receiver = receiver;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return new AsyncFinishBundleContext(doFn, receiver);
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ }
+
+ // FinishBundleContext subclass bound to the enclosing DoFn instance to prevent compiler crashes.
+ private class AsyncFinishBundleContext extends DoFn<InputT, OutputT>.FinishBundleContext {
+ private final AccumulatingOutputReceiver<OutputT> receiver;
+
+ AsyncFinishBundleContext(
+ DoFn<InputT, OutputT> doFn, AccumulatingOutputReceiver<OutputT> receiver) {
+ doFn.super();
+ this.receiver = receiver;
+ }
+
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public void output(OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ }
+
+ // BaseArgumentProvider supplying element-level context to the invoker.
+ private class ProcessArgProvider extends DoFnInvoker.BaseArgumentProvider<InputT, OutputT> {
+ private final KV<K, InputT> element;
+ private final BoundedWindow window;
+ private final Instant timestamp;
+ private final OutputReceiver<OutputT> receiver;
+
+ ProcessArgProvider(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ OutputReceiver<OutputT> receiver) {
+ this.element = element;
+ this.window = window;
+ this.timestamp = timestamp;
+ this.receiver = receiver;
+ }
+
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ }
+
+ // Schedule an element to the thread pool, retries with backoff if the buffer is full.
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ logBackpressure(element, sleep, totalSleep);
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % (long) HASH_MODULO_LIMIT) / HASH_MODULO_LIMIT;
+ double timerFrequencySec = timerFrequency.getMillis() / MS_PER_SEC;
+ double nowSec = System.currentTimeMillis() / MS_PER_SEC;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * MS_PER_SEC));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ K key = null;
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ logProcessingTimer(key);
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ List<List<OutputT>> toReturn = new ArrayList<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ logRescheduling(element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (!finishedElementIds.contains(elementId)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ // Outputs use processing-time timestamps matching Python behavior
+ for (List<OutputT> outputs : toReturn) {
+ for (OutputT out : outputs) {
+ receiver.output(out);
+ }
+ }
+
+ logFinishedItems(itemsFinished, itemsNotYetFinished, itemsRescheduled, itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ private void logProcessingTimer(@Nullable K key) {
Review Comment:
Sry i probably should've explained more earlier, I was thinking having one util log function that's something like this
```
void logInfo(String s) {
if (verboseLogging) { LOG.info(s); }
}
```
So we would be able to cut down the duplicate ifs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291622499
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
@gemini-code-assist how would you best fix this issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4579831509
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4636464547
## [Codecov](https://app.codecov.io/gh/apache/beam/pull/38609?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `74.22680%` with `75 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 58.08%. Comparing base ([`1ca1faf`](https://app.codecov.io/gh/apache/beam/commit/1ca1fafcd62b0e3f1bd6a071842cab7566366d1c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`17e66d1`](https://app.codecov.io/gh/apache/beam/commit/17e66d1e05cbd2c95942ba09ae63db046c24c49f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 140 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/beam/pull/38609?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...java/org/apache/beam/sdk/transforms/AsyncDoFn.java](https://app.codecov.io/gh/apache/beam/pull/38609?src=pr&el=tree&filepath=sdks%2Fjava%2Fcore%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fbeam%2Fsdk%2Ftransforms%2FAsyncDoFn.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c2Rrcy9qYXZhL2NvcmUvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2JlYW0vc2RrL3RyYW5zZm9ybXMvQXN5bmNEb0ZuLmphdmE=) | 74.22% | [62 Missing and 13 partials :warning: ](https://app.codecov.io/gh/apache/beam/pull/38609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #38609 +/- ##
============================================
- Coverage 58.73% 58.08% -0.65%
- Complexity 7418 13016 +5598
============================================
Files 1649 2513 +864
Lines 212369 262612 +50243
Branches 6330 10709 +4379
============================================
+ Hits 124727 152541 +27814
- Misses 83095 104350 +21255
- Partials 4547 5721 +1174
```
| [Flag](https://app.codecov.io/gh/apache/beam/pull/38609/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [java](https://app.codecov.io/gh/apache/beam/pull/38609/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `64.19% <74.22%> (-8.31%)` | :arrow_down: |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/beam/pull/38609?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321171203
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
Review Comment:
While maintaining parity with the Python SDK is a valid goal, blocking the runner's timer thread in the Java SDK is a critical anti-pattern. The timer thread is typically shared across multiple bundles and keys; blocking it can lead to pipeline-wide stalls, watermark delays, and potential deadlocks if the background thread pool tasks depend on the timer thread's progress. Instead of `Thread.sleep`, you should implement an exponential backoff by rescheduling the timer to fire again in the future. This allows the timer thread to continue processing other work while waiting for buffer capacity.
```suggestion
// Reschedule missing elements without blocking the timer thread
for (KV<K, InputT> element : toReschedule) {
scheduleIfRoom(element, GlobalWindow.INSTANCE, fireTimestamp, false);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3319377748
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
@gemini-code-assist is this issue still there in the most recent commit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291591501
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
Review Comment:
@gemini-code-assist I pushed a fix for this in my latest commit. Does this resolve the issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291590212
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
Review Comment:
@gemini-code-assist
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291639713
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
@AMOOOMA
I kept a single static class lock because it matches the exact synchronization model used by Python's AsyncWrapper (_lock = RLock()).
If you prefer, we can change this to a per-transform lock or leverage ConcurrentHashMap to reduce global thread contention, but I wanted to keep it aligned with the Python SDK first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3319474440
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
Review Comment:
### Critical Concurrency Issue: Shared `syncFn` Instance Across Background Threads\n\nThe `syncFn` instance is shared globally across all background threads in the executor pool. If `parallelism > 1`, multiple threads will concurrently invoke `processElement`, `startBundle`, and `finishBundle` on the exact same `syncFn` instance.\n\nSince standard `DoFn`s in Apache Beam are **not** thread-safe and often maintain state (especially during bundle lifecycle methods), this will lead to severe race conditions, state corruption, and unpredictable runtime failures.\n\n#### Recommended Solution:\nUse a `ThreadLocal` to clone `syncFn` per thread using `SerializableUtils.clone(syncFn)`. This ensures each thread has its own isolated, thread-safe instance of the `DoFn`.\n\n```java\nprivate transient ThreadLocal<DoFn<InputT, OutputT>> localSyncFn;\nprivate transient ConcurrentHashMap.KeySetView<DoFnInvoker<InputT, OutputT>, Boolean> localInvokers;\n```\n\nIn `@Setup`, initialize the `ThreadLoca
l` and track the cloned invokers:\n```java\nthis.localInvokers = ConcurrentHashMap.newKeySet();\nthis.localSyncFn = ThreadLocal.withInitial(() -> {\n DoFn<InputT, OutputT> clone = SerializableUtils.clone(syncFn);\n DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(clone);\n localInvokers.add(invoker);\n invoker.invokeSetup(new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {\n @Override\n public PipelineOptions pipelineOptions() {\n return options;\n }\n @Override\n public String getErrorContext() {\n return \"AsyncDoFn/Setup\";\n }\n });\n return clone;\n});\n```\n\nIn `@Teardown`, invoke teardown on all cloned invokers:\n```java\nif (localInvokers != null) {\n for (DoFnInvoker<InputT, OutputT> invoker : localInvokers) {\n try {\n invoker.invokeTeardown();\n } catch (Exception e) {\n LOG.warn(\"Error tearing down cloned DoFn\", e);\n }\n }\n}\n```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
Review Comment:
Declare a static reference counter map to track active instances of `AsyncDoFn` per `uuid` for safe teardown.
```java
private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =\n new ConcurrentHashMap<>();\n private static final ConcurrentHashMap<String, AtomicInteger> refCounts =\n new ConcurrentHashMap<>();
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
Review Comment:
Increment the reference count for the current `uuid` during `@Setup` to ensure safe teardown.
```java
lock.lock();\n try {\n pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));\n processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());\n itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));\n refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();\n } finally {\n lock.unlock();\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ if (!finishedItems.contains(element)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ for (List<TimestampedOutput<OutputT>> outputs : toReturn) {
+ for (TimestampedOutput<OutputT> out : outputs) {
+ if (out.timestamp != null) {
+ receiver.outputWithTimestamp(out.value, out.timestamp);
+ } else {
+ receiver.output(out.value);
+ }
+ }
+ }
+
+ LOG.info(
+ "Items finished: {}, not yet finished: {}, " + "rescheduled: {}, in processing state: {}",
+ itemsFinished,
+ itemsNotYetFinished,
+ itemsRescheduled,
+ itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ // Package-private helper methods for testing direct execution without Pipeline / ProcessContext
+ // boilerplate
+ void processDirect(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer) {
+ scheduleItem(element, window, timestamp);
+ toProcessState.add(element);
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ List<OutputT> commitFinishedItemsDirect(
+ Instant fireTimestamp, BagState<KV<K, InputT>> toProcessState, Timer timer) {
+ AccumulatingOutputReceiver<OutputT> receiver = new AccumulatingOutputReceiver<>();
+ commitFinishedItems(fireTimestamp, toProcessState, timer, receiver);
+ return receiver.getOutputs();
+ }
+
+ boolean isEmpty() {
+ return getItemsInBuffer().get() == 0;
+ }
+
+ int getItemsInBufferCount() {
+ return getItemsInBuffer().get();
+ }
+
+ static void resetState() {
+ lock.lock();
+ try {
+ for (Map.Entry<String, ExecutorService> entry : pool.entrySet()) {
+ entry.getValue().shutdownNow();
+ }
+ pool.clear();
+ processingElements.clear();
+ itemsInBuffer.clear();
+ } finally {
+ lock.unlock();
+ }
+ }
Review Comment:
Clear the static `refCounts` map when resetting state.
```java
static void resetState() {\n lock.lock();\n try {\n for (Map.Entry<String, ExecutorService> entry : pool.entrySet()) {\n entry.getValue().shutdownNow();\n }\n pool.clear();\n processingElements.clear();\n itemsInBuffer.clear();\n refCounts.clear();\n } finally {\n lock.unlock();\n }\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
Review Comment:
### High Severity Bug: Shared Resource Teardown Race Condition\n\nSince runners clone `DoFn` instances on the same worker node, multiple active instances of `AsyncDoFn` with the same `uuid` will share the same `ExecutorService` and static registries.\n\nWhen any of these cloned instances is torn down, its `@Teardown` method is called, which immediately removes the shared `threadPool` from the static map and shuts it down. This will cause all other active instances sharing the same `uuid` to fail with `IllegalStateException` when they attempt to access the thread pool or processing elements map.\n\n#### Solution:\nImplement reference counting to ensure that the shared resources are only removed and shut down when the last active instance of `AsyncDoFn` for a given `uuid` is torn down.
```java
ExecutorService threadPool = null;\n lock.lock();\n try {\n AtomicInteger refCount = refCounts.get(uuid);\n if (refCount != null && refCount.decrementAndGet() == 0) {\n refCounts.remove(uuid);\n threadPool = pool.remove(uuid);\n processingElements.remove(uuid);\n itemsInBuffer.remove(uuid);\n }\n } finally {\n lock.unlock();\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ }
+ } else {
Review Comment:
### Critical Correctness Issue: Silent Data Loss on Exception\n\nWhen an exception is thrown during the execution of the async task, it is caught, logged, and then the element is added to `finishedItems` and removed from the persistent state (`toProcessState`).\n\nThis causes **silent data loss** because the failed element is permanently discarded and the pipeline continues as if it succeeded. In Apache Beam, exceptions should be propagated to fail the bundle/timer so that the runner can retry the execution.
```java
} catch (Exception e) {\n LOG.error(\"Error executing async task for element {}\", element, e);\n throw new RuntimeException(\"Error executing async task for element \" + element, e);\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
Review Comment:
### High Severity Bug: Division by Zero / Infinite Loop\n\nIf `timerFrequency` is set to `Duration.ZERO` or a negative duration, `timerFrequencySec` will be `0.0` or negative. This will cause a division by zero in `nextTimeToFire`, resulting in `Double.NaN` or `Double.POSITIVE_INFINITY`, which ultimately sets the timer to `Instant.ofEpochMilli(0)`. This causes the timer to fire immediately and repeatedly in an infinite loop, hanging the pipeline.\n\nWe should validate that `timerFrequency` is strictly positive in the constructor.
```java
this.parallelism = parallelism;\n if (timerFrequency.getMillis() <= 0) {\n throw new IllegalArgumentException(\"timerFrequency must be greater than zero\");\n }\n this.timerFrequency = timerFrequency;
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
Review Comment:
### Medium Severity: Document Limitations and Requirements in Javadoc\n\nThe class Javadoc should clearly document the requirements and limitations of `AsyncDoFn` to prevent users from encountering hard-to-debug runtime issues:\n1. **Thread-Safety**: The wrapped `DoFn` must be thread-safe because its methods are invoked concurrently by multiple background threads.\n2. **No Multi-Output Support**: Tagged outputs and `MultiOutputReceiver` are not supported.\n3. **Bundle Lifecycle**: `startBundle` and `finishBundle` are invoked per element, so any batching or aggregation logic in them will not behave as expected.
```java
/**\n * Class that wraps a {@link DoFn} and converts it from one which processes elements synchronously to one\n * which processes them asynchronously.\n *\n * <p>For synchronous DoFns, the default settings mean that many (100s) of elements will be processed\n * in parallel and that processing an element will block all other work on that key. In addition,\n * runners are optimized for latencies less than a few seconds, and longer operations can result in\n * high retry rates. Async should be considered when the default parallelism is not correct and/or\n * items are expected to take longer than a few seconds to process.\n *\n * <h3>Limitations & Requirements:</h3>\n * <ul>\n * <li><b>Thread-Safety:</b> The wrapped {@code DoFn} must be thread-safe because its methods\n * (including {@code processElement}) will be invoked concurrently by multiple background threads.</li>\n * <li><b>No Multi-Output Support:</b> Tagged outputs and {@code MultiOutputReceiver} are not supporte
d.\n * Attempting to use them will result in an {@link UnsupportedOperationException}.</li>\n * <li><b>Bundle Lifecycle:</b> {@code startBundle} and {@code finishBundle} are invoked per element\n * within the background tasks, meaning any batching or aggregation logic implemented in them\n * will not function as expected.</li>\n * </ul>\n */
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "AMOOOMA (via GitHub)" <gi...@apache.org>.
AMOOOMA commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4675748393
Thanks! Tagging Danny for the final check R: @damccorm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4638040604
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3384315258
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
Review Comment:
Moved these into private util functions to clean up the scheduling logic in most recent commit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321261380
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
Review Comment:
@gemini-code-assist Intentionally maintaining parity with the Python SDK's tracking mechanism (async_dofn.py), which also uses a standard dictionary without an expiration or eviction policy for its in-memory task tracking
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4654725858
Ran end-to-end tests in google3 using Dataflow Runner and successfully verified stable streaming execution
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3384287638
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (!finishedElementIds.contains(elementId)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ // Outputs use processing-time timestamps matching Python behavior
+ for (List<OutputT> outputs : toReturn) {
+ for (OutputT out : outputs) {
+ receiver.output(out);
+ }
+ }
+
+ LOG.info(
+ "Items finished: {}, not yet finished: {}, " + "rescheduled: {}, in processing state: {}",
+ itemsFinished,
+ itemsNotYetFinished,
+ itemsRescheduled,
+ itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ // Package-private helper methods for testing direct execution without Pipeline / ProcessContext
+ // boilerplate
+ void processDirect(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer) {
+ scheduleItem(element, window, timestamp);
+ toProcessState.add(element);
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ List<OutputT> commitFinishedItemsDirect(
+ Instant fireTimestamp, BagState<KV<K, InputT>> toProcessState, Timer timer) {
+ AccumulatingOutputReceiver<OutputT> receiver = new AccumulatingOutputReceiver<>();
+ commitFinishedItems(fireTimestamp, toProcessState, timer, receiver);
+ return receiver.getOutputs();
+ }
+
+ boolean isEmpty() {
+ return getItemsInBuffer().get() == 0;
Review Comment:
getItemsInBuffer returns an Atomic Integer which I don't think has an isEmpty method.
and I'm using Atomic Integer because Java doesn't have Global Interpreter Lock like Python so using regular int for counter might cause race conditions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3383729124
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
Review Comment:
The naming isn't confirmed. The only thing Reuven mentioned was to change it to something other than AsyncDoFn. But yes we can send a rename PR later.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4568825751
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3327177403
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs = new ArrayList<>();
+ private final Instant inputTimestamp;
+ private final BoundedWindow window;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp, BoundedWindow window) {
+ this.inputTimestamp = inputTimestamp;
+ this.window = window;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(window))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp, window);
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
Review Comment:
When an asynchronous task fails, `inFlight.future.get()` throws an `ExecutionException`. Wrapping this directly in a `RuntimeException` preserves the `ExecutionException` wrapper, which makes the stack trace noisier and harder to debug. Unwrapping the cause of the `ExecutionException` provides a much cleaner and more actionable error message.
```suggestion
} catch (Exception e) {
LOG.error("Error executing async task for element {}", element, e);
Throwable cause = (e instanceof java.util.concurrent.ExecutionException) ? e.getCause() : e;
throw new RuntimeException("Error executing async task for element " + element, cause);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4523777641
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291646007
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
To address the issue of timestamp loss in `AccumulatingOutputReceiver`, you need to capture the timestamp provided by the `DoFn` during execution and propagate it when emitting the results. Since `OutputReceiver` does not natively support timestamps, you should store the output as a `TimestampedValue` container.
However, please note that `AsyncDoFn` is a `DoFn` wrapper, and in the `onTimer` method, you only have access to `OnTimerContext`, which does not provide a `ProcessContext` to emit with a timestamp. If the downstream `DoFn` expects specific windowing or timestamp behavior, this limitation might require a more fundamental change to how `AsyncDoFn` handles output emission (e.g., by using a `MultiOutputReceiver` or ensuring the runner supports timestamped output in the timer context).
If you are certain the downstream receiver supports it, you can cast the receiver to a type that supports timestamped output, but this is generally not recommended as it breaks the `OutputReceiver` abstraction. A safer approach is to ensure the `syncFn` itself handles the windowing/timestamping logic if possible, or to use a different mechanism to emit the results.
```suggestion
private static class TimestampedValue<T> {
final T value;
final Instant timestamp;
TimestampedValue(T value, Instant timestamp) {
this.value = value;
this.timestamp = timestamp;
}
}
private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
private final List<TimestampedValue<T>> outputs = Collections.synchronizedList(new ArrayList<>());
@Override
public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
return org.apache.beam.sdk.values.WindowedValues.<T>builder()
.setValue(value)
.setTimestamp(Instant.now())
.setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
.setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
.setReceiver(windowedValue -> outputs.add(new TimestampedValue<>(windowedValue.getValue(), windowedValue.getTimestamp())));
}
@Override
public void output(T output) {
outputs.add(new TimestampedValue<>(output, Instant.now()));
}
@Override
public void outputWithTimestamp(T output, Instant timestamp) {
outputs.add(new TimestampedValue<>(output, timestamp));
}
public List<TimestampedValue<T>> getOutputs() {
return outputs;
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321254612
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
Review Comment:
The author is correct that the Python SDK shares the `_sync_fn` instance. However, in the Java SDK, `DoFn` instances are generally expected to be thread-safe if accessed concurrently, or the runner is expected to provide thread-local instances.
If you want to maintain parity with the Python SDK's behavior, you should be aware that this implementation relies on the user's `DoFn` being thread-safe. If a user provides a non-thread-safe `DoFn` (e.g., one with mutable instance variables), this will cause data corruption.
To safely support non-thread-safe `DoFn`s while keeping the current design, you could consider using a `ThreadLocal` or a pool of `DoFn` instances, as suggested previously. If you choose to proceed with the shared instance, I strongly recommend adding a clear warning in the class Javadoc stating that the wrapped `DoFn` **must** be thread-safe.
```suggestion
/**
* Class that wraps a dofn and converts it from one which process elements synchronously to one
* which processes them asynchronously.
*
* <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
* in parallel and that processing an element will block all other work on that key. In addition
* runners are optimized for latencies less than a few seconds and longer operations can result in
* high retry rates. Async should be considered when the default parallelism is not correct and/or
* items are expected to take longer than a few seconds to process.
*
* <p>NOTE: 1) The wrapped syncFn **must** be thread-safe as it is invoked concurrently across
* multiple background threads. 2) Tagged output multi-outputs are unsupported. 3)
* StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
* behave as expected.
*/
public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3327032932
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
Review Comment:
@gemini-code-assist
Last review you told me the exact opposite and to implement ConcurrentLinkedQueue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3326842627
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
@gemini-code-assist This is a huge design change that greatly diverges from the Python SDK async_dofn.py. I intentionally used a single lock. The actual slow, heavy work is done outside the lock in the background thread pool. Actual lock contention is very low in practice and will not bottleneck a standard pipeline. You are highly exaggerating the practical impact saying things like it defeats the purpose of asynchronous processing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4580296504
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4682549765
Rerunning some workflows to confirm this doesn't break them before merging
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3392016457
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java:
##########
@@ -0,0 +1,796 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+ private static final double HASH_MODULO_LIMIT = 1000000.0;
+ private static final double MS_PER_SEC = 1000.0;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ // Setup is called by the runner exactly once on each worker node when this DoFn is initialized.
+ // It is responsible for setting up the wrapped synchronous DoFn
+ // and initializing the shared JVM-wide thread pool and registries.
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ logDuplicateElement(element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ bundleArgProvider(receiver);
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ processArgProvider(element, window, timestamp, receiver);
+
+ invoker.invokeStartBundle(bundleArgProvider);
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ // Add element to active elements map and increment buffer counter
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider(
+ AccumulatingOutputReceiver<OutputT> receiver) {
+ return new BundleArgProvider(receiver);
+ }
+
+ private DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ OutputReceiver<OutputT> receiver) {
+ return new ProcessArgProvider(element, window, timestamp, receiver);
+ }
+
+ // Named BaseArgumentProvider supplying bundle-level lifecycle context to the invoker.
+ private class BundleArgProvider extends DoFnInvoker.BaseArgumentProvider<InputT, OutputT> {
+ private final AccumulatingOutputReceiver<OutputT> receiver;
+
+ BundleArgProvider(AccumulatingOutputReceiver<OutputT> receiver) {
+ this.receiver = receiver;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return new AsyncFinishBundleContext(doFn, receiver);
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ }
+
+ // FinishBundleContext subclass bound to the enclosing DoFn instance to prevent compiler crashes.
+ private class AsyncFinishBundleContext extends DoFn<InputT, OutputT>.FinishBundleContext {
+ private final AccumulatingOutputReceiver<OutputT> receiver;
+
+ AsyncFinishBundleContext(
+ DoFn<InputT, OutputT> doFn, AccumulatingOutputReceiver<OutputT> receiver) {
+ doFn.super();
+ this.receiver = receiver;
+ }
+
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public void output(OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ }
+
+ // BaseArgumentProvider supplying element-level context to the invoker.
+ private class ProcessArgProvider extends DoFnInvoker.BaseArgumentProvider<InputT, OutputT> {
+ private final KV<K, InputT> element;
+ private final BoundedWindow window;
+ private final Instant timestamp;
+ private final OutputReceiver<OutputT> receiver;
+
+ ProcessArgProvider(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ OutputReceiver<OutputT> receiver) {
+ this.element = element;
+ this.window = window;
+ this.timestamp = timestamp;
+ this.receiver = receiver;
+ }
+
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ }
+
+ // Schedule an element to the thread pool, retries with backoff if the buffer is full.
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ logBackpressure(element, sleep, totalSleep);
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % (long) HASH_MODULO_LIMIT) / HASH_MODULO_LIMIT;
+ double timerFrequencySec = timerFrequency.getMillis() / MS_PER_SEC;
+ double nowSec = System.currentTimeMillis() / MS_PER_SEC;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * MS_PER_SEC));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ K key = null;
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ logProcessingTimer(key);
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ List<List<OutputT>> toReturn = new ArrayList<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ logRescheduling(element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (!finishedElementIds.contains(elementId)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ // Outputs use processing-time timestamps matching Python behavior
+ for (List<OutputT> outputs : toReturn) {
+ for (OutputT out : outputs) {
+ receiver.output(out);
+ }
+ }
+
+ logFinishedItems(itemsFinished, itemsNotYetFinished, itemsRescheduled, itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ private void logProcessingTimer(@Nullable K key) {
Review Comment:
`void logInfo(String s) {
if (verboseLogging) { LOG.info(s); }
}`
Beam's complier enforces a rule called `[Slf4jFormatShouldBeConst]` where basically the first argument of any LOG.info call has to be a string literal and the build fails if it is not.
As a workaround I could do this though:
```
private void logInfo(String s) {
if (verboseLogging) {
LOG.info("{}", s);
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321007943
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
The 'lock' is defined as a static member: 'private static final ReentrantLock lock = new ReentrantLock();'. Because this lock is static, **all** instances of 'AsyncDoFn' across the entire JVM share the same lock. Every element processed by any 'AsyncDoFn' step in the pipeline will contend for this single JVM-wide lock during 'scheduleIfRoom' and 'commitFinishedItems'. This introduces a massive performance bottleneck and defeats the purpose of asynchronous, high-throughput processing.\n\nInstead of a single static lock, use a lock per 'uuid' (e.g., stored in a static 'ConcurrentHashMap<String, ReentrantLock>') or make the lock instance-scoped if 'uuid' isolation is sufficient.
```java
private static final ConcurrentHashMap<String, ReentrantLock> locks = new ConcurrentHashMap<>();\n\n private ReentrantLock getLock() {\n return locks.computeIfAbsent(uuid, k -> new ReentrantLock());\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
Review Comment:
In 'commitFinishedItems' (which runs on the runner's timer execution thread), rescheduled elements are processed by calling 'scheduleItem'. 'scheduleItem' contains a blocking loop with 'Thread.sleep' that waits until there is room in the buffer. Blocking the runner's timer thread is extremely dangerous; it can freeze the entire pipeline, delay watermarks, and cause SDK harness timeouts/crashes.\n\nInstead of blocking, you should attempt to schedule the element once using 'scheduleIfRoom'. If there is no room, do **not** block or sleep. Simply leave the element in 'toProcessState' (by not adding it to 'finishedItems'), and let the next timer firing retry the scheduling.
```suggestion
// Reschedule missing elements without blocking the timer thread\n for (KV<K, InputT> element : toReschedule) {\n scheduleIfRoom(element, GlobalWindow.INSTANCE, fireTimestamp, false);\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
Review Comment:
The 'processingElements' map is a static 'ConcurrentHashMap' that holds active and completed futures. If a key or window is garbage collected or abandoned by the runner (e.g., due to window expiration), the timer for that key might never fire. Consequently, the corresponding elements will never be removed from 'processingElements', leading to a progressive memory leak and eventual 'OutOfMemoryError' in long-running pipelines.\n\nTo prevent this, use a Guava 'Cache' with an expiration policy (e.g., 'expireAfterWrite' of 1 hour) instead of a raw 'ConcurrentHashMap' for 'processingElements'. This ensures that abandoned or leaked futures are eventually evicted.
```java
private static final ConcurrentHashMap<String, com.google.common.cache.Cache<Object, InFlightElement<?>>>\n processingElements = new ConcurrentHashMap<>();
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
Review Comment:
The 'nextTimeToFire' method manually calculates the absolute firing time using 'System.currentTimeMillis()'. This makes the class non-deterministic and extremely difficult to test using Beam's mock clocks (like 'TestStream' or 'TestPipeline' with processing time).\n\nInstead of manual absolute time calculations, leverage Beam's idiomatic relative timer API: 'timer.offset(duration).setRelative()'. You can still calculate the jitter/stagger duration and apply it as an offset.
```java
private Duration getTimerOffset(@Nullable K key) {\n long seed = (key == null) ? 0 : key.hashCode();\n double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;\n long jitterMillis = (long) (fractionalOffset * timerFrequency.getMillis());\n return timerFrequency.plus(Duration.millis(jitterMillis));\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
Review Comment:
In the 'whenComplete' callback of the future, the global static 'lock' is acquired solely to decrement 'itemsInBuffer'. Since 'getItemsInBuffer()' returns an 'AtomicInteger', the 'decrementAndGet()' operation is already completely thread-safe, atomic, and lock-free. Acquiring a heavy JVM-wide lock inside an asynchronous completion callback is unnecessary and introduces severe thread contention.
```java
CompletableFuture<List<TimestampedOutput<OutputT>>> unused =\n future.whenComplete(\n (res, ex) -> {\n getItemsInBuffer().decrementAndGet();\n });
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
Review Comment:
The wrapped 'syncFn' is invoked concurrently across multiple background threads if 'parallelism > 1'. In Apache Beam, 'DoFn' instances are **not** thread-safe by default; the runner guarantees single-threaded access per instance. If the user's 'DoFn' has any instance variables or non-thread-safe state, concurrent execution will lead to race conditions and silent data corruption.\n\nTo ensure thread safety, you should clone the 'syncFn' per thread using a 'ThreadLocal' or a pool of deserialized instances. Since 'DoFn' is 'Serializable', you can serialize it once during 'setup' and deserialize it to create thread-local copies.
```suggestion
// Use a ThreadLocal to ensure each thread has its own isolated clone of the syncFn\n private transient ThreadLocal<DoFn<InputT, OutputT>> threadLocalSyncFn;\n\n // Note: Initialize threadLocalSyncFn in setup() by serializing/deserializing syncFn\n DoFn<InputT, OutputT> localFn = threadLocalSyncFn.get();\n DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(localFn);
```
##########
sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncDoFnTest.java:
##########
@@ -0,0 +1,945 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.ReadableState;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for verifying async processing structures and logic. */
+@RunWith(JUnit4.class)
+public class AsyncDoFnTest implements Serializable {
+
+ @Rule public final transient TestPipeline p = TestPipeline.create();
Review Comment:
The 'TestPipeline p' rule is declared but never used in any of the unit tests. Since all tests execute the 'AsyncDoFn' methods directly via helper methods ('processDirect', 'commitFinishedItemsDirect'), this unused rule should be removed to keep the test suite clean.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321262188
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
Review Comment:
While maintaining parity with the Python SDK is a valid design goal, it is important to consider the differences in memory management between the two runtimes. In Java, static `ConcurrentHashMap` entries that are never evicted will persist for the lifetime of the JVM, leading to inevitable `OutOfMemoryError` in long-running streaming pipelines. Even if the Python implementation currently lacks an eviction policy, it is safer to implement one in the Java SDK to ensure production stability. If you prefer not to use Guava, you could implement a periodic cleanup task or use a `LinkedHashMap` with access-order eviction, but some form of lifecycle management is necessary for a robust Java implementation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291594230
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+ List<KV<K, InputT>> finishedItems = new ArrayList<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+ int itemsCancelled = toCancel.size();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<K, InputT, OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ finishedItems.add(element);
+ activeElements.remove(elementId);
+ }
+ } else {
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ if (!finishedItems.contains(element)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
Review Comment:
@gemini-code-assist I pushed a fix for this in my latest commit. Does this resolve the issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4663397500
All core unit tests and correctness checks have been updated and are passing cleanly
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291594791
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+ List<KV<K, InputT>> finishedItems = new ArrayList<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+ int itemsCancelled = toCancel.size();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<K, InputT, OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ finishedItems.add(element);
+ activeElements.remove(elementId);
+ }
+ } else {
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ if (!finishedItems.contains(element)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
Review Comment:
Thank you for the update. I have reviewed the changes in your latest commit. Using a `Set` to track `elementId`s and ensuring each ID is processed only once per reconciliation cycle effectively resolves the O(N^2) complexity and the logic bug regarding duplicate elements. The implementation now correctly handles the reconciliation loop.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291622499
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
@gemini-code-assist how would you best fix this issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4566433416
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4578646669
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3327008790
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return new ArrayList<>(outputs);
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp);
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
Review Comment:
Hardcoding `GlobalWindow.INSTANCE` and using `fireTimestamp` when rescheduling missing elements loses the original window and shifts the element's timestamp forward. This can break downstream windowing and watermark tracking. Suggest storing the original window and timestamp along with the element in `BagState` (e.g., by wrapping them or using `WindowedValue`) and restoring them when rescheduling.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
Review Comment:
Since `AccumulatingOutputReceiver` is instantiated per element and accessed only by a single background thread executing that element's task, there is no concurrent write access to `outputs`. Using `ConcurrentLinkedQueue` introduces unnecessary synchronization and GC overhead. A standard `ArrayList` is safe and much more efficient. Additionally, pass the original `BoundedWindow` to the constructor to avoid hardcoding `GlobalWindow.INSTANCE` downstream.
```java
private final List<TimestampedOutput<T>> outputs = new java.util.ArrayList<>();\n private final Instant inputTimestamp;\n private final BoundedWindow window;\n\n AccumulatingOutputReceiver(Instant inputTimestamp, BoundedWindow window) {\n this.inputTimestamp = inputTimestamp;\n this.window = window;\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return new ArrayList<>(outputs);
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp);
Review Comment:
Pass the original `window` to the `AccumulatingOutputReceiver` constructor to preserve the windowing context of the element.
```java
AccumulatingOutputReceiver<OutputT> receiver =\n new AccumulatingOutputReceiver<>(timestamp, window);
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
Review Comment:
Hardcoding `GlobalWindow.INSTANCE` in the `OutputBuilder` forces all outputs into the global window, even if the input element was in a custom window. Use the passed `window` instead.
```suggestion
.setWindows(java.util.Collections.singletonList(window))
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
The static global `lock` is shared across all instances of `AsyncDoFn` in the JVM. This will cause severe thread contention under load, especially since it is acquired in `scheduleIfRoom` and `commitFinishedItems`. Suggest using a per-UUID lock stored in a static `ConcurrentHashMap` to isolate locking to instances sharing the same UUID.
```suggestion
private static final ConcurrentHashMap<String, ReentrantLock> locks = new ConcurrentHashMap<>();
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return new ArrayList<>(outputs);
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp);
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
Review Comment:
Iterating over `activeElements.entrySet()` scans all active elements across *all* keys in the JVM for every single timer firing. This is an $O(N)$ operation that will cause severe performance degradation under high key cardinality. Suggest grouping `processingElements` by key (e.g., `ConcurrentHashMap<Object, ConcurrentHashMap<Object, InFlightElement<OutputT>>>`) so you only look up and iterate over active elements for the current key.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4580018012
/gemini review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "AMOOOMA (via GitHub)" <gi...@apache.org>.
AMOOOMA commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4674299019
I also just realized you have files duplicated and the naming for the class should match AsyncWrapper instead of calling AsyncDoFn
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4675754732
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`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm merged PR #38609:
URL: https://github.com/apache/beam/pull/38609
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "AMOOOMA (via GitHub)" <gi...@apache.org>.
AMOOOMA commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3383460846
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
Review Comment:
Is the naming confirmed? If not prefer to stick to the AsyncWrapper which is the same as the python SDK. We can also try to send a rename PR later to do them both at the same time.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
Review 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.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (!finishedElementIds.contains(elementId)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ // Outputs use processing-time timestamps matching Python behavior
+ for (List<OutputT> outputs : toReturn) {
+ for (OutputT out : outputs) {
+ receiver.output(out);
+ }
+ }
+
+ LOG.info(
+ "Items finished: {}, not yet finished: {}, " + "rescheduled: {}, in processing state: {}",
+ itemsFinished,
+ itemsNotYetFinished,
+ itemsRescheduled,
+ itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ // Package-private helper methods for testing direct execution without Pipeline / ProcessContext
+ // boilerplate
+ void processDirect(
Review Comment:
if for testing only consider making the function name processDirect_test or similar that fits the java coding style.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
Review Comment:
Can we have a util function for this? it affects readability a bit.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
Review Comment:
same here.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
Review Comment:
If possible add some comments and make static const for these numbers likes maybe for miliseconds you can have MILISECONDS or something to make it readable.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (!finishedElementIds.contains(elementId)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ // Outputs use processing-time timestamps matching Python behavior
+ for (List<OutputT> outputs : toReturn) {
+ for (OutputT out : outputs) {
+ receiver.output(out);
+ }
+ }
+
+ LOG.info(
+ "Items finished: {}, not yet finished: {}, " + "rescheduled: {}, in processing state: {}",
+ itemsFinished,
+ itemsNotYetFinished,
+ itemsRescheduled,
+ itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ // Package-private helper methods for testing direct execution without Pipeline / ProcessContext
+ // boilerplate
+ void processDirect(
+ KV<K, InputT> element,
+ BoundedWindow window,
+ Instant timestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer) {
+ scheduleItem(element, window, timestamp);
+ toProcessState.add(element);
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ List<OutputT> commitFinishedItemsDirect(
+ Instant fireTimestamp, BagState<KV<K, InputT>> toProcessState, Timer timer) {
+ AccumulatingOutputReceiver<OutputT> receiver = new AccumulatingOutputReceiver<>();
+ commitFinishedItems(fireTimestamp, toProcessState, timer, receiver);
+ return receiver.getOutputs();
+ }
+
+ boolean isEmpty() {
+ return getItemsInBuffer().get() == 0;
Review Comment:
nit: we can use .isempty() is there's one, i am not too familiar with the packages so feel free to ignore if not.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #38609:
URL: https://github.com/apache/beam/pull/38609#issuecomment-4523058477
## Summary of Changes
Hello, I'm Gemini Code Assist[^1]! 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.
<details>
<summary><b>New Features</b></summary>
🧠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](https://codeassist.google/code-review/login) to enable Memory in your admin console.**
</details>
<details>
<summary><b>Using Gemini Code Assist</b></summary>
<br>
The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
<b>Invoking Gemini</b>
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.
<b>Customization</b>
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](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
<b>Limitations & Feedback</b>
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with :thumbsup: and :thumbsdown: 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](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
</details>
[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3327033953
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
+ new java.util.concurrent.ConcurrentLinkedQueue<>();
+
+ // Store original timestamp
+ private final Instant inputTimestamp;
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
Review Comment:
You are correct that I previously suggested `ConcurrentLinkedQueue`. I apologize for the inconsistency. In the context of the `AccumulatingOutputReceiver` being instantiated per element and accessed by a single background thread, `ArrayList` is indeed safer and more efficient as it avoids the overhead of thread-safe synchronization. Thank you for pointing out the contradiction.
```suggestion
private final List<TimestampedOutput<T>> outputs = new java.util.ArrayList<>();
// Store original timestamp
private final Instant inputTimestamp;
AccumulatingOutputReceiver(Instant inputTimestamp) {
this.inputTimestamp = inputTimestamp;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3326843882
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,740 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
I appreciate the clarification regarding the design alignment with the Python SDK. My concern is specifically about the scope of the static `lock`.
Because `lock` is a single static instance, it is shared by *all* `AsyncDoFn` instances within the same JVM. If a pipeline contains multiple `AsyncDoFn` transforms, every element processed by any of them must acquire this same lock to call `scheduleIfRoom` or `commitFinishedItems`.
Even if the heavy processing happens outside the lock, the serialization of the scheduling and reconciliation logic across all `AsyncDoFn` transforms in the pipeline creates a global bottleneck that scales poorly as the number of async transforms increases. Using a `ConcurrentHashMap<String, ReentrantLock>` keyed by the `uuid` would allow instances to lock independently, maintaining the desired thread safety while removing the global contention point.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291590212
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
Review Comment:
@gemini-code-assist
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291622499
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
@gemini-coding-assistant how would you best fix this issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321253876
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
Review Comment:
@gemini-code-assist Python SDK also concurrently shares its _sync_fn instance. Keeping the shared syncFn in Java maintains perfect behavioral parity with Python
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321305011
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
Review Comment:
While maintaining parity with the Python SDK is a valid design goal, hardcoding `System.currentTimeMillis()` makes the `AsyncDoFn` difficult to test deterministically. A common pattern in the Beam Java SDK to resolve this is to inject a `TimeProvider` interface. This allows you to use `System.currentTimeMillis()` in production while providing a mock implementation for tests, preserving the exact logic you need for parity while ensuring testability.
```suggestion
private Instant nextTimeToFire(@Nullable K key) {
long seed = (key == null) ? 0 : key.hashCode();
double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
double nowSec = System.currentTimeMillis() / 1000.0;
double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
double offset = fractionalOffset * timerFrequencySec;
return Instant.ofEpochMilli((long) ((base + offset) * 1000));
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3367030447
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
Review Comment:
In the current implementation of `teardown()`, `DoFnInvokers.invokerFor(syncFn).invokeTeardown()` is called before shutting down the thread pool. This introduces a critical race condition: active background threads in the pool might still be executing `invokeProcessElement` or `invokeFinishBundle` on the `syncFn` instance while its resources (e.g., database connections, clients) are being torn down.\n\nTo prevent this, the thread pool should be shut down and awaited first (when the reference count reaches 0), and only then should `invokeTeardown()` be called on the `syncFn`.
```java
public void teardown() {\n ExecutorService threadPoolToShutdown = null;\n lock.lock();\n try {\n AtomicInteger refCount = refCounts.get(uuid);\n if (refCount != null && refCount.decrementAndGet() == 0) {\n refCounts.remove(uuid);\n threadPoolToShutdown = pool.remove(uuid);\n processingElements.remove(uuid);\n itemsInBuffer.remove(uuid);\n }\n } finally {\n lock.unlock();\n }\n if (threadPoolToShutdown != null) {\n threadPoolToShutdown.shutdown();\n try {\n if (!threadPoolToShutdown.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {\n threadPoolToShutdown.shutdownNow();\n }\n } catch (InterruptedException e) {\n threadPoolToShutdown.shutdownNow();\n Thread.currentThread().interrupt();\n }\n }\n DoFnInvokers.invokerFor(syncFn).invokeTeardown();\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
The single static `lock` is shared across all instances of `AsyncDoFn` in the JVM. This causes unnecessary lock contention across different pipeline steps or different instances of `AsyncDoFn` running in the same JVM.\n\nTo eliminate this bottleneck, we should use a per-UUID lock (e.g., stored in a static `ConcurrentHashMap<String, ReentrantLock>`) so that different `AsyncDoFn` transforms do not block each other.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
Review Comment:
Currently, `processingElements` is a flat map of all active elements across all keys for a given `uuid`. In `commitFinishedItems`, the code iterates over the entire map to find elements matching the current key. This results in O(N) complexity per timer fire, where N is the total number of active elements across all keys in the JVM. Under heavy load with many keys, this creates a massive O(N * K) performance bottleneck.\n\nWe should restructure `processingElements` to group active elements by partition key `K`. This reduces the lookup complexity to O(1) and avoids scanning elements of other keys entirely.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
Review Comment:
The thread pool is created using `Executors.newFixedThreadPool(parallelism)` with the default thread factory, which produces generic thread names like `pool-X-thread-Y`. In production, this makes thread dumps extremely difficult to read and debug.\n\nUsing a custom thread factory to name the threads with the `uuid` and an incrementing index will greatly improve maintainability and operational visibility.
```java
if (useThreadPool) {\n pool.computeIfAbsent(\n uuid,\n k -> {\n final AtomicInteger threadIdx = new AtomicInteger(0);\n return Executors.newFixedThreadPool(\n parallelism,\n r -> {\n Thread t = new Thread(r);\n t.setName("async-dofn-" + uuid + "-worker-" + threadIdx.getAndIncrement());\n t.setDaemon(true);\n return t;\n });\n });\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
Review Comment:
During timer execution (`onTimer` / `commitFinishedItems`), calling the blocking `scheduleItem` method when the buffer is full will block the runner's main thread. If there are multiple rescheduled elements, this can block the main thread for several seconds, leading to severe pipeline lag, runner timeouts, or SDK harness crashes.\n\nSince elements that fail to schedule are already safely preserved in `toProcessState` (because they are not in `finishedElementIds`), we can attempt to schedule them non-blockingly using `scheduleIfRoom`. If the buffer is full, they will simply remain in the state and be retried when the next timer fires.
```suggestion
// Reschedule missing elements non-blockingly to avoid blocking the runner's thread\n for (KV<K, InputT> element : toReschedule) {\n scheduleIfRoom(element, GlobalWindow.INSTANCE, fireTimestamp, false);\n }
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
Review Comment:
When `useThreadPool` is false, the code falls back to `ForkJoinPool.commonPool()`. Since `AsyncDoFn` is specifically designed for high-latency, I/O-heavy operations, running these tasks on the common pool can easily saturate it and starve other parallel tasks in the JVM.\n\nWe should strongly recommend making `useThreadPool` default to `true` or warning against setting it to `false` in production environments.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3384303747
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
+ activeElements.remove(cancelId);
+ }
+ }
+
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (!finishedElementIds.contains(elementId)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
+
+ // Emit completed outputs
+ // (Emit completed tasks immediately; do not wait for all active tasks to finish).
+ // Outputs use processing-time timestamps matching Python behavior
+ for (List<OutputT> outputs : toReturn) {
+ for (OutputT out : outputs) {
+ receiver.output(out);
+ }
+ }
+
+ LOG.info(
+ "Items finished: {}, not yet finished: {}, " + "rescheduled: {}, in processing state: {}",
+ itemsFinished,
+ itemsNotYetFinished,
+ itemsRescheduled,
+ itemsInProcessingState);
+
+ if (itemsInProcessingState > 0) {
+ Instant timeToFire = nextTimeToFire(key);
+ timer.set(timeToFire);
+ }
+ }
+
+ // Package-private helper methods for testing direct execution without Pipeline / ProcessContext
+ // boilerplate
+ void processDirect(
Review Comment:
Kept it camelCase and added @VisibleForTesting to match Java style.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3384308182
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
Review Comment:
Yup. Added readable static constants for all these hard coded numbers.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291461494
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
The use of a single static `ReentrantLock` across all instances of `AsyncDoFn` can lead to significant thread contention if multiple asynchronous transforms are used in the same pipeline. Since the registries are already partitioned by `uuid`, consider using a per-transform lock or leveraging `ConcurrentHashMap` atomic operations to reduce global contention.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
Review Comment:
This block iterates over the entire `activeElements` map for the transform, which contains elements for ALL keys. Since this runs every time a timer fires for ANY key, the performance will degrade quadratically as the number of keys and in-flight elements increases. This 'orphan cleanup' is likely unnecessary as tasks will naturally finish or be rediscovered by retried bundles via their `elementId`. Consider removing this global scan.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
The `AccumulatingOutputReceiver` ignores timestamps, which can lead to data loss or incorrect windowing if the runner's default timestamp (e.g., the timer's firing time) falls outside the element's window. Both `output` and `outputWithTimestamp` should preserve the relevant timestamp of the element being processed or the explicitly provided timestamp.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ List<List<OutputT>> toReturn = new ArrayList<>();
+ List<KV<K, InputT>> finishedItems = new ArrayList<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+ int itemsCancelled = toCancel.size();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<K, InputT, OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ finishedItems.add(element);
+ activeElements.remove(elementId);
+ }
+ } else {
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
+
+ // Update State: keep only unfinished items
+ toProcessState.clear();
+ int itemsInProcessingState = 0;
+ for (KV<K, InputT> element : stateList) {
+ if (!finishedItems.contains(element)) {
+ toProcessState.add(element);
+ itemsInProcessingState++;
+ }
+ }
Review Comment:
The reconciliation loop has a logic bug when `stateList` contains duplicate elements (e.g., the same ID processed multiple times before a timer fires). The first instance of an element will remove it from `activeElements` (line 589), causing subsequent instances of the same element in the same loop to be incorrectly added to `toReschedule` (line 602). Additionally, the state update at line 619 uses `List.contains` in a loop, resulting in O(N^2) complexity. Consider using a `Set` of `elementId`s to track finished tasks and ensure each ID is only processed once per reconciliation cycle.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291590818
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
Review Comment:
The performance issue stems from the global scan of `activeElements` inside `commitFinishedItems`. Since this method is triggered by a timer scoped to a specific key, you should only inspect elements associated with that key.
To optimize this, I recommend maintaining a secondary index: `ConcurrentHashMap<K, Set<Object>> activeElementIdsByKey`.
1. **Update `scheduleIfRoom`**: When adding an element to `activeElements`, also add its `elementId` to the set for that key in `activeElementIdsByKey`.
2. **Update `commitFinishedItems`**: Instead of iterating over the entire `activeElements` map, iterate only over the `elementId`s in `activeElementIdsByKey.get(key)`.
3. **Cleanup**: Ensure you remove the `elementId` from the set when the task completes or is cancelled.
This changes the complexity from O(N) (where N is the total number of elements across all keys) to O(M) (where M is the number of elements for the specific key being processed).
```suggestion
// Add this to your class fields
private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, Set<Object>>> activeElementIdsByKey = new ConcurrentHashMap<>();
// In scheduleIfRoom, after activeElements.put(elementId, ...):
activeElementIdsByKey.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>())
.computeIfAbsent(element.getKey(), k -> Collections.synchronizedSet(new HashSet<>()))
.add(elementId);
// In commitFinishedItems, replace the global scan with:
Set<Object> activeIds = activeElementIdsByKey.get(uuid).get(key);
if (activeIds != null) {
for (Object elementId : activeIds) {
// Process only elements for this key
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291592056
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?, ?, ?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>>)
+ (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+
+ ExecutorService threadPool;
+ lock.lock();
+ try {
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ } finally {
+ lock.unlock();
+ }
+
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.output(output);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element, future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+ sleepTime *= 2;
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ Random random = new Random(seed);
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = random.nextDouble() * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ private void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<K, InputT, OutputT>> activeElements =
+ getProcessingElements();
+ Set<Object> stateIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancel = new ArrayList<>();
+ lock.lock();
+ try {
+ // Cancel any active elements for this key that are no longer in runner's state
+ for (Map.Entry<Object, InFlightElement<K, InputT, OutputT>> entry :
+ activeElements.entrySet()) {
+ Object elementId = entry.getKey();
+ InFlightElement<K, InputT, OutputT> inFlight = entry.getValue();
+
+ if (Objects.equals(inFlight.element.getKey(), key) && !stateIds.contains(elementId)) {
+ inFlight.future.cancel(true);
+ toCancel.add(elementId);
+ LOG.info("Cancelling item {} which is no longer in state", inFlight.element);
+ }
+ }
+ for (Object elementId : toCancel) {
+ activeElements.remove(elementId);
+ }
+ } finally {
+ lock.unlock();
+ }
Review Comment:
Yes, the change to iterate over `stateList` (the elements for the current key) instead of the entire `activeElements` map resolves the quadratic performance issue. By scoping the lookup to the elements associated with the current key, the complexity is now proportional to the number of elements for that specific key, which is the correct approach for this timer-based reconciliation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291641572
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
@gemini-code-assist how would you best fix this issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3291639713
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
@AMOOOMA
I kept a single static class lock because it matches the exact synchronization model used by Python's AsyncWrapper (_lock = RLock()).
If you prefer, we can change this to a per-transform lock or leverage ConcurrentHashMap to reduce global thread contention in Java, but I wanted to keep it aligned with the Python SDK first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3319379439
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,689 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<
+ String, ConcurrentHashMap<Object, InFlightElement<?, ?, ?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<K, InputT, OutputT> {
+ final KV<K, InputT> element;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(KV<K, InputT> element, CompletableFuture<List<OutputT>> future) {
+ this.element = element;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
Review Comment:
The issue of timestamp loss in `AccumulatingOutputReceiver` remains in the current implementation. The `output` and `outputWithTimestamp` methods still do not preserve the original element's timestamp, which is critical for correct windowing behavior in Beam. To resolve this, you should store the output along with its timestamp (e.g., using a `TimestampedValue` container) and emit it with the correct timestamp when the results are reconciled. The previous suggestion provides a robust way to handle this by wrapping the output values.
```suggestion
private static class TimestampedValue<T> {
final T value;
final Instant timestamp;
TimestampedValue(T value, Instant timestamp) {
this.value = value;
this.timestamp = timestamp;
}
}
private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
private final List<TimestampedValue<T>> outputs = Collections.synchronizedList(new ArrayList<>());
@Override
public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
return org.apache.beam.sdk.values.WindowedValues.<T>builder()
.setValue(value)
.setTimestamp(Instant.now())
.setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
.setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
.setReceiver(windowedValue -> outputs.add(new TimestampedValue<>(windowedValue.getValue(), windowedValue.getTimestamp())));
}
@Override
public void output(T output) {
outputs.add(new TimestampedValue<>(output, Instant.now()));
}
@Override
public void outputWithTimestamp(T output, Instant timestamp) {
outputs.add(new TimestampedValue<>(output, timestamp));
}
public List<TimestampedValue<T>> getOutputs() {
return outputs;
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3321169672
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,715 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn requires thread-safety ONLY if BOTH parallelism > 1 AND the
+ * DoFn is stateful (keeps instance state). 2) Tagged output multi-outputs are unsupported. 3)
+ * StartBundle/finishBundle are invoked per element so any batching or aggregation logic will not
+ * behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(Instant.now())
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, null));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ lock.lock();
+ try {
+ getItemsInBuffer().decrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+ Set<KV<K, InputT>> finishedItems = new HashSet<>();
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ for (KV<K, InputT> element : stateList) {
+ Object elementId = idFn.apply(element.getValue());
+
+ // Skip processing if we already completed, rescheduled, or found this elementId active in
+ // this cycle
+ if (finishedElementIds.contains(elementId)
+ || rescheduledElementIds.contains(elementId)
+ || inFlightElementIds.contains(elementId)) {
+ continue;
+ }
+
+ if (activeElements.containsKey(elementId)) {
+ InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ if (inFlight.future.isDone()) {
+ try {
+ if (!inFlight.future.isCancelled()) {
+ toReturn.add(inFlight.future.get());
+ }
+ finishedItems.add(element);
+ finishedElementIds.add(elementId);
+ activeElements.remove(elementId);
+ itemsFinished++;
+ } catch (Exception e) {
+ LOG.error("Error executing async task for element {}", element, e);
+ throw new RuntimeException("Error executing async task for element " + element, e);
+ }
+ } else {
+ inFlightElementIds.add(elementId);
+ itemsNotYetFinished++;
+ }
+ } else {
+ LOG.info(
+ "Item {} found in state but not in local active elements, scheduling now", element);
+ toReschedule.add(element);
+ rescheduledElementIds.add(elementId);
+ itemsRescheduled++;
+ }
+ }
+ } finally {
+ lock.unlock();
+ }
+
+ // Reschedule missing elements
+ for (KV<K, InputT> element : toReschedule) {
+ scheduleItem(element, GlobalWindow.INSTANCE, fireTimestamp);
+ }
Review Comment:
@gemini-code-assist The blocking loop with Thread.sleep is intentional to maintain 1-1 logic parity with the Python SDK implementation in async_dofn.py, which also blocks until room opens or a timeout is reached.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3326909234
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
Review Comment:
Currently, `processingElements` is a flat map of all active elements across all keys in the JVM. In `commitFinishedItems`, the code iterates over the entire map (`activeElements.entrySet()`) to find and cancel elements belonging to the current key. This is an $O(N)$ operation where $N$ is the total number of active elements across all keys in the JVM. Since `commitFinishedItems` is called per key when its timer fires, this results in $O(K \times N)$ complexity, which will severely degrade performance under heavy load with many keys.\n\nWe can optimize this to $O(1)$ lookup and $O(M)$ iteration (where $M$ is the number of active elements for the specific key) by grouping `processingElements` by the partition key `key`. Since `ConcurrentHashMap` does not support `null` keys, we can use a sentinel object to represent `null` keys.
```suggestion
private static final Object NULL_KEY = new Object();
private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, ConcurrentHashMap<Object, InFlightElement<?>>>>
processingElements = new ConcurrentHashMap<>();
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
+
+ private final Instant inputTimestamp; // <-- Store original timestamp
+
+ AccumulatingOutputReceiver(Instant inputTimestamp) {
+ this.inputTimestamp = inputTimestamp;
+ }
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setTimestamp(inputTimestamp)
+ .setWindows(java.util.Collections.singletonList(GlobalWindow.INSTANCE))
+ .setPaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING)
+ .setReceiver(
+ windowedValue ->
+ outputs.add(
+ new TimestampedOutput<>(
+ windowedValue.getValue(), windowedValue.getTimestamp())));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(new TimestampedOutput<>(output, inputTimestamp));
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(new TimestampedOutput<>(output, timestamp));
+ }
+
+ public List<T> getOutputs() {
+ List<T> rawOutputs = new ArrayList<>();
+ for (TimestampedOutput<T> out : outputs) {
+ rawOutputs.add(out.value);
+ }
+ return rawOutputs;
+ }
+
+ public List<TimestampedOutput<T>> getTimestampedOutputs() {
+ return outputs;
+ }
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public AsyncDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<TimestampedOutput<OutputT>>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>(timestamp);
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for AsyncDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "AsyncDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getTimestampedOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<TimestampedOutput<OutputT>>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
+ LOG.info("processing timer for key: {}", key);
+ }
+
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+
+ List<List<TimestampedOutput<OutputT>>> toReturn = new ArrayList<>();
+
+ List<KV<K, InputT>> toReschedule = new ArrayList<>();
+
+ int itemsFinished = 0;
+ int itemsNotYetFinished = 0;
+ int itemsRescheduled = 0;
+
+ Set<Object> finishedElementIds = new HashSet<>();
+ Set<Object> inFlightElementIds = new HashSet<>();
+ Set<Object> rescheduledElementIds = new HashSet<>();
+
+ lock.lock();
+ try {
+ Set<Object> stateElementIds = new HashSet<>();
+ for (KV<K, InputT> element : stateList) {
+ stateElementIds.add(idFn.apply(element.getValue()));
+ }
+
+ List<Object> toCancelIds = new ArrayList<>();
+ for (Map.Entry<Object, InFlightElement<OutputT>> entry : activeElements.entrySet()) {
+ InFlightElement<OutputT> inFlight = entry.getValue();
+ if (java.util.Objects.equals(inFlight.key, key)
+ && !stateElementIds.contains(entry.getKey())) {
+ toCancelIds.add(entry.getKey());
+ }
+ }
+
+ for (Object cancelId : toCancelIds) {
+ InFlightElement<OutputT> inFlight = activeElements.get(cancelId);
+ if (inFlight != null) {
+ inFlight.future.cancel(true);
Review Comment:
The cancellation logic calls `inFlight.future.cancel(true)`. In Java, `CompletableFuture.cancel(boolean mayInterruptIfRunning)` does **not** interrupt the thread executing the task, even if `mayInterruptIfRunning` is set to `true`. The background thread executing the wrapped `DoFn` will continue to run to completion, wasting CPU resources and potentially causing side effects, although its outputs will be ignored.\n\nIf true thread interruption is desired to stop the background task immediately, you should submit the task to the `ExecutorService` directly using `submit()`, which returns a standard `Future` whose `cancel(true)` method does support thread interruption, and then wrap or coordinate that with the `CompletableFuture`.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class TimestampedOutput<T> {
+ final T value;
+ final @Nullable Instant timestamp;
+
+ TimestampedOutput(T value, @Nullable Instant timestamp) {
+ this.value = value;
+ this.timestamp = timestamp;
+ }
+ }
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<TimestampedOutput<OutputT>>> future;
+
+ InFlightElement(
+ @Nullable Object key, CompletableFuture<List<TimestampedOutput<OutputT>>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<TimestampedOutput<T>> outputs =
+ Collections.synchronizedList(new ArrayList<>());
Review Comment:
`AccumulatingOutputReceiver` uses `Collections.synchronizedList(new ArrayList<>())` to store outputs. However, in `getOutputs()`, it iterates over this list using an enhanced for-loop without synchronizing on the list instance. According to the Java documentation for `Collections.synchronizedList`, iterating over a synchronized list without manual synchronization on the list object is not thread-safe and can throw `ConcurrentModificationException` or cause race conditions if another thread is concurrently adding to the list.\n\nTo make this fully thread-safe and lock-free, we should use `ConcurrentLinkedQueue` instead of a synchronized list.
```suggestion
private final java.util.concurrent.ConcurrentLinkedQueue<TimestampedOutput<T>> outputs =
new java.util.concurrent.ConcurrentLinkedQueue<>();
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncDoFn.java:
##########
@@ -0,0 +1,749 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class AsyncDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
Review Comment:
The lock used for synchronization is a static `ReentrantLock`. Because this lock is static, all instances of `AsyncDoFn` across the entire JVM (including completely unrelated transforms in the same pipeline) will contend for this single lock. This can severely limit throughput and cause thread starvation under high load.\n\nSince all shared states are already keyed by `uuid`, we only need to synchronize accesses for the same `uuid`. We can resolve this by maintaining a map of locks per `uuid` and retrieving the lock dynamically.
```java
private static final ConcurrentHashMap<String, ReentrantLock> locks = new ConcurrentHashMap<>();
private ReentrantLock getLock() {
return locks.computeIfAbsent(uuid, k -> new ReentrantLock());
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement Asynchronous wrapper for DoFn in Java SDK [beam]
Posted by "tejasiyer-dev (via GitHub)" <gi...@apache.org>.
tejasiyer-dev commented on code in PR #38609:
URL: https://github.com/apache/beam/pull/38609#discussion_r3384298810
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PerKeyConcurrentDoFn.java:
##########
@@ -0,0 +1,714 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.transforms;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class that wraps a dofn and converts it from one which process elements synchronously to one
+ * which processes them asynchronously.
+ *
+ * <p>For synchronous dofns the default settings mean that many (100s) of elements will be processed
+ * in parallel and that processing an element will block all other work on that key. In addition
+ * runners are optimized for latencies less than a few seconds and longer operations can result in
+ * high retry rates. Async should be considered when the default parallelism is not correct and/or
+ * items are expected to take longer than a few seconds to process.
+ *
+ * <p>/* NOTE: 1) The wrapped syncFn REQUIRES thread-safety if BOTH parallelism > 1 and the DoFn is
+ * stateful. 2) Tagged output multi-outputs are unsupported. 3) StartBundle/finishBundle are invoked
+ * per element so any batching or aggregation logic will not behave as expected.
+ */
+public class PerKeyConcurrentDoFn<K, InputT, OutputT> extends DoFn<KV<K, InputT>, OutputT> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PerKeyConcurrentDoFn.class);
+
+ private static final int DEFAULT_MIN_BUFFER_CAPACITY = 10;
+ private static final int DEFAULT_TIMEOUT_SEC = 1;
+ private static final int DEFAULT_MAX_WAIT_TIME_MS = 500;
+ private static final int TEARDOWN_AWAIT_SEC = 5;
+ private static final int INITIAL_BACKOFF_SLEEP_MS = 10;
+ private static final int BACKPRESSURE_LOG_THRESHOLD_MS = 10000;
+
+ @StateId("to_process")
+ private final StateSpec<BagState<KV<K, InputT>>> toProcessSpec;
+
+ @TimerId("timer")
+ private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final DoFn<InputT, OutputT> syncFn;
+ private final int parallelism;
+ private final Duration timerFrequency;
+ private final int maxItemsToBuffer;
+ private final Duration timeout;
+ private final Duration maxWaitTime;
+ private final SerializableFunction<InputT, Object> idFn;
+ private final boolean useThreadPool;
+ private final String uuid;
+
+ private transient volatile @Nullable PipelineOptions pipelineOptions;
+
+ // Shared JVM-Wide States (Static Registries)
+ // Map-backed registry holding shared resources across serialized worker instances. Since runners
+ // clone DoFn instances on the same worker node, static maps ensure safe JVM-wide resource reuse.
+ private static final ConcurrentHashMap<String, ExecutorService> pool = new ConcurrentHashMap<>();
+ // activeElements (processingElements) is global JVM memory (all keys)
+ private static final ConcurrentHashMap<String, ConcurrentHashMap<Object, InFlightElement<?>>>
+ processingElements = new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<String, AtomicInteger> itemsInBuffer =
+ new ConcurrentHashMap<>();
+ // Reference counts for cloned instances sharing the same UUID. Coordinates safe,
+ // leak-free thread pool shutdown during teardown without crashing active sibling clones.
+ private static final ConcurrentHashMap<String, AtomicInteger> refCounts =
+ new ConcurrentHashMap<>();
+
+ // If contention becomes a bottleneck, this can be replaced with per-uuid locks
+ // in a ConcurrentHashMap
+ private static final ReentrantLock lock = new ReentrantLock();
+ private static final boolean verboseLogging = false;
+
+ private static class InFlightElement<OutputT> {
+ final @Nullable Object key;
+ final CompletableFuture<List<OutputT>> future;
+
+ InFlightElement(@Nullable Object key, CompletableFuture<List<OutputT>> future) {
+ this.key = key;
+ this.future = future;
+ }
+ }
+
+ // The In-Memory Accumulating Receiver
+ // Accumulates elements in-memory during asynchronous background worker execution.
+ // Buffered elements are only committed downstream once the parent task completes successfully
+ // and the timer fires.
+ private static class AccumulatingOutputReceiver<T> implements OutputReceiver<T> {
+ private final List<T> outputs = new ArrayList<>();
+
+ AccumulatingOutputReceiver() {}
+
+ @Override
+ public org.apache.beam.sdk.values.OutputBuilder<T> builder(T value) {
+ return org.apache.beam.sdk.values.WindowedValues.<T>builder()
+ .setValue(value)
+ .setReceiver(windowedValue -> outputs.add(windowedValue.getValue()));
+ }
+
+ // Bypasses the nested anonymous OutputBuilder instantiation for standard outputs.
+ // JVM optimization to prevent garbage collection pressure under high pipeline throughput.
+ @Override
+ public void output(T output) {
+ outputs.add(output);
+ }
+
+ @Override
+ public void outputWithTimestamp(T output, Instant timestamp) {
+ outputs.add(output);
+ }
+
+ public List<T> getOutputs() {
+ return outputs;
+ }
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool) {
+ this(
+ syncFn,
+ parallelism,
+ timerFrequency,
+ maxItemsToBuffer,
+ timeout,
+ maxWaitTime,
+ idFn,
+ useThreadPool,
+ null);
+ }
+
+ public PerKeyConcurrentDoFn(
+ DoFn<InputT, OutputT> syncFn,
+ int parallelism,
+ Duration timerFrequency,
+ @Nullable Integer maxItemsToBuffer,
+ @Nullable Duration timeout,
+ @Nullable Duration maxWaitTime,
+ @Nullable SerializableFunction<InputT, Object> idFn,
+ boolean useThreadPool,
+ @Nullable Coder<KV<K, InputT>> coder) {
+ this.syncFn = syncFn;
+ this.parallelism = parallelism;
+ if (timerFrequency.getMillis() <= 0) {
+ throw new IllegalArgumentException("timerFrequency must be greater than zero");
+ }
+ this.timerFrequency = timerFrequency;
+ this.maxItemsToBuffer =
+ (maxItemsToBuffer != null)
+ ? maxItemsToBuffer
+ : Math.max(parallelism * 2, DEFAULT_MIN_BUFFER_CAPACITY);
+ this.timeout = (timeout != null) ? timeout : Duration.standardSeconds(DEFAULT_TIMEOUT_SEC);
+ this.maxWaitTime =
+ (maxWaitTime != null) ? maxWaitTime : Duration.millis(DEFAULT_MAX_WAIT_TIME_MS);
+ this.idFn =
+ (idFn != null)
+ ? idFn
+ : (SerializableFunction<InputT, Object>)
+ input -> java.util.Objects.requireNonNull(input);
+ this.useThreadPool = useThreadPool;
+ this.uuid = UUID.randomUUID().toString();
+ this.toProcessSpec = (coder != null) ? StateSpecs.bag(coder) : StateSpecs.bag();
+ }
+
+ private ExecutorService getThreadPool() {
+ ExecutorService threadPool = pool.get(uuid);
+ if (threadPool == null) {
+ throw new IllegalStateException("Thread pool not initialized for UUID: " + uuid);
+ }
+ return threadPool;
+ }
+
+ @SuppressWarnings("unchecked")
+ private ConcurrentHashMap<Object, InFlightElement<OutputT>> getProcessingElements() {
+ ConcurrentHashMap<Object, InFlightElement<?>> elements = processingElements.get(uuid);
+ if (elements == null) {
+ throw new IllegalStateException("Processing elements map not initialized for UUID: " + uuid);
+ }
+ return (ConcurrentHashMap<Object, InFlightElement<OutputT>>) (ConcurrentHashMap<?, ?>) elements;
+ }
+
+ private AtomicInteger getItemsInBuffer() {
+ AtomicInteger buffer = itemsInBuffer.get(uuid);
+ if (buffer == null) {
+ throw new IllegalStateException("Buffer counter not initialized for UUID: " + uuid);
+ }
+ return buffer;
+ }
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ this.pipelineOptions = options;
+
+ // Setup the wrapped DoFn
+ DoFnInvokers.invokerFor(syncFn)
+ .invokeSetup(
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Setup";
+ }
+ });
+
+ if (useThreadPool) {
+ LOG.info("Using thread pool for asynchronous execution with parallelism {}", parallelism);
+ }
+
+ lock.lock();
+ try {
+ if (useThreadPool) {
+ pool.computeIfAbsent(uuid, k -> Executors.newFixedThreadPool(parallelism));
+ }
+ processingElements.computeIfAbsent(uuid, k -> new ConcurrentHashMap<>());
+ itemsInBuffer.computeIfAbsent(uuid, k -> new AtomicInteger(0));
+ refCounts.computeIfAbsent(uuid, k -> new AtomicInteger(0)).incrementAndGet();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ // Clean up JVM-wide shared resources to prevent thread leaks on the worker
+ @Teardown
+ public void teardown() {
+ DoFnInvokers.invokerFor(syncFn).invokeTeardown();
+ ExecutorService threadPool = null;
+ lock.lock();
+ try {
+ AtomicInteger refCount = refCounts.get(uuid);
+ if (refCount != null && refCount.decrementAndGet() == 0) {
+ refCounts.remove(uuid);
+ threadPool = pool.remove(uuid);
+ processingElements.remove(uuid);
+ itemsInBuffer.remove(uuid);
+ }
+ } finally {
+ lock.unlock();
+ }
+ if (threadPool != null) {
+ threadPool.shutdown();
+ try {
+ if (!threadPool.awaitTermination(TEARDOWN_AWAIT_SEC, TimeUnit.SECONDS)) {
+ threadPool.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ threadPool.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ // Asynchronous Scheduling & Deduplication
+ // Submits tasks to the background thread pool. If an element with the same ID is already
+ // in-flight,
+ // the submission is silently ignored to enforce exactly-once semantics.
+ private boolean scheduleIfRoom(
+ KV<K, InputT> element, BoundedWindow window, Instant timestamp, boolean ignoreBuffer) {
+ lock.lock();
+ try {
+ ConcurrentHashMap<Object, InFlightElement<OutputT>> activeElements = getProcessingElements();
+ Object elementId = idFn.apply(element.getValue());
+
+ if (activeElements.containsKey(elementId)) {
+ LOG.info("Item {} already in processing elements", element);
+ return true;
+ }
+
+ int currentBuffer = getItemsInBuffer().get();
+ if (currentBuffer < maxItemsToBuffer || ignoreBuffer) {
+ java.util.concurrent.Executor executor =
+ useThreadPool ? getThreadPool() : java.util.concurrent.ForkJoinPool.commonPool();
+
+ // Pending asynchronous task that will produce a list of outputs
+ CompletableFuture<List<OutputT>> future =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ AccumulatingOutputReceiver<OutputT> receiver =
+ new AccumulatingOutputReceiver<>();
+ DoFnInvoker<InputT, OutputT> invoker = DoFnInvokers.invokerFor(syncFn);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> bundleArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
+ DoFn<InputT, OutputT> doFn) {
+ return doFn.new FinishBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions();
+ }
+
+ @Override
+ public void output(
+ OutputT output, Instant timestamp, BoundedWindow window) {
+ receiver.outputWithTimestamp(output, timestamp);
+ }
+
+ @Override
+ public <T> void output(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ BoundedWindow window) {
+ throw new UnsupportedOperationException(
+ "Tagged output not supported in "
+ + "FinishBundleContext for PerKeyConcurrentDoFn");
+ }
+ };
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Bundle";
+ }
+ };
+
+ invoker.invokeStartBundle(bundleArgProvider);
+
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> processArgProvider =
+ new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+ @Override
+ public InputT element(DoFn<InputT, OutputT> doFn) {
+ return element.getValue();
+ }
+
+ @Override
+ public OutputReceiver<OutputT> outputReceiver(
+ DoFn<InputT, OutputT> doFn) {
+ return receiver;
+ }
+
+ @Override
+ public BoundedWindow window() {
+ return window;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public PipelineOptions pipelineOptions() {
+ PipelineOptions options = pipelineOptions;
+ if (options == null) {
+ throw new IllegalStateException("PipelineOptions not set");
+ }
+ return options;
+ }
+
+ @Override
+ public String getErrorContext() {
+ return "PerKeyConcurrentDoFn/Process";
+ }
+ };
+
+ invoker.invokeProcessElement(processArgProvider);
+ invoker.invokeFinishBundle(bundleArgProvider);
+
+ return receiver.getOutputs();
+ } catch (Exception e) {
+ throw new CompletionException(e);
+ }
+ },
+ executor);
+
+ // Assigned to 'unused' to satisfy ErrorProne while preserving parent future for
+ // cancellation
+ CompletableFuture<List<OutputT>> unused =
+ future.whenComplete(
+ (res, ex) -> {
+ getItemsInBuffer().decrementAndGet();
+ });
+
+ activeElements.put(elementId, new InFlightElement<>(element.getKey(), future));
+ getItemsInBuffer().incrementAndGet();
+ return true;
+ }
+
+ return false;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void scheduleItem(KV<K, InputT> element, BoundedWindow window, Instant timestamp) {
+ boolean done = false;
+ long sleepTime = INITIAL_BACKOFF_SLEEP_MS;
+ long totalSleep = 0;
+ long timeoutMs = timeout.getMillis();
+
+ while (!done && totalSleep < timeoutMs) {
+ done = scheduleIfRoom(element, window, timestamp, false);
+ if (!done) {
+ long sleep = Math.min(maxWaitTime.getMillis(), sleepTime);
+ if (verboseLogging || totalSleep > BACKPRESSURE_LOG_THRESHOLD_MS) {
+ LOG.info(
+ "buffer is full for item {}, {} waiting {} ms. Have waited for {} ms.",
+ element,
+ getItemsInBuffer().get(),
+ sleep,
+ totalSleep);
+ }
+ try {
+ Thread.sleep(sleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while waiting for space in buffer", e);
+ }
+
+ // Prevents long overflow possibility
+ if (sleepTime < maxWaitTime.getMillis()) {
+ sleepTime *= 2;
+ }
+
+ totalSleep += sleep;
+ }
+ }
+ // Timeout: element skips JVM pool but stays in BagState for timer to reschedule later.
+ }
+
+ // Uses hashcode based jitter instead of random for deterministic rescheduling
+ // Satisfies lint check
+ private Instant nextTimeToFire(@Nullable K key) {
+ long seed = (key == null) ? 0 : key.hashCode();
+ double fractionalOffset = Math.abs(seed % 1000000) / 1000000.0;
+ double timerFrequencySec = timerFrequency.getMillis() / 1000.0;
+ double nowSec = System.currentTimeMillis() / 1000.0;
+
+ double base = Math.floor((nowSec + timerFrequencySec) / timerFrequencySec) * timerFrequencySec;
+ double offset = fractionalOffset * timerFrequencySec;
+
+ return Instant.ofEpochMilli((long) ((base + offset) * 1000));
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ BoundedWindow window,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer) {
+
+ KV<K, InputT> element = c.element();
+ scheduleItem(element, window, c.timestamp());
+ toProcessState.add(element);
+
+ Instant timeToFire = nextTimeToFire(element.getKey());
+ timer.set(timeToFire);
+ }
+
+ @OnTimer("timer")
+ public void onTimer(
+ OnTimerContext c,
+ @StateId("to_process") BagState<KV<K, InputT>> toProcessState,
+ @TimerId("timer") Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ commitFinishedItems(c.fireTimestamp(), toProcessState, timer, receiver);
+ }
+
+ // Synchronizes local task results with the runner's persistent state container.
+ // Emits successfully completed elements, cancels rolled-back tasks, and reschedules lost work.
+ void commitFinishedItems(
+ Instant fireTimestamp,
+ BagState<KV<K, InputT>> toProcessState,
+ Timer timer,
+ OutputReceiver<OutputT> receiver) {
+
+ Iterable<KV<K, InputT>> toProcessLocal = toProcessState.read();
+ if (toProcessLocal == null || !toProcessLocal.iterator().hasNext()) {
+ // Early Exit: if BagState is empty, we skip checking activeElements for this key.
+ return;
+ }
+
+ // Since fireTimestamp is key-scoped, we determine the current key from the first element in
+ // state
+ List<KV<K, InputT>> stateList = new ArrayList<>();
+ K key = null;
+ for (KV<K, InputT> element : toProcessLocal) {
+ stateList.add(element);
+ if (key == null) {
+ key = element.getKey();
+ }
+ }
+
+ if (verboseLogging) {
Review Comment:
Added verbose logging checks everywhere and made private util functions for logging in latest commit. Let me know if I should change it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org