You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by gu...@apache.org on 2023/07/04 11:56:28 UTC
[flink] 01/03: [FLINK-31645][network] Introduce the ProducerMergedPartitionFileReader
This is an automated email from the ASF dual-hosted git repository.
guoweijie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 25eb0cef638ff1562669e42be151dad082db0518
Author: Wencong Liu <li...@163.com>
AuthorDate: Tue Jul 4 10:16:59 2023 +0800
[FLINK-31645][network] Introduce the ProducerMergedPartitionFileReader
---
.../file/ProducerMergedPartitionFileReader.java | 209 ++++++++++++++++++++-
.../tiered/common/TieredStorageTestUtils.java | 54 ++++++
.../ProducerMergedPartitionFileReaderTest.java | 179 ++++++++++++++++++
.../ProducerMergedPartitionFileWriterTest.java | 28 +--
.../tiered/file/TestingPartitionFileReader.java | 105 +++++++++++
5 files changed, 542 insertions(+), 33 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReader.java
index ce17f1c4933..d82aa78659f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReader.java
@@ -18,14 +18,30 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.file;
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferRecycler;
+import org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.IOUtils;
import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil.positionToNextBuffer;
+import static org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil.readFromByteChannel;
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.ProducerMergedPartitionFileIndex.Region;
+import static org.apache.flink.util.Preconditions.checkNotNull;
/**
* The implementation of {@link PartitionFileReader} with producer-merge mode. In this mode, the
@@ -36,9 +52,51 @@ import java.nio.file.Path;
*/
public class ProducerMergedPartitionFileReader implements PartitionFileReader {
+ /**
+ * Max number of caches.
+ *
+ * <p>The constant defines the maximum number of caches that can be created. Its value is set to
+ * 10000, which is considered sufficient for most parallel jobs. Each cache only contains
+ * references and numerical variables and occupies a minimal amount of memory so the value is
+ * not excessively large.
+ */
+ private static final int DEFAULT_MAX_CACHE_NUM = 10000;
+
+ /**
+ * Buffer offset caches stored in map.
+ *
+ * <p>The key is the combination of {@link TieredStorageSubpartitionId} and buffer index. The
+ * value is the buffer offset cache, which includes file offset of the buffer index, the region
+ * containing the buffer index and next buffer index to consume.
+ */
+ private final Map<Tuple2<TieredStorageSubpartitionId, Integer>, BufferOffsetCache>
+ bufferOffsetCaches;
+
+ private final ByteBuffer reusedHeaderBuffer = BufferReaderWriterUtil.allocatedHeaderBuffer();
+
+ private final Path dataFilePath;
+
+ private final ProducerMergedPartitionFileIndex dataIndex;
+
+ private final int maxCacheNumber;
+
+ private volatile FileChannel fileChannel;
+
+ /** The current number of caches. */
+ private int numCaches;
+
ProducerMergedPartitionFileReader(
- Path dataFilePath, ProducerMergedPartitionFileIndex partitionFileIndex) {
- // TODO, implement the ProducerMergePartitionFileReader
+ Path dataFilePath, ProducerMergedPartitionFileIndex dataIndex) {
+ this(dataFilePath, dataIndex, DEFAULT_MAX_CACHE_NUM);
+ }
+
+ @VisibleForTesting
+ ProducerMergedPartitionFileReader(
+ Path dataFilePath, ProducerMergedPartitionFileIndex dataIndex, int maxCacheNumber) {
+ this.dataFilePath = dataFilePath;
+ this.dataIndex = dataIndex;
+ this.bufferOffsetCaches = new HashMap<>();
+ this.maxCacheNumber = maxCacheNumber;
}
@Override
@@ -50,8 +108,32 @@ public class ProducerMergedPartitionFileReader implements PartitionFileReader {
MemorySegment memorySegment,
BufferRecycler recycler)
throws IOException {
- // TODO, implement the ProducerMergePartitionFileReader
- return null;
+
+ lazyInitializeFileChannel();
+ Tuple2<TieredStorageSubpartitionId, Integer> cacheKey =
+ Tuple2.of(subpartitionId, bufferIndex);
+ Optional<BufferOffsetCache> cache = tryGetCache(cacheKey, true);
+ if (!cache.isPresent()) {
+ return null;
+ }
+ fileChannel.position(cache.get().getFileOffset());
+ Buffer buffer =
+ readFromByteChannel(fileChannel, reusedHeaderBuffer, memorySegment, recycler);
+ boolean hasNextBuffer =
+ cache.get()
+ .advance(
+ checkNotNull(buffer).readableBytes()
+ + BufferReaderWriterUtil.HEADER_LENGTH);
+ if (hasNextBuffer) {
+ int nextBufferIndex = bufferIndex + 1;
+ // TODO: introduce the LRU cache strategy in the future to restrict the total
+ // cache number. Testing to prevent cache leaks has been implemented.
+ if (numCaches < maxCacheNumber) {
+ bufferOffsetCaches.put(Tuple2.of(subpartitionId, nextBufferIndex), cache.get());
+ numCaches++;
+ }
+ }
+ return buffer;
}
@Override
@@ -60,9 +142,124 @@ public class ProducerMergedPartitionFileReader implements PartitionFileReader {
TieredStorageSubpartitionId subpartitionId,
int segmentId,
int bufferIndex) {
- return 0;
+ lazyInitializeFileChannel();
+ Tuple2<TieredStorageSubpartitionId, Integer> cacheKey =
+ Tuple2.of(subpartitionId, bufferIndex);
+ return tryGetCache(cacheKey, false)
+ .map(BufferOffsetCache::getFileOffset)
+ .orElse(Long.MAX_VALUE);
}
@Override
- public void release() {}
+ public void release() {
+ if (fileChannel != null) {
+ try {
+ fileChannel.close();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e, "Failed to close file channel.");
+ }
+ }
+ IOUtils.deleteFileQuietly(dataFilePath);
+ }
+
+ /**
+ * Initialize the file channel in a lazy manner, which can reduce usage of the file descriptor
+ * resource.
+ */
+ private void lazyInitializeFileChannel() {
+ if (fileChannel == null) {
+ try {
+ fileChannel = FileChannel.open(dataFilePath, StandardOpenOption.READ);
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e, "Failed to open file channel.");
+ }
+ }
+ }
+
+ /**
+ * Try to get the cache according to the key.
+ *
+ * <p>If the relevant buffer offset cache exists, it will be returned and subsequently removed.
+ * However, if the buffer offset cache does not exist, a new cache will be created using the
+ * data index and returned.
+ *
+ * @param cacheKey the key of cache.
+ * @param removeKey boolean decides whether to remove key.
+ * @return returns the relevant buffer offset cache if it exists, otherwise return {@link
+ * Optional#empty()}.
+ */
+ private Optional<BufferOffsetCache> tryGetCache(
+ Tuple2<TieredStorageSubpartitionId, Integer> cacheKey, boolean removeKey) {
+ BufferOffsetCache bufferOffsetCache = bufferOffsetCaches.remove(cacheKey);
+ if (bufferOffsetCache == null) {
+ Optional<Region> regionOpt = dataIndex.getRegion(cacheKey.f0, cacheKey.f1);
+ return regionOpt.map(region -> new BufferOffsetCache(cacheKey.f1, region));
+ } else {
+ if (removeKey) {
+ numCaches--;
+ } else {
+ bufferOffsetCaches.put(cacheKey, bufferOffsetCache);
+ }
+ return Optional.of(bufferOffsetCache);
+ }
+ }
+
+ /**
+ * The {@link BufferOffsetCache} represents the file offset cache for a buffer index. Each cache
+ * includes file offset of the buffer index, the region containing the buffer index and next
+ * buffer index to consume.
+ */
+ private class BufferOffsetCache {
+
+ private final Region region;
+
+ private long fileOffset;
+
+ private int nextBufferIndex;
+
+ private BufferOffsetCache(int bufferIndex, Region region) {
+ this.nextBufferIndex = bufferIndex;
+ this.region = region;
+ moveFileOffsetToBuffer(bufferIndex);
+ }
+
+ /**
+ * Get the file offset.
+ *
+ * @return the file offset.
+ */
+ private long getFileOffset() {
+ return fileOffset;
+ }
+
+ /**
+ * Updates the {@link BufferOffsetCache} upon the retrieval of a buffer from the file using
+ * the file offset in the {@link BufferOffsetCache}.
+ *
+ * @param bufferSize denotes the size of the buffer.
+ * @return return true if there are remaining buffers in the region, otherwise return false.
+ */
+ private boolean advance(long bufferSize) {
+ nextBufferIndex++;
+ fileOffset += bufferSize;
+ return nextBufferIndex < (region.getFirstBufferIndex() + region.getNumBuffers());
+ }
+
+ /**
+ * Relocates the file channel offset to the position of the specified buffer index.
+ *
+ * @param bufferIndex denotes the index of the buffer.
+ */
+ private void moveFileOffsetToBuffer(int bufferIndex) {
+ try {
+ checkNotNull(fileChannel).position(region.getRegionFileOffset());
+ for (int i = 0; i < (bufferIndex - region.getFirstBufferIndex()); ++i) {
+ positionToNextBuffer(fileChannel, reusedHeaderBuffer);
+ }
+ fileOffset = fileChannel.position();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e, "Failed to move file offset");
+ }
+ }
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/common/TieredStorageTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/common/TieredStorageTestUtils.java
new file mode 100644
index 00000000000..11650b359a9
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/common/TieredStorageTestUtils.java
@@ -0,0 +1,54 @@
+/*
+ * 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.io.network.partition.hybrid.tiered.common;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilderTestUtils;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/** Test utils for tiered storage. */
+public class TieredStorageTestUtils {
+
+ public static List<PartitionFileWriter.SubpartitionBufferContext> generateBuffersToWrite(
+ int numSubpartitions, int numSegments, int numBuffersPerSegment, int bufferSizeBytes) {
+ List<PartitionFileWriter.SubpartitionBufferContext> subpartitionBuffers = new ArrayList<>();
+ for (int i = 0; i < numSubpartitions; i++) {
+ List<PartitionFileWriter.SegmentBufferContext> segmentBufferContexts =
+ new ArrayList<>();
+ for (int j = 0; j < numSegments; j++) {
+ List<Tuple2<Buffer, Integer>> bufferAndIndexes = new ArrayList<>();
+ for (int k = 0; k < numBuffersPerSegment; k++) {
+ bufferAndIndexes.add(
+ new Tuple2<>(
+ BufferBuilderTestUtils.buildSomeBuffer(bufferSizeBytes), k));
+ }
+ segmentBufferContexts.add(
+ new PartitionFileWriter.SegmentBufferContext(j, bufferAndIndexes, false));
+ }
+
+ subpartitionBuffers.add(
+ new PartitionFileWriter.SubpartitionBufferContext(i, segmentBufferContexts));
+ }
+ return subpartitionBuffers;
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReaderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReaderTest.java
new file mode 100644
index 00000000000..7d6f3b25da0
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileReaderTest.java
@@ -0,0 +1,179 @@
+/*
+ * 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.io.network.partition.hybrid.tiered.file;
+
+import org.apache.flink.core.memory.MemorySegment;
+import org.apache.flink.core.memory.MemorySegmentFactory;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageIdMappingUtils;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil.HEADER_LENGTH;
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageTestUtils.generateBuffersToWrite;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link ProducerMergedPartitionFileReader}. */
+class ProducerMergedPartitionFileReaderTest {
+
+ private static final int DEFAULT_NUM_SUBPARTITION = 1;
+
+ private static final int DEFAULT_SEGMENT_NUM = 1;
+
+ private static final int DEFAULT_SEGMENT_ID = 0;
+
+ private static final int DEFAULT_BUFFER_NUMBER = 5;
+
+ private static final int DEFAULT_BUFFER_SIZE = 3;
+
+ private static final String DEFAULT_TEST_FILE_NAME = "testFile";
+
+ private static final TieredStoragePartitionId DEFAULT_PARTITION_ID =
+ TieredStorageIdMappingUtils.convertId(new ResultPartitionID());
+
+ private static final TieredStorageSubpartitionId DEFAULT_SUBPARTITION_ID =
+ new TieredStorageSubpartitionId(0);
+
+ @TempDir private Path tempFolder;
+
+ private Path testFilePath;
+
+ private ProducerMergedPartitionFileReader partitionFileReader;
+
+ @BeforeEach
+ void before() throws ExecutionException, InterruptedException {
+ ProducerMergedPartitionFileIndex partitionFileIndex =
+ new ProducerMergedPartitionFileIndex(DEFAULT_NUM_SUBPARTITION);
+ testFilePath = new File(tempFolder.toFile(), DEFAULT_TEST_FILE_NAME).toPath();
+ ProducerMergedPartitionFileWriter partitionFileWriter =
+ new ProducerMergedPartitionFileWriter(testFilePath, partitionFileIndex);
+ // Write buffers to disk by writer
+ List<PartitionFileWriter.SubpartitionBufferContext> subpartitionBuffers =
+ generateBuffersToWrite(
+ DEFAULT_NUM_SUBPARTITION,
+ DEFAULT_SEGMENT_NUM,
+ DEFAULT_BUFFER_NUMBER,
+ DEFAULT_BUFFER_SIZE);
+ partitionFileWriter.write(DEFAULT_PARTITION_ID, subpartitionBuffers).get();
+ partitionFileReader =
+ new ProducerMergedPartitionFileReader(testFilePath, partitionFileIndex);
+ }
+
+ @Test
+ void testReadBuffer() throws IOException {
+ for (int bufferIndex = 0; bufferIndex < DEFAULT_BUFFER_NUMBER; ++bufferIndex) {
+ Buffer buffer = readBuffer(bufferIndex, DEFAULT_SUBPARTITION_ID);
+ assertThat(buffer).isNotNull();
+ buffer.recycleBuffer();
+ }
+ MemorySegment memorySegment =
+ MemorySegmentFactory.allocateUnpooledSegment(DEFAULT_BUFFER_SIZE);
+ assertThat(
+ partitionFileReader.readBuffer(
+ DEFAULT_PARTITION_ID,
+ DEFAULT_SUBPARTITION_ID,
+ DEFAULT_SEGMENT_ID,
+ DEFAULT_BUFFER_NUMBER + 1,
+ memorySegment,
+ FreeingBufferRecycler.INSTANCE))
+ .isNull();
+ }
+
+ @Test
+ void testGetPriority() throws IOException {
+ int currentFileOffset = 0;
+ for (int bufferIndex = 0; bufferIndex < DEFAULT_BUFFER_NUMBER; ++bufferIndex) {
+ Buffer buffer = readBuffer(bufferIndex, DEFAULT_SUBPARTITION_ID);
+ assertThat(buffer).isNotNull();
+ assertThat(
+ partitionFileReader.getPriority(
+ DEFAULT_PARTITION_ID,
+ DEFAULT_SUBPARTITION_ID,
+ DEFAULT_SEGMENT_ID,
+ bufferIndex))
+ .isEqualTo(currentFileOffset);
+ currentFileOffset += (HEADER_LENGTH + DEFAULT_BUFFER_SIZE);
+ buffer.recycleBuffer();
+ }
+ }
+
+ @Test
+ void testCacheExceedMaxNumber() throws IOException {
+ int cacheNumber = 3;
+ AtomicInteger indexQueryTime = new AtomicInteger(0);
+ TestingProducerMergedPartitionFileIndex partitionFileIndex =
+ new TestingProducerMergedPartitionFileIndex.Builder()
+ .setGetRegionFunction(
+ (subpartitionId, integer) -> {
+ indexQueryTime.incrementAndGet();
+ return Optional.of(
+ new ProducerMergedPartitionFileIndex.Region(0, 0, 2));
+ })
+ .build();
+ partitionFileReader =
+ new ProducerMergedPartitionFileReader(
+ testFilePath, partitionFileIndex, cacheNumber);
+ // Read different subpartitions from the reader and make cache reach max number.
+ for (int subpartitionId = 0; subpartitionId < cacheNumber * 2; ++subpartitionId) {
+ assertThat(
+ readBuffer(
+ subpartitionId < cacheNumber ? 0 : 1,
+ new TieredStorageSubpartitionId(subpartitionId % cacheNumber)))
+ .isNotNull();
+ }
+ // The following buffer reading from other subpartitions can only query the index.
+ assertThat(readBuffer(0, new TieredStorageSubpartitionId(3))).isNotNull();
+ assertThat(readBuffer(0, new TieredStorageSubpartitionId(3))).isNotNull();
+ assertThat(indexQueryTime).hasValue(5);
+ }
+
+ @Test
+ void testRelease() {
+ assertThat(testFilePath.toFile().exists()).isTrue();
+ partitionFileReader.release();
+ assertThat(testFilePath.toFile().exists()).isFalse();
+ }
+
+ private Buffer readBuffer(int bufferIndex, TieredStorageSubpartitionId subpartitionId)
+ throws IOException {
+ MemorySegment memorySegment =
+ MemorySegmentFactory.allocateUnpooledSegment(DEFAULT_BUFFER_SIZE);
+ return partitionFileReader.readBuffer(
+ DEFAULT_PARTITION_ID,
+ subpartitionId,
+ DEFAULT_SEGMENT_ID,
+ bufferIndex,
+ memorySegment,
+ FreeingBufferRecycler.INSTANCE);
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileWriterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileWriterTest.java
index 65ad5980d30..ddde5377567 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileWriterTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/ProducerMergedPartitionFileWriterTest.java
@@ -18,9 +18,6 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.file;
-import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.runtime.io.network.buffer.Buffer;
-import org.apache.flink.runtime.io.network.buffer.BufferBuilderTestUtils;
import org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageIdMappingUtils;
@@ -32,11 +29,11 @@ import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
-import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageTestUtils.generateBuffersToWrite;
import static org.assertj.core.api.Assertions.assertThat;
/** Tests for {@link ProducerMergedPartitionFileWriter}. */
@@ -94,27 +91,4 @@ class ProducerMergedPartitionFileWriterTest {
partitionFileWriter.release();
assertThat(isReleased).isTrue();
}
-
- private static List<PartitionFileWriter.SubpartitionBufferContext> generateBuffersToWrite(
- int numSubpartitions, int numSegments, int numBuffersPerSegment, int bufferSizeBytes) {
- List<PartitionFileWriter.SubpartitionBufferContext> subpartitionBuffers = new ArrayList<>();
- for (int i = 0; i < numSubpartitions; i++) {
- List<PartitionFileWriter.SegmentBufferContext> segmentBufferContexts =
- new ArrayList<>();
- for (int j = 0; j < numSegments; j++) {
- List<Tuple2<Buffer, Integer>> bufferAndIndexes = new ArrayList<>();
- for (int k = 0; k < numBuffersPerSegment; k++) {
- bufferAndIndexes.add(
- new Tuple2<>(
- BufferBuilderTestUtils.buildSomeBuffer(bufferSizeBytes), k));
- }
- segmentBufferContexts.add(
- new PartitionFileWriter.SegmentBufferContext(j, bufferAndIndexes, false));
- }
-
- subpartitionBuffers.add(
- new PartitionFileWriter.SubpartitionBufferContext(i, segmentBufferContexts));
- }
- return subpartitionBuffers;
- }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/TestingPartitionFileReader.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/TestingPartitionFileReader.java
new file mode 100644
index 00000000000..572f3b97d3b
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/TestingPartitionFileReader.java
@@ -0,0 +1,105 @@
+/*
+ * 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.io.network.partition.hybrid.tiered.file;
+
+import org.apache.flink.core.memory.MemorySegment;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferRecycler;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import java.io.IOException;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+
+/** Testing implementation for {@link PartitionFileReader}. */
+public class TestingPartitionFileReader implements PartitionFileReader {
+
+ private final BiFunction<Integer, Integer, Buffer> readBufferFunction;
+
+ private final Function<Integer, Long> getPriorityFunction;
+
+ private final Runnable releaseRunnable;
+
+ private TestingPartitionFileReader(
+ BiFunction<Integer, Integer, Buffer> readBufferFunction,
+ Function<Integer, Long> getPriorityFunction,
+ Runnable releaseRunnable) {
+ this.readBufferFunction = readBufferFunction;
+ this.getPriorityFunction = getPriorityFunction;
+ this.releaseRunnable = releaseRunnable;
+ }
+
+ @Override
+ public Buffer readBuffer(
+ TieredStoragePartitionId partitionId,
+ TieredStorageSubpartitionId subpartitionId,
+ int segmentId,
+ int bufferIndex,
+ MemorySegment memorySegment,
+ BufferRecycler recycler)
+ throws IOException {
+ return readBufferFunction.apply(bufferIndex, segmentId);
+ }
+
+ @Override
+ public long getPriority(
+ TieredStoragePartitionId partitionId,
+ TieredStorageSubpartitionId subpartitionId,
+ int segmentId,
+ int bufferIndex) {
+ return getPriorityFunction.apply(subpartitionId.getSubpartitionId());
+ }
+
+ @Override
+ public void release() {
+ releaseRunnable.run();
+ }
+
+ /** Builder for {@link TestingPartitionFileReader}. */
+ public static class Builder {
+ private BiFunction<Integer, Integer, Buffer> readBufferSupplier =
+ (bufferIndex, segmentId) -> null;
+
+ private Function<Integer, Long> prioritySupplier = bufferIndex -> 0L;
+
+ private Runnable releaseNotifier = () -> {};
+
+ public Builder setReadBufferSupplier(
+ BiFunction<Integer, Integer, Buffer> readBufferSupplier) {
+ this.readBufferSupplier = readBufferSupplier;
+ return this;
+ }
+
+ public Builder setPrioritySupplier(Function<Integer, Long> prioritySupplier) {
+ this.prioritySupplier = prioritySupplier;
+ return this;
+ }
+
+ public Builder setReleaseNotifier(Runnable releaseNotifier) {
+ this.releaseNotifier = releaseNotifier;
+ return this;
+ }
+
+ public TestingPartitionFileReader build() {
+ return new TestingPartitionFileReader(
+ readBufferSupplier, prioritySupplier, releaseNotifier);
+ }
+ }
+}