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:30 UTC
[flink] 03/03: [FLINK-31645][network] Introduce the DiskIOScheduler to DiskTierProducerAgent
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 c4e55b588d1e5831702a76adbca7c1428d754a9b
Author: Wencong Liu <li...@163.com>
AuthorDate: Tue Jul 4 10:19:09 2023 +0800
[FLINK-31645][network] Introduce the DiskIOScheduler to DiskTierProducerAgent
---
.../netty/TieredStorageNettyServiceImpl.java | 6 +-
.../netty/TieredStorageResultSubpartitionView.java | 8 +-
.../partition/hybrid/tiered/tier/TierFactory.java | 10 ++-
.../tiered/tier/disk/DiskTierConsumerAgent.java | 85 ++++++++++++++++++++++
.../hybrid/tiered/tier/disk/DiskTierFactory.java | 21 +++++-
.../tiered/tier/disk/DiskTierProducerAgent.java | 39 +++++++++-
.../tiered/tier/memory/MemoryTierFactory.java | 10 ++-
.../netty/TieredStorageConsumerClientTest.java | 51 +++++++++----
.../hybrid/tiered/storage/TestingTierFactory.java | 10 ++-
.../tier/disk/DiskTierProducerAgentTest.java | 12 ++-
10 files changed, 224 insertions(+), 28 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageNettyServiceImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageNettyServiceImpl.java
index c42116cc02d..7f2f0dce8ba 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageNettyServiceImpl.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageNettyServiceImpl.java
@@ -176,11 +176,15 @@ public class TieredStorageNettyServiceImpl implements TieredStorageNettyService
NettyConnectionReaderAvailabilityAndPriorityHelper helper) {
List<NettyConnectionReaderRegistration> registrations =
getReaderRegistration(partitionId, subpartitionId);
+ boolean hasSetChannel = false;
for (NettyConnectionReaderRegistration registration : registrations) {
if (registration.trySetChannel(index, inputChannelProvider, helper)) {
- return;
+ hasSetChannel = true;
}
}
+ if (hasSetChannel) {
+ return;
+ }
NettyConnectionReaderRegistration registration = new NettyConnectionReaderRegistration();
registration.trySetChannel(index, inputChannelProvider, helper);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageResultSubpartitionView.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageResultSubpartitionView.java
index 5417bc616e6..56c43735a91 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageResultSubpartitionView.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageResultSubpartitionView.java
@@ -108,9 +108,11 @@ public class TieredStorageResultSubpartitionView implements ResultSubpartitionVi
@Override
public void notifyRequiredSegmentId(int segmentId) {
- requiredSegmentId = segmentId;
- stopSendingData = false;
- availabilityListener.notifyDataAvailable();
+ if (segmentId > requiredSegmentId) {
+ requiredSegmentId = segmentId;
+ stopSendingData = false;
+ availabilityListener.notifyDataAvailable();
+ }
}
@Override
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/TierFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/TierFactory.java
index 564238dc2b2..b5e6d3d0e3e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/TierFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/TierFactory.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier;
+import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileReader;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
@@ -26,7 +27,9 @@ import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.Tiere
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageMemoryManager;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageResourceRegistry;
+import java.time.Duration;
import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
/** A factory that creates all the components of a tier. */
public interface TierFactory {
@@ -44,7 +47,12 @@ public interface TierFactory {
PartitionFileReader partitionFileReader,
TieredStorageMemoryManager storageMemoryManager,
TieredStorageNettyService nettyService,
- TieredStorageResourceRegistry resourceRegistry);
+ TieredStorageResourceRegistry resourceRegistry,
+ BatchShuffleReadBufferPool bufferPool,
+ ScheduledExecutorService ioExecutor,
+ int maxRequestedBuffers,
+ Duration bufferRequestTimeout,
+ int maxBufferReadAhead);
/** Creates the consumer-side agent of a Tier. */
TierConsumerAgent createConsumerAgent(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierConsumerAgent.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierConsumerAgent.java
new file mode 100644
index 00000000000..77faf4f9711
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierConsumerAgent.java
@@ -0,0 +1,85 @@
+/*
+ * 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.tier.disk;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+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.runtime.io.network.partition.hybrid.tiered.netty.NettyConnectionReader;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.TieredStorageNettyService;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageConsumerSpec;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierConsumerAgent;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+
+/** The data client is used to fetch data from disk tier. */
+public class DiskTierConsumerAgent implements TierConsumerAgent {
+
+ private final Map<
+ TieredStoragePartitionId,
+ Map<TieredStorageSubpartitionId, CompletableFuture<NettyConnectionReader>>>
+ nettyConnectionReaders = new HashMap<>();
+
+ public DiskTierConsumerAgent(
+ List<TieredStorageConsumerSpec> tieredStorageConsumerSpecs,
+ TieredStorageNettyService nettyService) {
+ for (TieredStorageConsumerSpec tieredStorageConsumerSpec : tieredStorageConsumerSpecs) {
+ TieredStoragePartitionId partitionId = tieredStorageConsumerSpec.getPartitionId();
+ TieredStorageSubpartitionId subpartitionId =
+ tieredStorageConsumerSpec.getSubpartitionId();
+ nettyConnectionReaders
+ .computeIfAbsent(partitionId, ignore -> new HashMap<>())
+ .put(
+ subpartitionId,
+ nettyService.registerConsumer(partitionId, subpartitionId));
+ }
+ }
+
+ @Override
+ public void start() {
+ // noop
+ }
+
+ @Override
+ public Optional<Buffer> getNextBuffer(
+ TieredStoragePartitionId partitionId,
+ TieredStorageSubpartitionId subpartitionId,
+ int segmentId) {
+ try {
+ return nettyConnectionReaders
+ .get(partitionId)
+ .get(subpartitionId)
+ .get()
+ .readBuffer(segmentId);
+ } catch (InterruptedException | ExecutionException e) {
+ throw new RuntimeException("Failed to get next buffer.", e);
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ // noop
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierFactory.java
index b7614e12ac0..f06d6ddd04f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierFactory.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.disk;
+import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileReader;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
@@ -31,7 +32,9 @@ import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierFact
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierProducerAgent;
+import java.time.Duration;
import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
/** The implementation of {@link TierFactory} for disk tier. */
public class DiskTierFactory implements TierFactory {
@@ -64,7 +67,12 @@ public class DiskTierFactory implements TierFactory {
PartitionFileReader partitionFileReader,
TieredStorageMemoryManager storageMemoryManager,
TieredStorageNettyService nettyService,
- TieredStorageResourceRegistry resourceRegistry) {
+ TieredStorageResourceRegistry resourceRegistry,
+ BatchShuffleReadBufferPool bufferPool,
+ ScheduledExecutorService ioExecutor,
+ int maxRequestedBuffers,
+ Duration bufferRequestTimeout,
+ int maxBufferReadAhead) {
return new DiskTierProducerAgent(
partitionId,
numSubpartitions,
@@ -74,16 +82,21 @@ public class DiskTierFactory implements TierFactory {
minReservedDiskSpaceFraction,
isBroadcastOnly,
partitionFileWriter,
+ partitionFileReader,
storageMemoryManager,
nettyService,
- resourceRegistry);
+ resourceRegistry,
+ bufferPool,
+ ioExecutor,
+ maxRequestedBuffers,
+ bufferRequestTimeout,
+ maxBufferReadAhead);
}
@Override
public TierConsumerAgent createConsumerAgent(
List<TieredStorageConsumerSpec> tieredStorageConsumerSpecs,
TieredStorageNettyService nettyService) {
- // TODO, create the disk tier consumer agent.
- return null;
+ return new DiskTierConsumerAgent(tieredStorageConsumerSpecs, nettyService);
}
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgent.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgent.java
index dd7ae9c5136..1346c9378dd 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgent.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgent.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.disk;
+import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.api.EndOfSegmentEvent;
import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
import org.apache.flink.runtime.io.network.buffer.Buffer;
@@ -25,6 +26,7 @@ import org.apache.flink.runtime.io.network.partition.PartitionNotFoundException;
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.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileReader;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.NettyConnectionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.NettyConnectionWriter;
@@ -40,10 +42,12 @@ import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
+import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledExecutorService;
import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageUtils.DATA_FILE_SUFFIX;
import static org.apache.flink.util.Preconditions.checkArgument;
@@ -73,6 +77,8 @@ public class DiskTierProducerAgent implements TierProducerAgent, NettyServicePro
/** Record the number of buffers currently written to each subpartition. */
private final int[] currentSubpartitionWriteBuffers;
+ private final DiskIOScheduler diskIOScheduler;
+
private volatile boolean isReleased;
DiskTierProducerAgent(
@@ -84,9 +90,15 @@ public class DiskTierProducerAgent implements TierProducerAgent, NettyServicePro
float minReservedDiskSpaceFraction,
boolean isBroadcastOnly,
PartitionFileWriter partitionFileWriter,
+ PartitionFileReader partitionFileReader,
TieredStorageMemoryManager storageMemoryManager,
TieredStorageNettyService nettyService,
- TieredStorageResourceRegistry resourceRegistry) {
+ TieredStorageResourceRegistry resourceRegistry,
+ BatchShuffleReadBufferPool bufferPool,
+ ScheduledExecutorService ioExecutor,
+ int maxRequestedBuffers,
+ Duration bufferRequestTimeout,
+ int maxBufferReadAhead) {
checkArgument(
numBytesPerSegment >= bufferSizeBytes,
"One segment should contain at least one buffer.");
@@ -111,6 +123,19 @@ public class DiskTierProducerAgent implements TierProducerAgent, NettyServicePro
isBroadcastOnly ? 1 : numSubpartitions,
storageMemoryManager,
partitionFileWriter);
+
+ this.diskIOScheduler =
+ new DiskIOScheduler(
+ partitionId,
+ bufferPool,
+ ioExecutor,
+ maxRequestedBuffers,
+ bufferRequestTimeout,
+ maxBufferReadAhead,
+ nettyService,
+ this::retrieveFirstBufferIndexInSegment,
+ partitionFileReader);
+
nettyService.registerProducer(partitionId, this);
resourceRegistry.registerResource(partitionId, this::releaseResources);
}
@@ -155,10 +180,13 @@ public class DiskTierProducerAgent implements TierProducerAgent, NettyServicePro
new PartitionNotFoundException(
TieredStorageIdMappingUtils.convertId(partitionId)));
}
+ diskIOScheduler.connectionEstablished(subpartitionId, nettyConnectionWriter);
}
@Override
- public void connectionBroken(NettyConnectionId connectionId) {}
+ public void connectionBroken(NettyConnectionId connectionId) {
+ diskIOScheduler.connectionBroken(connectionId);
+ }
@Override
public void close() {
@@ -184,8 +212,15 @@ public class DiskTierProducerAgent implements TierProducerAgent, NettyServicePro
private void releaseResources() {
if (!isReleased) {
+ firstBufferIndexInSegment.clear();
diskCacheManager.release();
isReleased = true;
}
}
+
+ private Integer retrieveFirstBufferIndexInSegment(int subpartitionId, int bufferIndex) {
+ return firstBufferIndexInSegment.size() > subpartitionId
+ ? firstBufferIndexInSegment.get(subpartitionId).get(bufferIndex)
+ : null;
+ }
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/memory/MemoryTierFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/memory/MemoryTierFactory.java
index aa05da64a3f..81a3706a49e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/memory/MemoryTierFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/memory/MemoryTierFactory.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.memory;
+import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileReader;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
@@ -31,7 +32,9 @@ import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierFact
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierProducerAgent;
+import java.time.Duration;
import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
/** The implementation of {@link TierFactory} for memory tier. */
public class MemoryTierFactory implements TierFactory {
@@ -61,7 +64,12 @@ public class MemoryTierFactory implements TierFactory {
PartitionFileReader partitionFileReader,
TieredStorageMemoryManager memoryManager,
TieredStorageNettyService nettyService,
- TieredStorageResourceRegistry resourceRegistry) {
+ TieredStorageResourceRegistry resourceRegistry,
+ BatchShuffleReadBufferPool bufferPool,
+ ScheduledExecutorService ioExecutor,
+ int maxRequestedBuffers,
+ Duration bufferRequestTimeout,
+ int maxBufferReadAhead) {
return new MemoryTierProducerAgent(
partitionID,
numSubpartitions,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageConsumerClientTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageConsumerClientTest.java
index fa3063b32e2..6fe6c7a69e1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageConsumerClientTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/netty/TieredStorageConsumerClientTest.java
@@ -27,13 +27,17 @@ import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.Tiered
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TestingTierFactory;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageConsumerClient;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageConsumerSpec;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierConsumerAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.memory.MemoryTierConsumerAgent;
import org.junit.jupiter.api.Test;
import java.util.Collections;
+import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
+import java.util.function.BiFunction;
+import java.util.function.Function;
import static org.assertj.core.api.Assertions.assertThat;
@@ -46,28 +50,47 @@ class TieredStorageConsumerClientTest {
private final TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
@Test
- void testGetNextBuffer() {
+ void testGetNextBufferFromMemoryTier() {
Buffer buffer = BufferBuilderTestUtils.buildSomeBuffer(0);
+ Optional<Buffer> nextBuffer =
+ createTieredStorageConsumerClient(segmentId -> buffer, MemoryTierConsumerAgent::new)
+ .getNextBuffer(partitionId, subpartitionId);
+ assertThat(nextBuffer).hasValue(buffer);
+ }
+
+ @Test
+ void testGetNextBufferFromDiskTier() {
+ Buffer buffer = BufferBuilderTestUtils.buildSomeBuffer(0);
+ Optional<Buffer> nextBuffer =
+ createTieredStorageConsumerClient(segmentId -> buffer, MemoryTierConsumerAgent::new)
+ .getNextBuffer(partitionId, subpartitionId);
+ assertThat(nextBuffer).hasValue(buffer);
+ }
+
+ private TieredStorageConsumerClient createTieredStorageConsumerClient(
+ Function<Integer, Buffer> readBufferFunction,
+ BiFunction<
+ List<TieredStorageConsumerSpec>,
+ TieredStorageNettyService,
+ TierConsumerAgent>
+ tierConsumerAgentSupplier) {
+
TestingTieredStorageNettyService nettyService =
new TestingTieredStorageNettyService.Builder()
.setRegisterConsumerFunction(
(partitionId, subpartitionId) ->
CompletableFuture.completedFuture(
new TestingNettyConnectionReader.Builder()
- .setReadBufferFunction(segmentId -> buffer)
+ .setReadBufferFunction(readBufferFunction)
.build()))
.build();
- TieredStorageConsumerClient tieredStorageConsumerClient =
- new TieredStorageConsumerClient(
- Collections.singletonList(
- new TestingTierFactory.Builder()
- .setTierConsumerAgentSupplier(MemoryTierConsumerAgent::new)
- .build()),
- Collections.singletonList(
- new TieredStorageConsumerSpec(partitionId, subpartitionId)),
- nettyService);
- Optional<Buffer> nextBuffer =
- tieredStorageConsumerClient.getNextBuffer(partitionId, subpartitionId);
- assertThat(nextBuffer).hasValue(buffer);
+ return new TieredStorageConsumerClient(
+ Collections.singletonList(
+ new TestingTierFactory.Builder()
+ .setTierConsumerAgentSupplier(tierConsumerAgentSupplier)
+ .build()),
+ Collections.singletonList(
+ new TieredStorageConsumerSpec(partitionId, subpartitionId)),
+ nettyService);
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/TestingTierFactory.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/TestingTierFactory.java
index 7211e200d42..29652fc1ca3 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/TestingTierFactory.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/TestingTierFactory.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileReader;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
@@ -27,7 +28,9 @@ import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierFact
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierProducerAgent;
+import java.time.Duration;
import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
import java.util.function.BiFunction;
import java.util.function.Supplier;
@@ -71,7 +74,12 @@ public class TestingTierFactory implements TierFactory {
PartitionFileReader partitionFileReader,
TieredStorageMemoryManager storageMemoryManager,
TieredStorageNettyService nettyService,
- TieredStorageResourceRegistry resourceRegistry) {
+ TieredStorageResourceRegistry resourceRegistry,
+ BatchShuffleReadBufferPool bufferPool,
+ ScheduledExecutorService ioExecutor,
+ int maxRequestedBuffers,
+ Duration bufferRequestTimeout,
+ int maxBufferReadAhead) {
return tierProducerAgentSupplier.get();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgentTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgentTest.java
index 630767c3677..32037c7abda 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgentTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/disk/DiskTierProducerAgentTest.java
@@ -18,6 +18,8 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.disk;
+import org.apache.flink.runtime.concurrent.ManuallyTriggeredScheduledExecutorService;
+import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.buffer.BufferBuilderTestUtils;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.TestingTieredStorageMemoryManager;
@@ -26,6 +28,7 @@ import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.Tiered
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageUtils;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.PartitionFileWriter;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.TestingPartitionFileReader;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.file.TestingPartitionFileWriter;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.TestingNettyServiceProducer;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.TestingTieredStorageNettyService;
@@ -38,6 +41,7 @@ import org.junit.jupiter.api.io.TempDir;
import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
+import java.time.Duration;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.assertj.core.api.Assertions.assertThat;
@@ -203,8 +207,14 @@ public class DiskTierProducerAgentTest {
minReservedDiskSpaceFraction,
isBroadcastOnly,
partitionFileWriter,
+ new TestingPartitionFileReader.Builder().build(),
memoryManager,
nettyService,
- resourceRegistry);
+ resourceRegistry,
+ new BatchShuffleReadBufferPool(1, 1),
+ new ManuallyTriggeredScheduledExecutorService(),
+ 0,
+ Duration.ofMinutes(5),
+ 0);
}
}