diff --git a/CHANGES.md b/CHANGES.md index e20194f9ab09..8bc74b958bad 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 647297394f52..efffa4fdbc4a 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 @@ private void translateStreamingImpulse( private 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 1298fd3105aa..3f8b5bdfe473 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.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.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 void mapPartition( List> 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> residualState = + sdfStateInternals.state( + timer.getNamespace(), + BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag( + inputCoder)); + WindowedValue 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 77dc4d795167..e4abae4c751d 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 @@ private void translateStreamingImpulse( private 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 6c13f8872e9d..f513bda9935a 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.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.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 void mapPartition( List> 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> residualState = + sdfStateInternals.state( + timer.getNamespace(), + BundleCheckpointHandlers.StateAndTimerBundleCheckpointHandler.residualStateTag( + inputCoder)); + WindowedValue 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 f7bc62a8ec5a..b4d5e43557c1 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.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 void onTimer( 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> residualState = + keyedStateInternals.state( + namespace, StateAndTimerBundleCheckpointHandler.residualStateTag(residualCoder)); + WindowedValue 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 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 34d952b15608..f869b047f339 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.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 implements BundleCheckpointHandler { @@ -82,6 +85,20 @@ private static String constructSdfCheckpointId(String id, int index) { 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 StateTag>> residualStateTag( + Coder> 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 void onCheckpoint(ProcessBundleResponse response) { 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 000000000000..24714dc498d6 --- /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> RESIDUAL_CODER = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + + /** + * Regression test for https://github.com/apache/beam/issues/27648. + * + *

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 stateDescriptorIds = new HashSet<>(); + StateInternalsFactory stateInternalsFactory = + key -> new RecordingStateInternals(InMemoryStateInternals.forKey(key), stateDescriptorIds); + TimerInternalsFactory timerInternalsFactory = key -> new InMemoryTimerInternals(); + + StateAndTimerBundleCheckpointHandler 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 descriptorIds; + + RecordingStateInternals(StateInternals delegate, Set descriptorIds) { + this.delegate = delegate; + this.descriptorIds = descriptorIds; + } + + @Override + public Object getKey() { + return delegate.getKey(); + } + + @Override + public T state( + StateNamespace namespace, StateTag address, StateContext c) { + descriptorIds.add(address.getId()); + return delegate.state(namespace, address, c); + } + } +}