This is an automated email from the ASF dual-hosted git repository. 1996fanrui pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit 86099d0c6142bfc6a8fbbab41decd99d53fc9bc3 Author: Rui Fan <[email protected]> AuthorDate: Mon Jul 6 01:58:14 2026 +0200 [FLINK-39524][checkpoint] Rewrite ChannelStateFilteringHandler to emit into a spill segment filterAndRewrite now appends each surviving record (length-prefixed via a 4B placeholder + writeIntUnsafe backfill) into a caller-provided DataOutputSerializer: no network buffers, no InterruptedException. The deserializer creation is inlined (always filterContext.getTmpDirectories(); the java.io.tmpdir fallback is gone). New SpillingWithFilteringHandler: getBuffer() is the reusable heap pre-filter-segment logic (verbatim from the v1 FilteringHandler); recover() routes filter output into the spill segment of the mapped channel. The v1 FilteringHandler stays present and factory-selected -- it is still the memory backend. Since it can no longer call the rewritten method, the old buffer-delivering filter loop (BufferSupplier, List<Buffer> overloads, the chunking machinery and its two GateFilterHandler fields) is kept as clearly marked FLINK-38544-transitional overloads; they die together with the class when the spilling backend lands. RecoveredStateFilteringLargeRecordITCase (explicitly flag-on, v1 path) still passes after this commit. AbstractInputChannelRecoveredStateHandler.create(...) gains the final String[] spillTmpDirectories parameter (unused by the transitional in-memory branches); SequentialChannelStateReaderImpl passes filterContext.getTmpDirectories() mechanically. Tests: ChannelStateFilteringHandlerTest, GateFilterHandlerTest and GateFilterHandlerBufferOwnershipTest (rewritten against the serializer sink), RecoveredChannelStateHandlerFilterRoutingTest, and the InputChannelRecoveredStateHandlerTest adaptation. The three helpers that in final form obtain spilling handlers through the factory construct them directly for now (transitional; the factory-switch commit restores the factory calls). --- .../channel/ChannelStateFilteringHandler.java | 127 +++++++++++++++++-- .../channel/RecoveredChannelStateHandler.java | 137 ++++++++++++++++++++- .../channel/SequentialChannelStateReaderImpl.java | 20 ++- .../channel/ChannelStateFilteringHandlerTest.java | 79 ++++++++++++ .../GateFilterHandlerBufferOwnershipTest.java | 112 ++++++----------- .../checkpoint/channel/GateFilterHandlerTest.java | 124 ++++++++++--------- .../InputChannelRecoveredStateHandlerTest.java | 86 +++++++++++-- 7 files changed, 516 insertions(+), 169 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandler.java index b257c3b4054..ef3db27bbff 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandler.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandler.java @@ -101,7 +101,43 @@ public class ChannelStateFilteringHandler implements Closeable { } /** - * Filters a recovered buffer from the specified virtual channel, returning new buffers + * Filters {@code sourceBuffer} through the virtual channel identified by {@code gateIndex} / + * {@code oldChannelIndex}, appending each surviving record (length-prefixed) into {@code + * outputSerializer}. One call may emit 0..N records depending on the filter result and whether + * records spanning previous buffers complete here. The caller owns the segment boundary. + */ + public void filterAndRewrite( + int gateIndex, + int oldSubtaskIndex, + int oldChannelIndex, + Buffer sourceBuffer, + DataOutputSerializer outputSerializer) + throws IOException { + + if (gateIndex < 0 || gateIndex >= gateHandlers.length) { + throw new IllegalStateException( + "Invalid gateIndex: " + + gateIndex + + ", number of gates: " + + gateHandlers.length); + } + + GateFilterHandler<?> gateHandler = gateHandlers[gateIndex]; + if (gateHandler == null) { + throw new IllegalStateException( + "No handler for gateIndex " + + gateIndex + + ". This gate is not a network input and should not have recovered buffers."); + } + gateHandler.filterAndRewrite( + oldSubtaskIndex, oldChannelIndex, sourceBuffer, outputSerializer); + } + + /** + * FLINK-38544 transitional: removed when the spilling backend lands (dies together with the + * in-memory {@code FilteringHandler}, its only caller). + * + * <p>Filters a recovered buffer from the specified virtual channel, returning new buffers * containing only the records that belong to the current subtask. * * <p>One source buffer may produce 0 to N result buffers: 0 if all records are filtered out, @@ -215,7 +251,8 @@ public class ChannelStateFilteringHandler implements Closeable { : VirtualChannelRecordFilterFactory.createPassThroughFilter(); RecordDeserializer<DeserializationDelegate<StreamElement>> deserializer = - createDeserializer(filterContext.getTmpDirectories()); + new SpillingAdaptiveSpanningRecordDeserializer<>( + filterContext.getTmpDirectories()); VirtualChannel<T> vc = new VirtualChannel<>(deserializer, recordFilter); gateVirtualChannels.put(key, vc); @@ -246,21 +283,14 @@ public class ChannelStateFilteringHandler implements Closeable { return oldIndexes.stream().mapToInt(Integer::intValue).toArray(); } - private static RecordDeserializer<DeserializationDelegate<StreamElement>> createDeserializer( - String[] tmpDirectories) { - if (tmpDirectories != null && tmpDirectories.length > 0) { - return new SpillingAdaptiveSpanningRecordDeserializer<>(tmpDirectories); - } else { - String[] defaultDirs = new String[] {System.getProperty("java.io.tmpdir")}; - return new SpillingAdaptiveSpanningRecordDeserializer<>(defaultDirs); - } - } - // ------------------------------------------------------------------------------------------- // Inner classes // ------------------------------------------------------------------------------------------- - /** Provides buffers for re-serializing filtered records. Implementations may block. */ + /** + * FLINK-38544 transitional: removed when the spilling backend lands. Provides buffers for + * re-serializing filtered records on the in-memory delivery path. Implementations may block. + */ @FunctionalInterface public interface BufferSupplier { Buffer requestBufferBlocking() throws IOException, InterruptedException; @@ -275,6 +305,9 @@ public class ChannelStateFilteringHandler implements Closeable { private final Map<SubtaskConnectionDescriptor, VirtualChannel<T>> virtualChannels; private final StreamElementSerializer<T> serializer; private final DeserializationDelegate<StreamElement> deserializationDelegate; + + // FLINK-38544 transitional: removed when the spilling backend lands (used only by the + // in-memory buffer-delivering filterAndRewrite overload below). private final DataOutputSerializer outputSerializer; private final byte[] lengthBuffer = new byte[4]; @@ -287,6 +320,72 @@ public class ChannelStateFilteringHandler implements Closeable { this.outputSerializer = new DataOutputSerializer(128); } + /** + * Deserializes records from {@code sourceBuffer}, applies the virtual channel's record + * filter, and re-serializes each surviving record into {@code outputSerializer}. No + * intermediate network buffer is used; the caller owns the segment boundary. + */ + void filterAndRewrite( + int oldSubtaskIndex, + int oldChannelIndex, + Buffer sourceBuffer, + DataOutputSerializer outputSerializer) + throws IOException { + + boolean sourceBufferOwnershipTransferred = false; + try { + SubtaskConnectionDescriptor key = + new SubtaskConnectionDescriptor(oldSubtaskIndex, oldChannelIndex); + VirtualChannel<T> vc = virtualChannels.get(key); + if (vc == null) { + throw new IllegalStateException( + "No VirtualChannel found for key: " + + key + + "; known channels are " + + virtualChannels.keySet()); + } + + vc.setNextBuffer(sourceBuffer); + sourceBufferOwnershipTransferred = true; + + while (true) { + DeserializationResult result = vc.getNextRecord(deserializationDelegate); + if (result.isFullRecord()) { + serializeElement(deserializationDelegate.getInstance(), outputSerializer); + } + if (result.isBufferConsumed()) { + break; + } + } + } catch (Throwable t) { + if (!sourceBufferOwnershipTransferred) { + sourceBuffer.recycleBuffer(); + } + throw t; + } + } + + /** + * Appends one stream element as a length-prefixed record. Reserves the 4B prefix, + * serializes the element, then backfills the length, because {@code outputSerializer} + * already holds the segment header and earlier records, so the prefix cannot be written + * from a fixed offset. + */ + private void serializeElement(StreamElement element, DataOutputSerializer outputSerializer) + throws IOException { + int startPos = outputSerializer.length(); + outputSerializer.writeInt(0); // length placeholder + serializer.serialize(element, outputSerializer); + int recordLength = outputSerializer.length() - startPos - Integer.BYTES; + outputSerializer.writeIntUnsafe(recordLength, startPos); + } + + // ----------------------------------------------------------------------------------- + // FLINK-38544 transitional: everything below until hasPartialData() is the in-memory + // buffer-delivering filter path, kept alive for the v1 FilteringHandler; removed when + // the spilling backend lands. + // ----------------------------------------------------------------------------------- + /** * Deserializes records from {@code sourceBuffer}, applies the virtual channel's record * filter, and immediately re-serializes each surviving record into output buffers. @@ -447,6 +546,8 @@ public class ChannelStateFilteringHandler implements Closeable { return currentBuffer; } + // ------------------------ end of FLINK-38544 transitional block ----------------------- + boolean hasPartialData() { return virtualChannels.values().stream().anyMatch(VirtualChannel::hasPartialData); } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java index 98660963afe..4a5c005c303 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java @@ -131,12 +131,13 @@ abstract class AbstractInputChannelRecoveredStateHandler InflightDataRescalingDescriptor channelMapping, boolean checkpointingDuringRecoveryEnabled, @Nullable ChannelStateFilteringHandler filteringHandler, - int memorySegmentSize) { + int memorySegmentSize, + String[] spillTmpDirectories) { if (!checkpointingDuringRecoveryEnabled) { return new NoSpillingHandler(inputGates, channelMapping, false); } // FLINK-38544 transitional: the flag-on path still uses the in-memory handlers until the - // spilling backend lands. + // spilling backend lands; spillTmpDirectories is unused until then. if (filteringHandler == null) { return new NoSpillingHandler(inputGates, channelMapping, true); } @@ -410,6 +411,8 @@ abstract class AbstractSpillingHandler extends AbstractInputChannelRecoveredStat files.add(filePath); } + // TODO: FLINK-38544 — wire this into SequentialChannelStateReaderImpl#readInputData when the + // handler factory starts selecting the Spilling* handlers; unused in production until then. @Override @Nullable FetchedChannelState getProducedChannelState() { @@ -496,6 +499,136 @@ class SpillingNoFilteringHandler extends AbstractSpillingHandler { /** * Recovery handler for the case where checkpointing during recovery is enabled and a filtering * handler is present. Uses a reusable heap-backed pre-filter buffer (isolated from the Network + * Buffer Pool) and writes filtered/rewritten output to the spill file via {@link + * ChannelStateFilteringHandler#filterAndRewrite}. + */ +class SpillingWithFilteringHandler extends AbstractSpillingHandler { + + 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 #getBuffer} call, reused for every subsequent call, and freed + * in {@link #closeInternal()}. + * + * <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; + + SpillingWithFilteringHandler( + InputGate[] inputGates, + InflightDataRescalingDescriptor channelMapping, + ChannelStateFilteringHandler filteringHandler, + int memorySegmentSize, + String[] spillTmpDirectories) { + super( + inputGates, + channelMapping, + spillTmpDirectories, + DEFAULT_SPILL_FILE_SIZE_BYTES, + DEFAULT_MAX_SEGMENT_SIZE_BYTES); + this.filteringHandler = filteringHandler; + checkArgument( + memorySegmentSize > 0, "memorySegmentSize must be positive: %s", memorySegmentSize); + this.memorySegmentSize = memorySegmentSize; + } + + /** + * 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 #closeInternal()}. + * + * <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. + */ + @Override + public BufferWithContext<Buffer> getBuffer(InputChannelInfo channelInfo) { + 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; + Buffer buffer = new NetworkBuffer(preFilterSegment, recycler); + return new BufferWithContext<>(wrap(buffer), buffer); + } + + @Override + public void recover( + InputChannelInfo channelInfo, + int oldSubtaskIndex, + BufferWithContext<Buffer> bufferWithContext) + throws IOException, InterruptedException { + Buffer buffer = bufferWithContext.context; + try { + if (buffer.readableBytes() > 0) { + filteringHandler.filterAndRewrite( + channelInfo.getGateIdx(), + oldSubtaskIndex, + channelInfo.getInputChannelIdx(), + buffer.retainBuffer(), + segmentSerializerFor(getMappedChannels(channelInfo).getChannelInfo())); + } + } finally { + buffer.recycleBuffer(); + } + } + + @VisibleForTesting + boolean isPreFilterBufferInUse() { + return preFilterBufferInUse; + } + + @VisibleForTesting + @Nullable + MemorySegment getPreFilterSegmentForTesting() { + return preFilterSegment; + } + + @Override + void closeInternal() throws IOException { + try { + super.closeInternal(); + } finally { + if (preFilterSegment != null) { + preFilterSegment.free(); + preFilterSegment = null; + preFilterBufferInUse = false; + } + } + } +} + +/** + * FLINK-38544 transitional: removed when the spilling backend lands (the factory then selects + * {@link SpillingWithFilteringHandler} instead). + * + * <p>Recovery handler for the case where checkpointing during recovery is enabled and a filtering + * handler is present. Uses a reusable heap-backed pre-filter buffer (isolated from the Network * Buffer Pool), filters recovered buffers through {@link * ChannelStateFilteringHandler#filterAndRewrite}, and delivers the filtered buffers directly into * the input channel via {@code onRecoveredStateBuffer}. diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java index 2bfad5d6133..784aa7c1c0d 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java @@ -79,7 +79,8 @@ public class SequentialChannelStateReaderImpl implements SequentialChannelStateR taskStateSnapshot.getInputRescalingDescriptor(), filterContext.isCheckpointingDuringRecoveryEnabled(), filteringHandler, - filterContext.getMemorySegmentSize())) { + filterContext.getMemorySegmentSize(), + filterContext.getTmpDirectories())) { boolean readAny = read( stateHandler, @@ -98,18 +99,15 @@ public class SequentialChannelStateReaderImpl implements SequentialChannelStateR !filteringHandler.hasPartialData(), "Not all data has been fully consumed during filtering"); } - // A recovered-state container is produced whenever the checkpointing-during-recovery - // path recovered any state, regardless of whether filtering was needed: on this path - // conversion must hand the recovered buffers to the physical channels in recovery mode - // (needsRecovery = state.isPresent()), so the signal must reflect "any recovered data - // was pushed", not "a filter ran". The no-checkpointing path pushes recovered buffers - // directly and produces nothing here, matching the caller's - // checkState(readInputData(...).isEmpty()). + // The container signals "any recovered data was pushed", not "a filter ran": the + // no-checkpointing path pushes recovered buffers directly and must produce nothing + // here, matching the caller's checkState(readInputData(...).isEmpty()). // // FLINK-38544 transitional in-memory backend: recovered buffers already live in the - // physical channels' queues, so the returned container is an empty placeholder that - // only signals "there is state to recover". The spilling backend returns a real, - // file-backed container here. + // physical channels' queues, so the returned container is an empty placeholder. + // TODO: FLINK-38544 — return stateHandler.getProducedChannelState() once the handler + // factory selects the Spilling* handlers; as-is their file-backed state would be + // silently dropped here. return filterContext.isCheckpointingDuringRecoveryEnabled() && readAny ? Optional.of(new FetchedChannelState(Collections.emptyList())) : Optional.empty(); diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandlerTest.java new file mode 100644 index 00000000000..e9581fa3451 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateFilteringHandlerTest.java @@ -0,0 +1,79 @@ +/* + * 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.runtime.checkpoint.channel; + +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor; +import org.apache.flink.runtime.checkpoint.RescaleMappings; +import org.apache.flink.runtime.io.network.partition.consumer.InputGate; +import org.apache.flink.runtime.io.network.partition.consumer.SingleInputGateBuilder; +import org.apache.flink.runtime.memory.MemoryManager; +import org.apache.flink.streaming.runtime.io.recovery.RecordFilterContext; +import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Path; +import java.util.HashSet; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link ChannelStateFilteringHandler}. */ +class ChannelStateFilteringHandlerTest { + + @TempDir Path tempDir; + + @Test + void testCreateFromContextUsesProvidedSpillDirectories() { + InputGate inputGate = new SingleInputGateBuilder().setNumberOfChannels(1).build(); + RecordFilterContext context = createRecordFilterContext(new String[] {tempDir.toString()}); + + ChannelStateFilteringHandler handler = + ChannelStateFilteringHandler.createFromContext( + context, new InputGate[] {inputGate}); + + assertThat(handler).isNotNull(); + handler.close(); + } + + private static RecordFilterContext createRecordFilterContext(String[] tmpDirectories) { + return new RecordFilterContext( + new RecordFilterContext.InputFilterConfig[] { + new RecordFilterContext.InputFilterConfig( + LongSerializer.INSTANCE, new ForwardPartitioner<>(), 1) + }, + new InflightDataRescalingDescriptor( + new InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor[] { + new InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor( + new int[] {0}, + RescaleMappings.identity(1, 1), + new HashSet<>(), + InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor + .MappingType.IDENTITY) + }), + 0, + 128, + tmpDirectories, + true, + MemoryManager.DEFAULT_PAGE_SIZE); + } +} diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerBufferOwnershipTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerBufferOwnershipTest.java index 85b4fd1d48e..ae7b722f683 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerBufferOwnershipTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerBufferOwnershipTest.java @@ -38,16 +38,14 @@ import org.junit.jupiter.api.Test; import java.io.IOException; import java.util.HashMap; -import java.util.List; import java.util.Map; -import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; /** * Tests buffer ownership semantics of {@link ChannelStateFilteringHandler.GateFilterHandler}. Each - * test verifies that buffers are properly recycled on both success and failure paths. + * test verifies that source buffers are properly recycled on both success and failure paths. */ class GateFilterHandlerBufferOwnershipTest { @@ -60,13 +58,9 @@ class GateFilterHandlerBufferOwnershipTest { createHandler(RecordFilter.acceptAll()); Buffer sourceBuffer = createBufferWithRecords(1L, 2L); - List<Buffer> result = handler.filterAndRewrite(0, 0, sourceBuffer, this::createEmptyBuffer); + handler.filterAndRewrite(0, 0, sourceBuffer, new DataOutputSerializer(BUFFER_SIZE)); - // sourceBuffer should be recycled by the deserializer after consumption assertThat(sourceBuffer.isRecycled()).isTrue(); - - // Clean up result buffers - result.forEach(Buffer::recycleBuffer); } @Test @@ -75,102 +69,60 @@ class GateFilterHandlerBufferOwnershipTest { ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(rejectAll); Buffer sourceBuffer = createBufferWithRecords(1L, 2L); - List<Buffer> result = handler.filterAndRewrite(0, 0, sourceBuffer, this::createEmptyBuffer); + handler.filterAndRewrite(0, 0, sourceBuffer, new DataOutputSerializer(BUFFER_SIZE)); - assertThat(result).isEmpty(); - // sourceBuffer should still be recycled even though no output was produced assertThat(sourceBuffer.isRecycled()).isTrue(); } @Test void testSourceBufferRecycledOnInvalidVirtualChannel() { - // Create handler with KEY=(0,0) but call with (1,1) to trigger IllegalStateException + // Create handler with KEY=(0,0) but call with (1,1) to trigger IllegalStateException. ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(RecordFilter.acceptAll()); Buffer sourceBuffer = createBufferWithRecords(1L); assertThatThrownBy( - () -> handler.filterAndRewrite(1, 1, sourceBuffer, this::createEmptyBuffer)) + () -> + handler.filterAndRewrite( + 1, 1, sourceBuffer, new DataOutputSerializer(BUFFER_SIZE))) .isInstanceOf(IllegalStateException.class); - // sourceBuffer must be recycled even when lookup fails before setNextBuffer - assertThat(sourceBuffer.isRecycled()).isTrue(); - } - - @Test - void testResultBuffersAndCurrentBufferRecycledOnSerializationError() throws Exception { - // Use a small buffer so that records span multiple buffers. The supplier fails on the - // second request, after the first output buffer has been filled and added to resultBuffers. - AtomicInteger bufferRequestCount = new AtomicInteger(0); - ChannelStateFilteringHandler.BufferSupplier failingSupplier = - () -> { - if (bufferRequestCount.incrementAndGet() > 1) { - throw new IOException("Simulated buffer allocation failure"); - } - return createEmptyBuffer(13); - }; - - ChannelStateFilteringHandler.GateFilterHandler<Long> handler = - createHandler(RecordFilter.acceptAll()); - - Buffer sourceBuffer = createBufferWithRecords(1L, 2L, 3L, 4L, 5L); - - // The exception should propagate; no buffer leak (no IllegalReferenceCountException - // from double-recycle). - assertThatThrownBy(() -> handler.filterAndRewrite(0, 0, sourceBuffer, failingSupplier)) - .isInstanceOf(IOException.class) - .hasMessage("Simulated buffer allocation failure"); - - // sourceBuffer ownership was transferred to the deserializer via setNextBuffer(). - // The deserializer may still hold it if it hasn't fully consumed the buffer before the - // error. Calling clear() triggers the cleanup chain: - // GateFilterHandler#clear() -> VirtualChannel#clear() -> deserializer.clear() - handler.clear(); + // sourceBuffer must be recycled even when lookup fails before setNextBuffer. assertThat(sourceBuffer.isRecycled()).isTrue(); } /** - * Tests the production cleanup path: when filterAndRewrite throws mid-processing, the - * deserializer may still hold sourceBuffer. In production, ChannelStateFilteringHandler is used - * in a try-with-resources block (see {@code SequentialChannelStateReaderImpl#readInputData}), - * so its close() is guaranteed to be called, which triggers clear() on all GateFilterHandlers - * and their deserializers. This test simulates that exact pattern. + * When filterAndRewrite throws mid-processing, the deserializer may still hold sourceBuffer. In + * production, ChannelStateFilteringHandler is used in a try-with-resources block (see {@code + * SequentialChannelStateReaderImpl#readInputData}), so its close() is guaranteed to be called, + * which triggers clear() on all GateFilterHandlers and their deserializers. This test simulates + * that exact pattern. */ @Test void testCloseRecyclesDeserializerHeldBufferAfterError() throws Exception { - AtomicInteger bufferRequestCount = new AtomicInteger(0); - ChannelStateFilteringHandler.BufferSupplier failingSupplier = - () -> { - if (bufferRequestCount.incrementAndGet() > 1) { - throw new IOException("Simulated buffer allocation failure"); - } - return createEmptyBuffer(13); - }; - ChannelStateFilteringHandler.GateFilterHandler<Long> gateHandler = createHandler(RecordFilter.acceptAll()); - // Wrap in ChannelStateFilteringHandler, the production-level owner ChannelStateFilteringHandler filteringHandler = new ChannelStateFilteringHandler( new ChannelStateFilteringHandler.GateFilterHandler<?>[] {gateHandler}); + // A serializer that throws while writing the second record's length prefix, triggering a + // mid-processing failure after the first record has already been emitted. + DataOutputSerializer failingSerializer = new FailingAfterFirstRecordSerializer(); Buffer sourceBuffer = createBufferWithRecords(1L, 2L, 3L, 4L, 5L); - // Simulate the production try-with-resources pattern assertThatThrownBy( () -> { try (ChannelStateFilteringHandler ignored = filteringHandler) { filteringHandler.filterAndRewrite( - 0, 0, 0, sourceBuffer, failingSupplier); + 0, 0, 0, sourceBuffer, failingSerializer); } }) .isInstanceOf(IOException.class) - .hasMessage("Simulated buffer allocation failure"); + .hasMessage("Simulated write failure"); - // After close(), the entire cleanup chain has fired: - // ChannelStateFilteringHandler.close() -> GateFilterHandler.clear() - // -> VirtualChannel.clear() -> deserializer.clear() -> sourceBuffer.recycleBuffer() + // After close(), the entire cleanup chain has fired. assertThat(sourceBuffer.isRecycled()).isTrue(); } @@ -219,12 +171,26 @@ class GateFilterHandlerBufferOwnershipTest { } } - private Buffer createEmptyBuffer() { - return createEmptyBuffer(BUFFER_SIZE); - } + /** + * A {@link DataOutputSerializer} that throws an IOException while writing the second record's + * length prefix, simulating a failure mid-stream to verify that the source buffer is still + * recycled via the filtering handler's close() cleanup chain. Each surviving record begins with + * a {@code writeInt} placeholder for its length, so the second {@code writeInt} marks the start + * of the second record. + */ + private static final class FailingAfterFirstRecordSerializer extends DataOutputSerializer { + private int writeIntCount = 0; + + FailingAfterFirstRecordSerializer() { + super(BUFFER_SIZE); + } - private Buffer createEmptyBuffer(int size) { - MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(size); - return new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE); + @Override + public void writeInt(int v) throws IOException { + if (++writeIntCount > 1) { + throw new IOException("Simulated write failure"); + } + super.writeInt(v); + } } } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerTest.java index f02ce35fd86..1646908727a 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/GateFilterHandlerTest.java @@ -57,10 +57,10 @@ class GateFilterHandlerTest { createHandler(RecordFilter.acceptAll()); Buffer sourceBuffer = createBufferWithRecords(1L, 2L, 3L); - List<Buffer> result = handler.filterAndRewrite(0, 0, sourceBuffer, this::createEmptyBuffer); + DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); + handler.filterAndRewrite(0, 0, sourceBuffer, output); - // deserializeBuffers consumes (recycles) each buffer via the deserializer - List<Long> values = deserializeBuffers(result); + List<Long> values = readRecordsFromSerializer(output); assertThat(values).containsExactly(1L, 2L, 3L); } @@ -70,9 +70,11 @@ class GateFilterHandlerTest { ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(rejectAll); Buffer sourceBuffer = createBufferWithRecords(1L, 2L, 3L); - List<Buffer> result = handler.filterAndRewrite(0, 0, sourceBuffer, this::createEmptyBuffer); + DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); + handler.filterAndRewrite(0, 0, sourceBuffer, output); - assertThat(result).isEmpty(); + // No bytes should be written when all records are filtered out. + assertThat(output.length()).isZero(); } @Test @@ -81,42 +83,50 @@ class GateFilterHandlerTest { ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(keepEven); Buffer sourceBuffer = createBufferWithRecords(1L, 2L, 3L, 4L, 5L); - List<Buffer> result = handler.filterAndRewrite(0, 0, sourceBuffer, this::createEmptyBuffer); + DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); + handler.filterAndRewrite(0, 0, sourceBuffer, output); - List<Long> values = deserializeBuffers(result); + List<Long> values = readRecordsFromSerializer(output); assertThat(values).containsExactly(2L, 4L); } @Test - void testSmallOutputBufferProducesMultipleBuffers() throws Exception { - // Use a very small output buffer size so records must span multiple buffers - int smallBufferSize = 8; + void testEmptyBuffer() throws Exception { ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(RecordFilter.acceptAll()); - Buffer sourceBuffer = createBufferWithRecords(1L, 2L, 3L); - List<Buffer> result = - handler.filterAndRewrite( - 0, 0, sourceBuffer, () -> createEmptyBuffer(smallBufferSize)); + Buffer emptyBuffer = createEmptyBuffer(); + emptyBuffer.setSize(0); - // Each Long record needs 4 bytes length + ~9 bytes data > 8-byte buffer - assertThat(result.size()).isGreaterThan(1); + DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); + handler.filterAndRewrite(0, 0, emptyBuffer, output); - List<Long> values = deserializeBuffers(result); - assertThat(values).containsExactly(1L, 2L, 3L); + // No data written for an empty source buffer. + assertThat(output.length()).isZero(); } @Test - void testEmptyBuffer() throws Exception { + void testSourceBufferRecycledOnSuccess() throws Exception { ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(RecordFilter.acceptAll()); - Buffer emptyBuffer = createEmptyBuffer(); - emptyBuffer.setSize(0); + Buffer sourceBuffer = createBufferWithRecords(1L, 2L); + DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); + handler.filterAndRewrite(0, 0, sourceBuffer, output); + + assertThat(sourceBuffer.isRecycled()).isTrue(); + } + + @Test + void testSourceBufferRecycledWhenAllRecordsFilteredOut() throws Exception { + RecordFilter<Long> rejectAll = record -> false; + ChannelStateFilteringHandler.GateFilterHandler<Long> handler = createHandler(rejectAll); - List<Buffer> result = handler.filterAndRewrite(0, 0, emptyBuffer, this::createEmptyBuffer); + Buffer sourceBuffer = createBufferWithRecords(1L, 2L); + DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); + handler.filterAndRewrite(0, 0, sourceBuffer, output); - assertThat(result).isEmpty(); + assertThat(sourceBuffer.isRecycled()).isTrue(); } // ------------------------------------------------------------------------------------------- @@ -141,23 +151,13 @@ class GateFilterHandlerTest { private Buffer createBufferWithRecords(Long... values) throws IOException { StreamElementSerializer<Long> serializer = new StreamElementSerializer<>(LongSerializer.INSTANCE); - return serializeRecordsToBuffer(serializer, values); - } - - /** Serializes records into a buffer using Flink's length-prefixed format. */ - private Buffer serializeRecordsToBuffer( - StreamElementSerializer<Long> serializer, Long... values) throws IOException { DataOutputSerializer output = new DataOutputSerializer(BUFFER_SIZE); for (Long value : values) { - // Serialize using the same length-prefixed format as Flink DataOutputSerializer recordOutput = new DataOutputSerializer(64); serializer.serialize(new StreamRecord<>(value), recordOutput); int recordLength = recordOutput.length(); - - // Write 4-byte big-endian length prefix output.writeInt(recordLength); - // Write record bytes output.write(recordOutput.getSharedBuffer(), 0, recordLength); } @@ -171,43 +171,49 @@ class GateFilterHandlerTest { } private Buffer createEmptyBuffer() { - return createEmptyBuffer(BUFFER_SIZE); - } - - private Buffer createEmptyBuffer(int size) { - MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(size); + MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE); return new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE); } - private List<Long> deserializeBuffers(List<Buffer> buffers) throws IOException { + /** + * Deserializes the records the handler appended into {@code output}. The body format is + * repeated (4B recordLen + N bytes of serialized StreamElement), which the deserializer reads + * directly. + */ + private List<Long> readRecordsFromSerializer(DataOutputSerializer output) throws Exception { + List<Long> values = new ArrayList<>(); StreamElementSerializer<Long> serializer = new StreamElementSerializer<>(LongSerializer.INSTANCE); + DeserializationDelegate<StreamElement> delegate = + new NonReusingDeserializationDelegate<>(serializer); + + byte[] bodyBytes = output.getCopyOfBuffer(); + if (bodyBytes.length == 0) { + return values; + } + MemorySegment memSeg = MemorySegmentFactory.allocateUnpooledSegment(bodyBytes.length); + memSeg.put(0, bodyBytes); + NetworkBuffer buf = new NetworkBuffer(memSeg, FreeingBufferRecycler.INSTANCE); + buf.setSize(bodyBytes.length); + SpillingAdaptiveSpanningRecordDeserializer<DeserializationDelegate<StreamElement>> deserializer = new SpillingAdaptiveSpanningRecordDeserializer<>( new String[] {System.getProperty("java.io.tmpdir")}); - DeserializationDelegate<StreamElement> delegate = - new NonReusingDeserializationDelegate<>(serializer); - - List<Long> values = new ArrayList<>(); - for (Buffer buffer : buffers) { - deserializer.setNextBuffer(buffer); - while (true) { - RecordDeserializer.DeserializationResult result = - deserializer.getNextRecord(delegate); - if (result.isFullRecord()) { - StreamElement element = delegate.getInstance(); - if (element.isRecord()) { - @SuppressWarnings("unchecked") - StreamRecord<Long> record = (StreamRecord<Long>) element; - values.add(record.getValue()); - } - } - if (result.isBufferConsumed()) { - break; + deserializer.setNextBuffer(buf); + + RecordDeserializer.DeserializationResult result; + do { + result = deserializer.getNextRecord(delegate); + if (result.isFullRecord()) { + StreamElement element = delegate.getInstance(); + if (element.isRecord()) { + @SuppressWarnings("unchecked") + StreamRecord<Long> record = (StreamRecord<Long>) element; + values.add(record.getValue()); } } - } + } while (!result.isBufferConsumed()); return values; } } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java index 76f82a82ff7..e398e41562a 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java @@ -33,7 +33,9 @@ import org.apache.flink.runtime.memory.MemoryManager; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import java.nio.file.Path; import java.util.Arrays; import java.util.HashSet; import java.util.List; @@ -46,6 +48,8 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test of different implementation of {@link AbstractInputChannelRecoveredStateHandler}. */ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandlerTest { + @TempDir private Path tmpDir; + private static final int preAllocatedSegments = 3; private NetworkBufferPool networkBufferPool; private SingleInputGate inputGate; @@ -89,7 +93,8 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler }), false, null, - MemoryManager.DEFAULT_PAGE_SIZE); + MemoryManager.DEFAULT_PAGE_SIZE, + null); } private AbstractInputChannelRecoveredStateHandler buildMultiChannelHandler() { @@ -118,19 +123,43 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler }), false, null, - MemoryManager.DEFAULT_PAGE_SIZE); + MemoryManager.DEFAULT_PAGE_SIZE, + null); } /** Builds a handler in filtering mode (non-null filtering handler, no-op stub). */ - private FilteringHandler buildFilteringInputChannelStateHandler() { + private SpillingWithFilteringHandler buildFilteringInputChannelStateHandler() { // Empty GateFilterHandler array: filtering is "enabled" structurally, but no gate-level // filter logic runs. Suitable for exercising getBuffer() routing only. - return buildFilteringInputChannelStateHandler( - inputGate, + ChannelStateFilteringHandler stubFilteringHandler = new ChannelStateFilteringHandler( - new ChannelStateFilteringHandler.GateFilterHandler[0])); + new ChannelStateFilteringHandler.GateFilterHandler[0]); + // FLINK-38544 transitional: constructed directly because the factory still routes the + // flag-on filtering case to the in-memory FilteringHandler; goes back through + // AbstractInputChannelRecoveredStateHandler.create(...) when the spilling backend lands. + return new SpillingWithFilteringHandler( + new InputGate[] {inputGate}, + new InflightDataRescalingDescriptor( + new InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor[] { + new InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor( + new int[] {1}, + RescaleMappings.identity(1, 1), + new HashSet<>(), + InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor + .MappingType.IDENTITY) + }), + stubFilteringHandler, + MemoryManager.DEFAULT_PAGE_SIZE, + new String[] {tmpDir.toAbsolutePath().toString()}); } + /** + * Builds the in-memory filtering handler through the factory: the flag-on filtering case still + * routes there until the spilling backend lands. + */ private FilteringHandler buildFilteringInputChannelStateHandler( SingleInputGate inputGate, ChannelStateFilteringHandler stubFilteringHandler) { return (FilteringHandler) @@ -150,7 +179,30 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler }), true, stubFilteringHandler, - MemoryManager.DEFAULT_PAGE_SIZE); + MemoryManager.DEFAULT_PAGE_SIZE, + null); + } + + private AbstractInputChannelRecoveredStateHandler buildSpillingNoFilteringHandler( + String[] spillTmpDirectories) { + // FLINK-38544 transitional: constructed directly because the factory still routes the + // flag-on no-filtering case to the in-memory NoSpillingHandler; goes back through + // AbstractInputChannelRecoveredStateHandler.create(...) when the spilling backend lands. + return new SpillingNoFilteringHandler( + new InputGate[] {inputGate}, + new InflightDataRescalingDescriptor( + new InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor[] { + new InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor( + new int[] {1}, + RescaleMappings.identity(1, 1), + new HashSet<>(), + InflightDataRescalingDescriptor + .InflightDataGateOrPartitionRescalingDescriptor + .MappingType.IDENTITY) + }), + spillTmpDirectories); } @Test @@ -201,7 +253,8 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler @Test void testPreFilterBufferIsolationFromNetworkBufferPool() throws Exception { - try (FilteringHandler filteringHandler = buildFilteringInputChannelStateHandler()) { + try (SpillingWithFilteringHandler filteringHandler = + buildFilteringInputChannelStateHandler()) { int availableBefore = networkBufferPool.getNumberOfAvailableMemorySegments(); RecoveredChannelStateHandler.BufferWithContext<Buffer> bufferWithContext = @@ -242,7 +295,8 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler @Test void testPreFilterSegmentReusedAcrossCalls() throws Exception { - try (FilteringHandler filteringHandler = buildFilteringInputChannelStateHandler()) { + try (SpillingWithFilteringHandler filteringHandler = + buildFilteringInputChannelStateHandler()) { // First getBuffer() lazily allocates the segment. RecoveredChannelStateHandler.BufferWithContext<Buffer> first = filteringHandler.getBuffer(channelInfo); @@ -271,7 +325,8 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler @Test void testGetBufferThrowsWhenPriorBufferNotRecycled() throws Exception { - try (FilteringHandler filteringHandler = buildFilteringInputChannelStateHandler()) { + try (SpillingWithFilteringHandler filteringHandler = + buildFilteringInputChannelStateHandler()) { RecoveredChannelStateHandler.BufferWithContext<Buffer> first = filteringHandler.getBuffer(channelInfo); try { @@ -337,7 +392,7 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler @Test void testPreFilterSegmentFreedOnClose() throws Exception { - FilteringHandler filteringHandler = buildFilteringInputChannelStateHandler(); + SpillingWithFilteringHandler filteringHandler = buildFilteringInputChannelStateHandler(); RecoveredChannelStateHandler.BufferWithContext<Buffer> bufferWithContext = filteringHandler.getBuffer(channelInfo); bufferWithContext.context.recycleBuffer(); @@ -351,4 +406,13 @@ class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandler assertThat(segment.isFreed()).isTrue(); assertThat(filteringHandler.getPreFilterSegmentForTesting()).isNull(); } + + @Test + void testSpillingHandlerRequiresSpillDirectories() { + assertThatThrownBy(() -> buildSpillingNoFilteringHandler(null)) + .isInstanceOf(NullPointerException.class); + assertThatThrownBy(() -> buildSpillingNoFilteringHandler(new String[0])) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("spillTmpDirectories must not be empty"); + } }
