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/18 02:02:13 UTC
[flink] 02/03: [FLINK-31647][network] Introduce the SegmentPartitionFileReader
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 1bf06cfefad9a751eed7ea84d8e63974ab0f0009
Author: Wencong Liu <li...@163.com>
AuthorDate: Mon Jul 10 11:27:16 2023 +0800
[FLINK-31647][network] Introduce the SegmentPartitionFileReader
---
.../network/partition/BufferReaderWriterUtil.java | 2 +-
.../hybrid/tiered/file/PartitionFileReader.java | 5 +-
.../tiered/file/SegmentPartitionFileReader.java | 119 ++++++++++++++++++-
.../file/SegmentPartitionFileReaderTest.java | 127 +++++++++++++++++++++
4 files changed, 246 insertions(+), 7 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferReaderWriterUtil.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferReaderWriterUtil.java
index 66eeba83c13..863be5b5b44 100755
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferReaderWriterUtil.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferReaderWriterUtil.java
@@ -294,7 +294,7 @@ public final class BufferReaderWriterUtil {
}
}
- static BufferHeader parseBufferHeader(ByteBuffer headerBuffer) {
+ public static BufferHeader parseBufferHeader(ByteBuffer headerBuffer) {
configureByteBuffer(headerBuffer);
short dataTypeIndex = headerBuffer.getShort();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/PartitionFileReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/PartitionFileReader.java
index ec1c75e6547..0a868c9002a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/PartitionFileReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/PartitionFileReader.java
@@ -24,6 +24,8 @@ 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 javax.annotation.Nullable;
+
import java.io.IOException;
/** {@link PartitionFileReader} defines the read logic for different types of shuffle files. */
@@ -38,8 +40,9 @@ public interface PartitionFileReader {
* @param bufferIndex the index of buffer
* @param memorySegment the empty buffer to store the read buffer
* @param recycler the buffer recycler
- * @return the read buffer
+ * @return null if there is no data otherwise a buffer.
*/
+ @Nullable
Buffer readBuffer(
TieredStoragePartitionId partitionId,
TieredStorageSubpartitionId subpartitionId,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReader.java
index 59281225bc9..9f79a78cbd8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReader.java
@@ -18,19 +18,59 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.file;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.core.fs.FileSystem;
+import org.apache.flink.core.fs.Path;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferHeader;
import org.apache.flink.runtime.io.network.buffer.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+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 java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+import static org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil.parseBufferHeader;
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.SegmentPartitionFile.getSegmentPath;
/** The implementation of {@link PartitionFileReader} with segment file mode. */
public class SegmentPartitionFileReader implements PartitionFileReader {
+ private final ByteBuffer reusedHeaderBuffer = BufferReaderWriterUtil.allocatedHeaderBuffer();
+
+ /**
+ * Opened file channels and segment id of related segment files stored in map.
+ *
+ * <p>The key is partition id and subpartition id. The value is file channel and segment id.
+ */
+ private final Map<
+ TieredStoragePartitionId,
+ Map<TieredStorageSubpartitionId, Tuple2<ReadableByteChannel, Integer>>>
+ openedChannelAndSegmentIds = new HashMap<>();
+
+ private final String dataFilePath;
+
+ private FileSystem fileSystem;
+
public SegmentPartitionFileReader(String dataFilePath) {
- // TODO, implement the SegmentPartitionFileReader
+ this.dataFilePath = dataFilePath;
+ try {
+ this.fileSystem = new Path(dataFilePath).getFileSystem();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e, "Failed to initialize the FileSystem.");
+ }
}
@Override
@@ -42,8 +82,45 @@ public class SegmentPartitionFileReader implements PartitionFileReader {
MemorySegment memorySegment,
BufferRecycler recycler)
throws IOException {
- // TODO, implement the SegmentPartitionFileReader
- return null;
+
+ // Get the channel of the segment file for a subpartition.
+ Map<TieredStorageSubpartitionId, Tuple2<ReadableByteChannel, Integer>> subpartitionInfo =
+ openedChannelAndSegmentIds.computeIfAbsent(partitionId, ignore -> new HashMap<>());
+ Tuple2<ReadableByteChannel, Integer> fileChannelAndSegmentId =
+ subpartitionInfo.getOrDefault(subpartitionId, Tuple2.of(null, -1));
+ ReadableByteChannel channel = fileChannelAndSegmentId.f0;
+
+ // Create the channel if there is a new segment file for a subpartition.
+ if (channel == null || fileChannelAndSegmentId.f1 != segmentId) {
+ if (channel != null) {
+ channel.close();
+ }
+ channel = openNewChannel(partitionId, subpartitionId, segmentId);
+ if (channel == null) {
+ // return null if the segment file doesn't exist.
+ return null;
+ }
+ subpartitionInfo.put(subpartitionId, Tuple2.of(channel, segmentId));
+ }
+
+ // Try to read a buffer from the channel.
+ reusedHeaderBuffer.clear();
+ int bufferHeaderResult = channel.read(reusedHeaderBuffer);
+ if (bufferHeaderResult == -1) {
+ channel.close();
+ openedChannelAndSegmentIds.get(partitionId).remove(subpartitionId);
+ return new NetworkBuffer(memorySegment, recycler, Buffer.DataType.END_OF_SEGMENT);
+ }
+ reusedHeaderBuffer.flip();
+ BufferHeader header = parseBufferHeader(reusedHeaderBuffer);
+ int dataBufferResult = channel.read(memorySegment.wrap(0, header.getLength()));
+ if (dataBufferResult != header.getLength()) {
+ channel.close();
+ throw new IOException("The length of data buffer is illegal.");
+ }
+ Buffer.DataType dataType = header.getDataType();
+ return new NetworkBuffer(
+ memorySegment, recycler, dataType, header.isCompressed(), header.getLength());
}
@Override
@@ -52,9 +129,41 @@ public class SegmentPartitionFileReader implements PartitionFileReader {
TieredStorageSubpartitionId subpartitionId,
int segmentId,
int bufferIndex) {
- return 0;
+ // noop
+ return -1;
+ }
+
+ private ReadableByteChannel openNewChannel(
+ TieredStoragePartitionId partitionId,
+ TieredStorageSubpartitionId subpartitionId,
+ int segmentId)
+ throws IOException {
+ Path currentSegmentPath =
+ getSegmentPath(
+ dataFilePath, partitionId, subpartitionId.getSubpartitionId(), segmentId);
+ if (!fileSystem.exists(currentSegmentPath)) {
+ return null;
+ }
+ return Channels.newChannel(fileSystem.open(currentSegmentPath));
}
@Override
- public void release() {}
+ public void release() {
+ openedChannelAndSegmentIds.values().stream()
+ .map(Map::values)
+ .flatMap(
+ (Function<
+ Collection<Tuple2<ReadableByteChannel, Integer>>,
+ Stream<Tuple2<ReadableByteChannel, Integer>>>)
+ Collection::stream)
+ .filter(Objects::nonNull)
+ .forEach(
+ channel -> {
+ try {
+ channel.f0.close();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e);
+ }
+ });
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReaderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReaderTest.java
new file mode 100644
index 00000000000..1ba7fb36070
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/file/SegmentPartitionFileReaderTest.java
@@ -0,0 +1,127 @@
+/*
+ * 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.fs.Path;
+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.util.List;
+
+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 SegmentPartitionFileReader}. */
+class SegmentPartitionFileReaderTest {
+
+ private static final TieredStoragePartitionId DEFAULT_PARTITION_ID =
+ TieredStorageIdMappingUtils.convertId(new ResultPartitionID());
+
+ private static final TieredStorageSubpartitionId DEFAULT_SUBPARTITION_ID =
+ new TieredStorageSubpartitionId(0);
+
+ private static final int DEFAULT_NUM_SUBPARTITION = 2;
+
+ private static final int DEFAULT_SEGMENT_NUM = 2;
+
+ private static final int DEFAULT_BUFFER_PER_SEGMENT = 3;
+
+ private static final int DEFAULT_BUFFER_SIZE = 1;
+
+ @TempDir private File tempFolder;
+
+ private SegmentPartitionFileReader partitionFileReader;
+
+ @BeforeEach
+ void before() {
+ Path tieredStorageDir = Path.fromLocalFile(tempFolder);
+ SegmentPartitionFileWriter partitionFileWriter =
+ new SegmentPartitionFileWriter(
+ tieredStorageDir.getPath(), DEFAULT_NUM_SUBPARTITION);
+
+ // Prepare the buffers to be written.
+ List<PartitionFileWriter.SubpartitionBufferContext> subpartitionBuffers =
+ generateBuffersToWrite(
+ DEFAULT_NUM_SUBPARTITION,
+ DEFAULT_SEGMENT_NUM,
+ DEFAULT_BUFFER_PER_SEGMENT,
+ DEFAULT_BUFFER_SIZE);
+
+ // Write the file.
+ partitionFileWriter.write(DEFAULT_PARTITION_ID, subpartitionBuffers);
+ partitionFileWriter.release();
+ partitionFileReader = new SegmentPartitionFileReader(tieredStorageDir.getPath());
+ }
+
+ @Test
+ void testReadBuffer() throws IOException {
+ for (int subpartitionId = 0; subpartitionId < DEFAULT_NUM_SUBPARTITION; ++subpartitionId) {
+ for (int segmentId = 0; segmentId < DEFAULT_SEGMENT_NUM; ++segmentId) {
+ for (int bufferIndex = 0; bufferIndex < DEFAULT_BUFFER_PER_SEGMENT; ++bufferIndex) {
+ Buffer buffer =
+ readBuffer(
+ bufferIndex,
+ new TieredStorageSubpartitionId(subpartitionId),
+ segmentId);
+ assertThat(buffer).isNotNull();
+ buffer.recycleBuffer();
+ }
+ }
+ }
+ }
+
+ @Test
+ void testGetPriority() throws IOException {
+ assertThat(
+ partitionFileReader.getPriority(
+ DEFAULT_PARTITION_ID, DEFAULT_SUBPARTITION_ID, 0, 0))
+ .isEqualTo(-1);
+ assertThat(readBuffer(0, DEFAULT_SUBPARTITION_ID, 0)).isNotNull();
+ assertThat(
+ partitionFileReader.getPriority(
+ DEFAULT_PARTITION_ID, DEFAULT_SUBPARTITION_ID, 0, 1))
+ .isEqualTo(-1);
+ }
+
+ private Buffer readBuffer(
+ int bufferIndex, TieredStorageSubpartitionId subpartitionId, int segmentId)
+ throws IOException {
+ MemorySegment memorySegment =
+ MemorySegmentFactory.allocateUnpooledSegment(DEFAULT_BUFFER_SIZE);
+ return partitionFileReader.readBuffer(
+ DEFAULT_PARTITION_ID,
+ subpartitionId,
+ segmentId,
+ bufferIndex,
+ memorySegment,
+ FreeingBufferRecycler.INSTANCE);
+ }
+}