rkhachatryan commented on code in PR #28001:
URL: https://github.com/apache/flink/pull/28001#discussion_r3144162459
##########
flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java:
##########
@@ -162,6 +238,11 @@ public void close() throws IOException {
for (final InputGate inputGate : inputGates) {
inputGate.finishReadRecoveredState();
}
+ if (preFilterSegment != null) {
+ preFilterSegment.free();
+ preFilterSegment = null;
+ preFilterBufferInUse = false;
+ }
Review Comment:
Should we check `preFilterBufferInUse` here to avoid double-release?
##########
flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java:
##########
@@ -82,24 +88,94 @@ class InputChannelRecoveredStateHandler
*/
@Nullable private final ChannelStateFilteringHandler filteringHandler;
+ /** Network buffer memory segment size in bytes. Used to size the reusable
pre-filter buffer. */
+ private final int memorySegmentSize;
+
+ /**
+ * Reusable heap memory segment backing the pre-filter buffer in filtering
mode. Lazily
+ * allocated on the first {@link #getPreFilterBuffer} call, reused for
every subsequent call,
+ * and freed in {@link #close()}.
+ *
+ * <p>Reuse is safe because at most one pre-filter buffer is in flight per
task at any moment.
+ * This invariant is enforced at runtime by {@link #preFilterBufferInUse}.
+ */
+ @Nullable private MemorySegment preFilterSegment;
+
+ /**
+ * Tracks whether {@link #preFilterSegment} is currently wrapped by a live
{@link Buffer} that
+ * has not yet been recycled. Flipped to {@code true} when a new buffer is
issued, and flipped
+ * back to {@code false} by the custom {@link BufferRecycler} when the
buffer is recycled.
+ */
+ private boolean preFilterBufferInUse;
+
InputChannelRecoveredStateHandler(
InputGate[] inputGates,
InflightDataRescalingDescriptor channelMapping,
- @Nullable ChannelStateFilteringHandler filteringHandler) {
+ @Nullable ChannelStateFilteringHandler filteringHandler,
+ int memorySegmentSize) {
this.inputGates = inputGates;
this.channelMapping = channelMapping;
this.filteringHandler = filteringHandler;
+ checkArgument(
+ memorySegmentSize > 0, "memorySegmentSize must be positive:
%s", memorySegmentSize);
+ this.memorySegmentSize = memorySegmentSize;
}
@Override
public BufferWithContext<Buffer> getBuffer(InputChannelInfo channelInfo)
throws IOException, InterruptedException {
- // request the buffer from any mapped channel as they all will receive
the same buffer
+ if (filteringHandler != null) {
+ return getPreFilterBuffer();
+ }
+ // Non-filtering mode: use existing network buffer pool allocation.
RecoveredInputChannel channel = getMappedChannels(channelInfo);
Buffer buffer = channel.requestBufferBlocking();
return new BufferWithContext<>(wrap(buffer), buffer);
}
+ /**
+ * Allocates a pre-filter buffer from a reusable heap segment (isolated
from the Network Buffer
+ * Pool) in filtering mode.
+ *
+ * <p>Memory management: a single {@link MemorySegment} per task is lazily
allocated on first
+ * invocation and reused across every subsequent call. The custom {@link
BufferRecycler} does
+ * not free the segment — it only flips {@link #preFilterBufferInUse} back
to {@code false} so
+ * the next call can reuse it. The segment itself is freed in {@link
#close()}.
+ *
+ * <p>Runtime invariant check: the one-at-a-time invariant on pre-filter
buffers is guaranteed
+ * by Flink's serial recovery loop and the deserializer's ownership
contract. This method
+ * asserts the invariant before issuing a buffer: if a previously issued
buffer has not yet been
+ * recycled, it throws {@link IllegalStateException} so any future
regression fails loudly
+ * instead of silently corrupting memory.
+ */
+ private BufferWithContext<Buffer> getPreFilterBuffer() {
+ checkState(
+ !preFilterBufferInUse,
+ "Previous pre-filter buffer has not been recycled. This
violates the "
+ + "one-buffer-at-a-time invariant of pre-filter
buffers.");
+
+ if (preFilterSegment == null) {
+ preFilterSegment =
MemorySegmentFactory.allocateUnpooledSegment(memorySegmentSize);
+ }
+ preFilterBufferInUse = true;
+
+ // The recycler keeps the segment alive for reuse; only flips the
in-use flag.
+ BufferRecycler recycler = segment -> preFilterBufferInUse = false;
Review Comment:
Do we need any thread-safety mechanisms for the recycler?
##########
flink-tests/src/test/java/org/apache/flink/test/checkpointing/RecoveredStateFilteringLargeRecordITCase.java:
##########
@@ -0,0 +1,225 @@
+/*
+ * 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.flink.test.checkpointing;
+
+import org.apache.flink.api.common.JobStatus;
+import org.apache.flink.api.common.eventtime.WatermarkStrategy;
+import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.connector.source.lib.NumberSequenceSource;
+import
org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.CheckpointingOptions;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.ExternalizedCheckpointRetention;
+import org.apache.flink.configuration.MemorySize;
+import org.apache.flink.configuration.RestartStrategyOptions;
+import org.apache.flink.configuration.StateRecoveryOptions;
+import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.connector.datagen.source.DataGeneratorSource;
+import org.apache.flink.core.execution.JobClient;
+import org.apache.flink.runtime.memory.MemoryManager;
+import org.apache.flink.runtime.minicluster.MiniCluster;
+import org.apache.flink.runtime.testutils.CommonTestUtils;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink;
+import org.apache.flink.test.junit5.InjectMiniCluster;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.util.TestLoggerExtension;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+import javax.annotation.Nullable;
+
+import java.io.File;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static
org.apache.flink.configuration.RestartStrategyOptions.RestartStrategyType.NO_RESTART_STRATEGY;
+
+/**
+ * Integration test: recovery with filtering must keep at most one pre-filter
source buffer in
+ * flight per task, even when records are large enough to span many source
buffers.
+ *
+ * <p>Strategy: configure the network memory segment to Flink's minimum page
size and generate
+ * records several times larger. Every record therefore forces the recovery
loop to allocate many
+ * consecutive pre-filter buffers, exercising the serialized getBuffer/recycle
cycle that must reuse
+ * a single heap segment. A sleeping downstream operator induces back-pressure
on the keyBy shuffle
+ * so that inflight buffers actually accumulate and are captured by the
unaligned checkpoint;
+ * otherwise the recovery path would have nothing to filter.
+ *
+ * <p>Filtering during recovery is gated by two config options. Both are
pinned on explicitly in
+ * {@link #getEnv} so the test deterministically exercises the filtering
handler.
+ *
+ * <p>If the one-at-a-time invariant were violated during recovery, the
runtime check in {@code
+ * InputChannelRecoveredStateHandler.getPreFilterBuffer} would fail with {@link
+ * IllegalStateException} and the job would fail. Reaching a completed
post-recovery checkpoint is
+ * therefore proof that the invariant held across all buffer cycles.
+ */
+@ExtendWith({TestLoggerExtension.class})
+class RecoveredStateFilteringLargeRecordITCase {
Review Comment:
This test still passes locally for me when I disable the allocation of
unpooled segments in `InputChannelRecoveredStateHandler.getBuffer` - is that
expected? Shouldn't it deadlock?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
