Add drain support for Dataflow and Flink#38786
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request adds support for draining pipelines in the Apache Beam Java SDK, enabling runners to stop accepting new input while finishing the processing of existing data. It updates the core SDK to include new pipeline states and provides implementations for both the Dataflow and Flink runners to handle these drain requests effectively. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a first-class “drain” lifecycle operation to PipelineResult, introducing explicit DRAINING / DRAINED states and updating Dataflow + Flink runners to support and test these semantics.
Changes:
- Added
PipelineResult#drain()default API and newPipelineResult.StatevaluesDRAININGandDRAINED. - Updated Dataflow state mapping and
DataflowPipelineJobto issue drain requests, plus added/updated tests. - Implemented drain support for Flink runner results (done + detached cases) with tests.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java | Introduces drain() API + DRAINING/DRAINED states in the core contract. |
| runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java | Maps Dataflow service drain states to the new Beam states. |
| runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java | Updates assertions to validate the new drain state mapping. |
| runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java | Implements drain via a generalized lifecycle request helper and logs terminal DRAINED. |
| runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java | Adds a unit test validating drain request behavior. |
| runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java | Treats DRAINED as a successful terminal condition for termination waits. |
| runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkRunnerResult.java | Implements a no-op drain() returning DONE for already-finished results. |
| runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java | Implements detached drain using stopWithSavepoint, exposing DRAINING/DRAINED. |
| runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java | Adds tests validating drain semantics for Flink (done + detached). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| private JobClient jobClient; | ||
| private int jobCheckIntervalInSecs; | ||
| private @Nullable CompletableFuture<String> drainSavepointFuture; |
| CompletableFuture<String> drainFuture = drainSavepointFuture; | ||
| if (drainFuture != null) { | ||
| return getDrainState(drainFuture); | ||
| } |
| public synchronized State drain() throws IOException { | ||
| CompletableFuture<String> drainFuture = drainSavepointFuture; | ||
| if (drainFuture == null) { | ||
| drainFuture = this.jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT); | ||
| drainSavepointFuture = drainFuture; | ||
| } | ||
| return getDrainState(drainFuture); | ||
| } |
| throw new RuntimeException("Fail to drain flink job", e); | ||
| } catch (ExecutionException e) { | ||
| throw new RuntimeException("Fail to drain flink job", e); | ||
| } |
| capitalizedAction, | ||
| state); | ||
| return state; | ||
| } else if (e.getMessage().contains("has terminated")) { |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for draining pipelines across Flink and Google Cloud Dataflow runners by implementing the drain() method and adding DRAINING and DRAINED states to the Beam SDK. Key feedback focuses on improving thread safety and robustness, specifically by declaring the drainSavepointFuture field as volatile, handling exceptional completion in getState(), allowing retries and throwing checked IOExceptions in the Flink runner's drain() implementation, and adding a null check to prevent a potential NullPointerException when reading exception messages in the Dataflow runner.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
|
||
| private JobClient jobClient; | ||
| private int jobCheckIntervalInSecs; | ||
| private @Nullable CompletableFuture<String> drainSavepointFuture; |
There was a problem hiding this comment.
The drainSavepointFuture field is read in getState() without synchronization, but written in the synchronized drain() method. To ensure proper visibility and thread safety across threads, this field should be declared as volatile.
| private @Nullable CompletableFuture<String> drainSavepointFuture; | |
| private volatile @Nullable CompletableFuture<String> drainSavepointFuture; |
| CompletableFuture<String> drainFuture = drainSavepointFuture; | ||
| if (drainFuture != null) { | ||
| return getDrainState(drainFuture); | ||
| } |
There was a problem hiding this comment.
If the drainSavepointFuture completes exceptionally (e.g., if the stop-with-savepoint operation fails), calling getState() will currently throw a RuntimeException and crash the caller. Instead, we should be more robust: if the drain operation failed, the job is likely still running or in its previous state. We can check isCompletedExceptionally() and, if true, fall back to querying the actual Flink job status via jobClient.getJobStatus().
| CompletableFuture<String> drainFuture = drainSavepointFuture; | |
| if (drainFuture != null) { | |
| return getDrainState(drainFuture); | |
| } | |
| CompletableFuture<String> drainFuture = drainSavepointFuture; | |
| if (drainFuture != null && !drainFuture.isCompletedExceptionally()) { | |
| return drainFuture.isDone() ? State.DRAINED : State.DRAINING; | |
| } |
| public synchronized State drain() throws IOException { | ||
| CompletableFuture<String> drainFuture = drainSavepointFuture; | ||
| if (drainFuture == null) { | ||
| drainFuture = this.jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT); | ||
| drainSavepointFuture = drainFuture; | ||
| } | ||
| return getDrainState(drainFuture); | ||
| } | ||
|
|
||
| private State getDrainState(CompletableFuture<String> drainFuture) { | ||
| if (!drainFuture.isDone()) { | ||
| return State.DRAINING; | ||
| } | ||
| try { | ||
| drainFuture.get(); | ||
| return State.DRAINED; | ||
| } catch (InterruptedException e) { | ||
| Thread.currentThread().interrupt(); | ||
| throw new RuntimeException("Fail to drain flink job", e); | ||
| } catch (ExecutionException e) { | ||
| throw new RuntimeException("Fail to drain flink job", e); | ||
| } | ||
| } |
There was a problem hiding this comment.
We can improve the robustness of the drain() method in several ways:
- Support Retries: If a previous drain attempt failed (i.e.,
drainFuture.isCompletedExceptionally()), we should allow the user to retry the drain operation rather than permanently returning the cached failure. - Throw Checked Exceptions: Since
drain()is declared to throwIOException, we should wrap and throwIOExceptioninstead ofRuntimeExceptionwhen the future fails. - Simplify Code: We can inline the
getDrainStatelogic directly intodrain()and remove the helper method entirely, making the code cleaner.
@Override
public synchronized State drain() throws IOException {
CompletableFuture<String> drainFuture = drainSavepointFuture;
if (drainFuture == null || drainFuture.isCompletedExceptionally()) {
drainFuture = this.jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT);
drainSavepointFuture = drainFuture;
}
if (!drainFuture.isDone()) {
return State.DRAINING;
}
try {
drainFuture.get();
return State.DRAINED;
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new IOException("Failed to drain flink job", e);
} catch (ExecutionException e) {
throw new IOException("Failed to drain flink job", e.getCause());
}
}| capitalizedAction, | ||
| state); | ||
| return state; | ||
| } else if (e.getMessage().contains("has terminated")) { |
There was a problem hiding this comment.
If e.getMessage() is null (which can happen with certain IOException subclasses or in test environments), calling contains() on it will throw a NullPointerException, masking the original exception. We should add a null check before calling contains().
| } else if (e.getMessage().contains("has terminated")) { | |
| } else if (e.getMessage() != null && e.getMessage().contains("has terminated")) { |
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
Adds Java SDK pipeline drain support for Dataflow and Flink.
This change adds
PipelineResult#drain()as a default SDK method and introduces two new pipeline states:DRAININGDRAINEDDataflow now requests drain by updating the job requested state to
JOB_STATE_DRAINED, and maps Dataflow drain states precisely:JOB_STATE_DRAINING->PipelineResult.State.DRAININGJOB_STATE_DRAINED->PipelineResult.State.DRAINEDFlink detached results now support drain through
JobClient#stopWithSavepoint(...), returningDRAININGwhile the stop/savepoint operation is pending andDRAINEDafter it completes.addresses #38771
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.