Skip to content

Commit a42f34d

Browse files
EliaaazzzEliaaazzz
committed
[Flink] Bound SDF self-checkpoint residual state on the portable runner
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
1 parent 8c0f0e0 commit a42f34d

5 files changed

Lines changed: 174 additions & 13 deletions

File tree

‎CHANGES.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@
9898
* 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)).
9999
* Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)).
100100
* 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)).
101+
* 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)).
101102

102103
## Security Fixes
103104

‎runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
import org.apache.beam.runners.core.InMemoryStateInternals;
3232
import org.apache.beam.runners.core.InMemoryTimerInternals;
3333
import org.apache.beam.runners.core.StateInternals;
34-
import org.apache.beam.runners.core.StateTags;
3534
import org.apache.beam.runners.core.TimerInternals;
3635
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
3736
import org.apache.beam.runners.flink.FlinkPipelineOptions;
@@ -56,6 +55,7 @@
5655
import org.apache.beam.sdk.fn.data.FnDataReceiver;
5756
import org.apache.beam.sdk.io.FileSystems;
5857
import org.apache.beam.sdk.options.PipelineOptions;
58+
import org.apache.beam.sdk.state.MapState;
5959
import org.apache.beam.sdk.transforms.join.RawUnionValue;
6060
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
6161
import org.apache.beam.sdk.util.construction.PTransformTranslation;
@@ -284,12 +284,17 @@ public void mapPartition(
284284
List<WindowedValue<InputT>> residuals = new ArrayList<>();
285285
TimerInternals.TimerData timer;
286286
while ((timer = sdfTimerInternals.removeNextProcessingTimer()) != null) {
287-
WindowedValue stateValue =
288-
sdfStateInternals
289-
.state(timer.getNamespace(), StateTags.value(timer.getTimerId(), inputCoder))
290-
.read();
291-
292-
residuals.add(stateValue);
287+
MapState<String, WindowedValue<InputT>> residualState =
288+
sdfStateInternals.state(
289+
timer.getNamespace(),
290+
BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag(
291+
inputCoder));
292+
WindowedValue<InputT> stateValue = residualState.get(timer.getTimerId()).read();
293+
residualState.remove(timer.getTimerId());
294+
// A pre-#27648 savepoint has no entry under the new tag, so the read is null; drop it.
295+
if (stateValue != null) {
296+
residuals.add(stateValue);
297+
}
293298
}
294299
processElements(residuals, bundle);
295300
}

