This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 956fe0c90e4 [Flink] Bound SDF self-checkpoint residual state on the
portable runner (#39191)
956fe0c90e4 is described below
commit 956fe0c90e41b0ca9ba9fc5af23b263e0a7f554f
Author: Elia Liu <[email protected]>
AuthorDate: Fri Jul 3 03:42:47 2026 +1000
[Flink] Bound SDF self-checkpoint residual state on the portable runner
(#39191)
A splittable DoFn that self-checkpoints, for example a polling source using
tracker.defer_remainder(), stored each residual under a new unique state
tag,
so every self-checkpoint registered another Flink keyed-state descriptor and
the job leaked heap without bound.
Store all residuals for a key and window under a single stable MapState
keyed
by the checkpoint id, and remove each entry once its timer fires and the
residual is re-processed. The streaming and batch read paths look the
residual
up by timer id.
Fixes #27648
---
CHANGES.md | 1 +
.../FlinkStreamingPortablePipelineTranslator.java | 1 -
.../functions/FlinkExecutableStageFunction.java | 19 ++--
.../FlinkStreamingPortablePipelineTranslator.java | 1 -
.../functions/FlinkExecutableStageFunction.java | 19 ++--
.../streaming/ExecutableStageDoFnOperator.java | 18 ++-
.../control/BundleCheckpointHandlers.java | 24 +++-
.../control/BundleCheckpointHandlersTest.java | 125 +++++++++++++++++++++
8 files changed, 186 insertions(+), 22 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index e20194f9ab0..8bc74b958ba 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -98,6 +98,7 @@
* Fixed GCS filesystem glob matching to correctly handle `/` in object names
and support `**` for recursive matching (Go)
([#38059](https://github.com/apache/beam/issues/38059)).
* Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when
multiple requests share the same enrichment key (Python)
([#38035](https://github.com/apache/beam/issues/38035)).
* Fixed IcebergIO writing manifest column bounds padded with trailing `0x00`
bytes, which broke equality predicate pushdown in some query engines (Java)
([#38580](https://github.com/apache/beam/issues/38580)).
+* Fixed unbounded checkpoint state growth for splittable DoFns that
self-checkpoint on the portable Flink runner (Java)
([#27648](https://github.com/apache/beam/issues/27648)).
## Security Fixes
diff --git
a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
index 647297394f5..efffa4fdbc4 100644
---
a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
+++
b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
@@ -733,7 +733,6 @@ public class FlinkStreamingPortablePipelineTranslator
private <InputT, OutputT> void translateExecutableStage(
String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext
context) {
- // TODO: Fail on splittable DoFns.
// TODO: Special-case single outputs to avoid multiplexing PCollections.
RunnerApi.Components components = pipeline.getComponents();
RunnerApi.PTransform transform = components.getTransformsOrThrow(id);
diff --git
a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
index 1298fd3105a..3f8b5bdfe47 100644
---
a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
+++
b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
@@ -31,7 +31,6 @@ import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.InMemoryStateInternals;
import org.apache.beam.runners.core.InMemoryTimerInternals;
import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
@@ -56,6 +55,7 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.construction.PTransformTranslation;
@@ -284,12 +284,17 @@ public class FlinkExecutableStageFunction<InputT> extends
AbstractRichFunction
List<WindowedValue<InputT>> residuals = new ArrayList<>();
TimerInternals.TimerData timer;
while ((timer = sdfTimerInternals.removeNextProcessingTimer()) !=
null) {
- WindowedValue stateValue =
- sdfStateInternals
- .state(timer.getNamespace(),
StateTags.value(timer.getTimerId(), inputCoder))
- .read();
-
- residuals.add(stateValue);
+ MapState<String, WindowedValue<InputT>> residualState =
+ sdfStateInternals.state(
+ timer.getNamespace(),
+
BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag(
+ inputCoder));
+ WindowedValue<InputT> stateValue =
residualState.get(timer.getTimerId()).read();
+ residualState.remove(timer.getTimerId());
+ // A pre-#27648 savepoint has no entry under the new tag, so the
read is null; drop it.
+ if (stateValue != null) {
+ residuals.add(stateValue);
+ }
}
processElements(residuals, bundle);
}
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
index 77dc4d79516..e4abae4c751 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
@@ -732,7 +732,6 @@ public class FlinkStreamingPortablePipelineTranslator
private <InputT, OutputT> void translateExecutableStage(
String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext
context) {
- // TODO: Fail on splittable DoFns.
// TODO: Special-case single outputs to avoid multiplexing PCollections.
RunnerApi.Components components = pipeline.getComponents();
RunnerApi.PTransform transform = components.getTransformsOrThrow(id);
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
index 6c13f8872e9..f513bda9935 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
@@ -31,7 +31,6 @@ import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.InMemoryStateInternals;
import org.apache.beam.runners.core.InMemoryTimerInternals;
import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
@@ -56,6 +55,7 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.construction.PTransformTranslation;
@@ -284,12 +284,17 @@ public class FlinkExecutableStageFunction<InputT> extends
AbstractRichFunction
List<WindowedValue<InputT>> residuals = new ArrayList<>();
TimerInternals.TimerData timer;
while ((timer = sdfTimerInternals.removeNextProcessingTimer()) !=
null) {
- WindowedValue stateValue =
- sdfStateInternals
- .state(timer.getNamespace(),
StateTags.value(timer.getTimerId(), inputCoder))
- .read();
-
- residuals.add(stateValue);
+ MapState<String, WindowedValue<InputT>> residualState =
+ sdfStateInternals.state(
+ timer.getNamespace(),
+
BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag(
+ inputCoder));
+ WindowedValue<InputT> stateValue =
residualState.get(timer.getTimerId()).read();
+ residualState.remove(timer.getTimerId());
+ // A pre-#27648 savepoint has no entry under the new tag, so the
read is null; drop it.
+ if (stateValue != null) {
+ residuals.add(stateValue);
+ }
}
processElements(residuals, bundle);
}
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
index f7bc62a8ec5..b4d5e43557c 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
@@ -88,6 +88,7 @@ import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.function.ThrowingFunction;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateContext;
import org.apache.beam.sdk.state.TimeDomain;
@@ -1009,9 +1010,20 @@ public class ExecutableStageDoFnOperator<InputT, OutputT>
Preconditions.checkNotNull(remoteBundle, "Call to onTimer outside of a
bundle");
if (StateAndTimerBundleCheckpointHandler.isSdfTimer(timerId)) {
StateNamespace namespace = StateNamespaces.window(windowCoder, window);
- WindowedValue stateValue =
- keyedStateInternals.state(namespace, StateTags.value(timerId,
residualCoder)).read();
- processElement(stateValue);
+ // Residuals from pre-#27648 savepoints used per-residual ValueState
tags and are not
+ // migrated; an older savepoint has no entry under the new tag, so the
read is null and the
+ // residual is dropped.
+ MapState<String, WindowedValue<InputT>> residualState =
+ keyedStateInternals.state(
+ namespace,
StateAndTimerBundleCheckpointHandler.residualStateTag(residualCoder));
+ WindowedValue<InputT> stateValue = residualState.get(timerId).read();
+ // Drop the consumed residual so a polling SDF's keyed state stays
bounded across
+ // self-checkpoints (apache/beam#27648). Flink commits this removal
only with the bundle's
+ // checkpoint, so on failure the timer re-fires and the residual is
re-read.
+ residualState.remove(timerId);
+ if (stateValue != null) {
+ processElement(stateValue);
+ }
} else {
KV<String, String> transformAndTimerFamilyId =
TimerReceiverFactory.decodeTimerDataTimerId(timerFamilyId);
diff --git
a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlers.java
b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlers.java
index 34d952b1560..f869b047f33 100644
---
a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlers.java
+++
b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlers.java
@@ -23,12 +23,15 @@ import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateInternalsFactory;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaces;
+import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternalsFactory;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.IdGenerators;
+import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.CoderUtils;
@@ -48,7 +51,7 @@ public class BundleCheckpointHandlers {
/**
* A {@link BundleCheckpointHandler} which uses {@link
* org.apache.beam.runners.core.TimerInternals.TimerData} and {@link
- * org.apache.beam.sdk.state.ValueState} to reschedule {@link
DelayedBundleApplication}.
+ * org.apache.beam.sdk.state.MapState} to reschedule {@link
DelayedBundleApplication}.
*/
public static class StateAndTimerBundleCheckpointHandler<T> implements
BundleCheckpointHandler {
@@ -82,6 +85,20 @@ public class BundleCheckpointHandlers {
return SDF_PREFIX + ":" + id + ":" + index;
}
+ /** The state id under which all SDF self-checkpoint residuals for a
key/window are stored. */
+ public static final String SDF_RESIDUAL_STATE = "sdf_checkpoint_residuals";
+
+ /**
+ * The single, stable state tag holding every SDF self-checkpoint residual
for a key/window,
+ * keyed by the per-residual checkpoint id. Storing all residuals under
one stable descriptor,
+ * instead of a fresh {@link StateTags#value} per residual, keeps a
polling SDF's keyed state
+ * bounded. See https://github.com/apache/beam/issues/27648.
+ */
+ public static <T> StateTag<MapState<String, WindowedValue<T>>>
residualStateTag(
+ Coder<WindowedValue<T>> residualCoder) {
+ return StateTags.map(SDF_RESIDUAL_STATE, StringUtf8Coder.of(),
residualCoder);
+ }
+
@Override
public void onCheckpoint(ProcessBundleResponse response) {
String id = idGenerator.getId();
@@ -126,8 +143,9 @@ public class BundleCheckpointHandlers {
Instant.ofEpochMilli(outputTimestamp),
TimeDomain.PROCESSING_TIME);
stateInternals
- .state(stateNamespace, StateTags.value(tag, residualCoder))
- .write(
+ .state(stateNamespace, residualStateTag(residualCoder))
+ .put(
+ tag,
WindowedValues.of(
stateValue.getValue(),
stateValue.getTimestamp(),
diff --git
a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlersTest.java
b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlersTest.java
new file mode 100644
index 00000000000..24714dc498d
--- /dev/null
+++
b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlersTest.java
@@ -0,0 +1,125 @@
+/*
+ * 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.runners.fnexecution.control;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.hasSize;
+
+import java.util.HashSet;
+import java.util.Set;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse;
+import org.apache.beam.runners.core.InMemoryStateInternals;
+import org.apache.beam.runners.core.InMemoryTimerInternals;
+import org.apache.beam.runners.core.StateInternals;
+import org.apache.beam.runners.core.StateInternalsFactory;
+import org.apache.beam.runners.core.StateNamespace;
+import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.core.TimerInternalsFactory;
+import
org.apache.beam.runners.fnexecution.control.BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.state.State;
+import org.apache.beam.sdk.state.StateContext;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.util.CoderUtils;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link BundleCheckpointHandlers}. */
+@RunWith(JUnit4.class)
+public class BundleCheckpointHandlersTest {
+
+ private static final Coder<WindowedValue<String>> RESIDUAL_CODER =
+ WindowedValues.getFullCoder(StringUtf8Coder.of(),
GlobalWindow.Coder.INSTANCE);
+
+ /**
+ * Regression test for https://github.com/apache/beam/issues/27648.
+ *
+ * <p>A polling SDF self-checkpoints on every poll via {@code
tracker.defer_remainder()}. Each
+ * self-checkpoint used to store its residual under a brand new, unique
state tag
+ * ("sdf_checkpoint:<id>:<index>"), so every poll registered
another keyed-state
+ * descriptor and the job leaked heap without bound. All residuals for a
key/window must instead
+ * live under a single, stable state descriptor.
+ */
+ @Test
+ public void repeatedSelfCheckpointsUseBoundedStateDescriptors() throws
Exception {
+ Set<String> stateDescriptorIds = new HashSet<>();
+ StateInternalsFactory<String> stateInternalsFactory =
+ key -> new RecordingStateInternals(InMemoryStateInternals.forKey(key),
stateDescriptorIds);
+ TimerInternalsFactory<String> timerInternalsFactory = key -> new
InMemoryTimerInternals();
+
+ StateAndTimerBundleCheckpointHandler<String> handler =
+ new StateAndTimerBundleCheckpointHandler<>(
+ timerInternalsFactory,
+ stateInternalsFactory,
+ RESIDUAL_CODER,
+ GlobalWindow.Coder.INSTANCE);
+
+ int selfCheckpoints = 100;
+ for (int i = 0; i < selfCheckpoints; i++) {
+ handler.onCheckpoint(residualResponse("key"));
+ }
+
+ // One stable descriptor for all residuals, no matter how many times the
SDF checkpointed.
+ assertThat(stateDescriptorIds, hasSize(1));
+ }
+
+ private static ProcessBundleResponse residualResponse(String key) throws
Exception {
+ // The residual element's value is used as the SDF key.
+ byte[] encodedElement =
+ CoderUtils.encodeToByteArray(RESIDUAL_CODER,
WindowedValues.valueInGlobalWindow(key));
+ return ProcessBundleResponse.newBuilder()
+ .addResidualRoots(
+ DelayedBundleApplication.newBuilder()
+ .setApplication(
+ BundleApplication.newBuilder()
+ .setElement(ByteString.copyFrom(encodedElement))
+ .build())
+ .build())
+ .build();
+ }
+
+ /** A {@link StateInternals} that records every state descriptor id it is
asked for. */
+ private static class RecordingStateInternals implements StateInternals {
+ private final StateInternals delegate;
+ private final Set<String> descriptorIds;
+
+ RecordingStateInternals(StateInternals delegate, Set<String>
descriptorIds) {
+ this.delegate = delegate;
+ this.descriptorIds = descriptorIds;
+ }
+
+ @Override
+ public Object getKey() {
+ return delegate.getKey();
+ }
+
+ @Override
+ public <T extends State> T state(
+ StateNamespace namespace, StateTag<T> address, StateContext<?> c) {
+ descriptorIds.add(address.getId());
+ return delegate.state(namespace, address, c);
+ }
+ }
+}