‎runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@
8888
import org.apache.beam.sdk.function.ThrowingFunction;
8989
import org.apache.beam.sdk.options.PipelineOptions;
9090
import org.apache.beam.sdk.state.BagState;
91+
import org.apache.beam.sdk.state.MapState;
9192
import org.apache.beam.sdk.state.State;
9293
import org.apache.beam.sdk.state.StateContext;
9394
import org.apache.beam.sdk.state.TimeDomain;
@@ -1009,9 +1010,20 @@ public <KeyT> void onTimer(
10091010
Preconditions.checkNotNull(remoteBundle, "Call to onTimer outside of a bundle");
10101011
if (StateAndTimerBundleCheckpointHandler.isSdfTimer(timerId)) {
10111012
StateNamespace namespace = StateNamespaces.window(windowCoder, window);
1012-
WindowedValue stateValue =
1013-
keyedStateInternals.state(namespace, StateTags.value(timerId, residualCoder)).read();
1014-
processElement(stateValue);
1013+
// Residuals from pre-#27648 savepoints used per-residual ValueState tags and are not
1014+
// migrated; an older savepoint has no entry under the new tag, so the read is null and the
1015+
// residual is dropped.
1016+
MapState<String, WindowedValue<InputT>> residualState =
1017+
keyedStateInternals.state(
1018+
namespace, StateAndTimerBundleCheckpointHandler.residualStateTag(residualCoder));
1019+
WindowedValue<InputT> stateValue = residualState.get(timerId).read();
1020+
// Drop the consumed residual so a polling SDF's keyed state stays bounded across
1021+
// self-checkpoints (apache/beam#27648). Flink commits this removal only with the bundle's
1022+
// checkpoint, so on failure the timer re-fires and the residual is re-read.
1023+
residualState.remove(timerId);
1024+
if (stateValue != null) {
1025+
processElement(stateValue);
1026+
}
10151027
} else {
10161028
KV<String, String> transformAndTimerFamilyId =
10171029
TimerReceiverFactory.decodeTimerDataTimerId(timerFamilyId);

‎runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlers.java‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,15 @@
2323
import org.apache.beam.runners.core.StateInternalsFactory;
2424
import org.apache.beam.runners.core.StateNamespace;
2525
import org.apache.beam.runners.core.StateNamespaces;
26+
import org.apache.beam.runners.core.StateTag;
2627
import org.apache.beam.runners.core.StateTags;
2728
import org.apache.beam.runners.core.TimerInternals;
2829
import org.apache.beam.runners.core.TimerInternalsFactory;
2930
import org.apache.beam.sdk.coders.Coder;
31+
import org.apache.beam.sdk.coders.StringUtf8Coder;
3032
import org.apache.beam.sdk.fn.IdGenerator;
3133
import org.apache.beam.sdk.fn.IdGenerators;
34+
import org.apache.beam.sdk.state.MapState;
3235
import org.apache.beam.sdk.state.TimeDomain;
3336
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
3437
import org.apache.beam.sdk.util.CoderUtils;
@@ -48,7 +51,7 @@ public class BundleCheckpointHandlers {
4851
/**
4952
* A {@link BundleCheckpointHandler} which uses {@link
5053
* org.apache.beam.runners.core.TimerInternals.TimerData} and {@link
51-
* org.apache.beam.sdk.state.ValueState} to reschedule {@link DelayedBundleApplication}.
54+
* org.apache.beam.sdk.state.MapState} to reschedule {@link DelayedBundleApplication}.
5255
*/
5356
public static class StateAndTimerBundleCheckpointHandler<T> implements BundleCheckpointHandler {
5457

@@ -82,6 +85,20 @@ private static String constructSdfCheckpointId(String id, int index) {
8285
return SDF_PREFIX + ":" + id + ":" + index;
8386
}
8487

88+
/** The state id under which all SDF self-checkpoint residuals for a key/window are stored. */
89+
public static final String SDF_RESIDUAL_STATE = "sdf_checkpoint_residuals";
90+
91+
/**
92+
* The single, stable state tag holding every SDF self-checkpoint residual for a key/window,
93+
* keyed by the per-residual checkpoint id. Storing all residuals under one stable descriptor,
94+
* instead of a fresh {@link StateTags#value} per residual, keeps a polling SDF's keyed state
95+
* bounded. See https://github.com/apache/beam/issues/27648.
96+
*/
97+
public static <T> StateTag<MapState<String, WindowedValue<T>>> residualStateTag(
98+
Coder<WindowedValue<T>> residualCoder) {
99+
return StateTags.map(SDF_RESIDUAL_STATE, StringUtf8Coder.of(), residualCoder);
100+
}
101+
85102
@Override
86103
public void onCheckpoint(ProcessBundleResponse response) {
87104
String id = idGenerator.getId();
@@ -126,8 +143,9 @@ public void onCheckpoint(ProcessBundleResponse response) {
126143
Instant.ofEpochMilli(outputTimestamp),
127144
TimeDomain.PROCESSING_TIME);
128145
stateInternals
129-
.state(stateNamespace, StateTags.value(tag, residualCoder))
130-
.write(
146+
.state(stateNamespace, residualStateTag(residualCoder))
147+
.put(
148+
tag,
131149
WindowedValues.of(
132150
stateValue.getValue(),
133151
stateValue.getTimestamp(),

‎runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/BundleCheckpointHandlersTest.java‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.fnexecution.control;
19+
20+
import static org.hamcrest.MatcherAssert.assertThat;
21+
import static org.hamcrest.Matchers.hasSize;
22+
23+
import java.util.HashSet;
24+
import java.util.Set;
25+
import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
26+
import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
27+
import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse;
28+
import org.apache.beam.runners.core.InMemoryStateInternals;
29+
import org.apache.beam.runners.core.InMemoryTimerInternals;
30+
import org.apache.beam.runners.core.StateInternals;
31+
import org.apache.beam.runners.core.StateInternalsFactory;
32+
import org.apache.beam.runners.core.StateNamespace;
33+
import org.apache.beam.runners.core.StateTag;
34+
import org.apache.beam.runners.core.TimerInternalsFactory;
35+
import org.apache.beam.runners.fnexecution.control.BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler;
36+
import org.apache.beam.sdk.coders.Coder;
37+
import org.apache.beam.sdk.coders.StringUtf8Coder;
38+
import org.apache.beam.sdk.state.State;
39+
import org.apache.beam.sdk.state.StateContext;
40+
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
41+
import org.apache.beam.sdk.util.CoderUtils;
42+
import org.apache.beam.sdk.values.WindowedValue;
43+
import org.apache.beam.sdk.values.WindowedValues;
44+
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
45+
import org.junit.Test;
46+
import org.junit.runner.RunWith;
47+
import org.junit.runners.JUnit4;
48+
49+
/** Tests for {@link BundleCheckpointHandlers}. */
50+
@RunWith(JUnit4.class)
51+
public class BundleCheckpointHandlersTest {
52+
53+
private static final Coder<WindowedValue<String>> RESIDUAL_CODER =
54+
WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE);
55+
56+
/**
57+
* Regression test for https://github.com/apache/beam/issues/27648.
58+
*
59+
* <p>A polling SDF self-checkpoints on every poll via {@code tracker.defer_remainder()}. Each
60+
* self-checkpoint used to store its residual under a brand new, unique state tag
61+
* ("sdf_checkpoint:&lt;id&gt;:&lt;index&gt;"), so every poll registered another keyed-state
62+
* descriptor and the job leaked heap without bound. All residuals for a key/window must instead
63+
* live under a single, stable state descriptor.
64+
*/
65+
@Test
66+
public void repeatedSelfCheckpointsUseBoundedStateDescriptors() throws Exception {
67+
Set<String> stateDescriptorIds = new HashSet<>();
68+
StateInternalsFactory<String> stateInternalsFactory =
69+
key -> new RecordingStateInternals(InMemoryStateInternals.forKey(key), stateDescriptorIds);
70+
TimerInternalsFactory<String> timerInternalsFactory = key -> new InMemoryTimerInternals();
71+
72+
StateAndTimerBundleCheckpointHandler<String> handler =
73+
new StateAndTimerBundleCheckpointHandler<>(
74+
timerInternalsFactory,
75+
stateInternalsFactory,
76+
RESIDUAL_CODER,
77+
GlobalWindow.Coder.INSTANCE);
78+
79+
int selfCheckpoints = 100;
80+
for (int i = 0; i < selfCheckpoints; i++) {
81+
handler.onCheckpoint(residualResponse("key"));
82+
}
83+
84+
// One stable descriptor for all residuals, no matter how many times the SDF checkpointed.
85+
assertThat(stateDescriptorIds, hasSize(1));
86+
}
87+
88+
private static ProcessBundleResponse residualResponse(String key) throws Exception {
89+
// The residual element's value is used as the SDF key.
90+
byte[] encodedElement =
91+
CoderUtils.encodeToByteArray(RESIDUAL_CODER, WindowedValues.valueInGlobalWindow(key));
92+
return ProcessBundleResponse.newBuilder()
93+
.addResidualRoots(
94+
DelayedBundleApplication.newBuilder()
95+
.setApplication(
96+
BundleApplication.newBuilder()
97+
.setElement(ByteString.copyFrom(encodedElement))
98+
.build())
99+
.build())
100+
.build();
101+
}
102+
103+
/** A {@link StateInternals} that records every state descriptor id it is asked for. */
104+
private static class RecordingStateInternals implements StateInternals {
105+
private final StateInternals delegate;
106+
private final Set<String> descriptorIds;
107+
108+
RecordingStateInternals(StateInternals delegate, Set<String> descriptorIds) {
109+
this.delegate = delegate;
110+
this.descriptorIds = descriptorIds;
111+
}
112+
113+
@Override
114+
public Object getKey() {
115+
return delegate.getKey();
116+
}
117+
118+
@Override
119+
public <T extends State> T state(
120+
StateNamespace namespace, StateTag<T> address, StateContext<?> c) {
121+
descriptorIds.add(address.getId());
122+
return delegate.state(namespace, address, c);
123+
}
124+
}
125+
}

0 commit comments

Comments
 (0)