You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by xt...@apache.org on 2022/10/31 03:33:12 UTC
[flink] branch master updated (e8e9db37e17 -> d11940c4a78)
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
from e8e9db37e17 [FLINK-29728] TablePlanner prevents Flink from starting is working directory is a symbolic link
new 840a5198df7 [hotfix] Migrate ResultPartitionFactoryTest to JUnit5 and AssertJ.
new 7f338e87a77 [FLINK-28889] Full spilling strategy is no longer consider consuming progress.
new 14392cf4ff2 [FLINK-28889] HsBufferContext supports multiple consumer.
new a2b402d66de [FLINK-28889] Introduce HsSubpartitionConsumerMemoryDataManager and let HsMemoryDataManager supports multiple consumer.
new f95c12e9a4a [FLINK-28889] HsFileDataManager supports multiple consumer.
new d65722f29a9 [FLINK-28889] Assign an unique identifier for each downstream consumer.
new c04985e6ac9 [FLINK-28889] fail fast when multiple consumer is not allowed but register many times
new d11940c4a78 [FLINK-28889] HsResultPartition support broadcast optimize
The 8 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../ResultPartitionDeploymentDescriptor.java | 5 +
.../network/partition/ResultPartitionFactory.java | 12 ++
.../network/partition/hybrid/HsBufferContext.java | 16 +-
.../network/partition/hybrid/HsConsumerId.java} | 43 ++---
.../io/network/partition/hybrid/HsDataView.java | 4 +-
.../partition/hybrid/HsFileDataManager.java | 8 +-
.../partition/hybrid/HsFullSpillingStrategy.java | 68 +++----
.../partition/hybrid/HsMemoryDataManager.java | 65 ++++---
.../hybrid/HsMemoryDataManagerOperation.java | 15 +-
.../partition/hybrid/HsResultPartition.java | 60 ++++--
.../hybrid/HsSelectiveSpillingStrategy.java | 16 +-
.../partition/hybrid/HsSpillingInfoProvider.java | 34 +++-
...titionView.java => HsSubpartitionConsumer.java} | 20 +-
... HsSubpartitionConsumerInternalOperations.java} | 6 +-
.../HsSubpartitionConsumerMemoryDataManager.java | 185 +++++++++++++++++++
.../partition/hybrid/HsSubpartitionFileReader.java | 3 +-
.../hybrid/HsSubpartitionFileReaderImpl.java | 16 +-
.../hybrid/HsSubpartitionMemoryDataManager.java | 191 +++++++------------
.../network/partition/ResultPartitionBuilder.java | 8 +
.../partition/ResultPartitionFactoryTest.java | 87 +++++----
.../partition/hybrid/HsBufferContextTest.java | 27 ++-
.../network/partition/hybrid/HsConsumerIdTest.java | 26 +--
.../partition/hybrid/HsFileDataManagerTest.java | 34 ++--
.../hybrid/HsFullSpillingStrategyTest.java | 39 +---
.../partition/hybrid/HsMemoryDataManagerTest.java | 20 ++
.../partition/hybrid/HsResultPartitionTest.java | 204 +++++++++++++++++++++
...sSubpartitionConsumerMemoryDataManagerTest.java | 197 ++++++++++++++++++++
.../hybrid/HsSubpartitionFileReaderImplTest.java | 84 ++++++---
.../HsSubpartitionMemoryDataManagerTest.java | 160 ++++++----------
.../partition/hybrid/HsSubpartitionViewTest.java | 60 +++---
.../hybrid/TestingMemoryDataManagerOperation.java | 26 ++-
.../hybrid/TestingSpillingInfoProvider.java | 10 +-
...tingSubpartitionConsumerInternalOperation.java} | 6 +-
.../shuffle/PartitionDescriptorBuilder.java | 9 +-
34 files changed, 1227 insertions(+), 537 deletions(-)
copy flink-runtime/src/main/java/org/apache/flink/runtime/{rest/handler/async/OperationKey.java => io/network/partition/hybrid/HsConsumerId.java} (54%)
rename flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/{HsSubpartitionView.java => HsSubpartitionConsumer.java} (94%)
rename flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/{HsSubpartitionViewInternalOperations.java => HsSubpartitionConsumerInternalOperations.java} (87%)
create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManager.java
copy flink-runtime-web/src/test/java/org/apache/flink/runtime/webmonitor/handlers/JarDeleteHeadersTest.java => flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerIdTest.java (56%)
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManagerTest.java
rename flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/{TestingSubpartitionViewInternalOperation.java => TestingSubpartitionConsumerInternalOperation.java} (88%)
[flink] 02/08: [FLINK-28889] Full spilling strategy is no longer consider consuming progress.
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 7f338e87a777abd102719dcb772b184770b62b88
Author: Weijie Guo <re...@163.com>
AuthorDate: Wed Oct 19 19:44:16 2022 +0800
[FLINK-28889] Full spilling strategy is no longer consider consuming progress.
---
.../partition/hybrid/HsFullSpillingStrategy.java | 58 ++++++++--------------
.../hybrid/HsFullSpillingStrategyTest.java | 39 +--------------
2 files changed, 21 insertions(+), 76 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java
index f0339ee152a..ac2d86fcc30 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java
@@ -22,14 +22,10 @@ import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvid
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
import java.util.ArrayDeque;
-import java.util.ArrayList;
import java.util.Deque;
-import java.util.List;
import java.util.Optional;
import java.util.TreeMap;
-import static org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingStrategyUtils.getBuffersByConsumptionPriorityInOrder;
-
/** A special implementation of {@link HsSpillingStrategy} that spilled all buffers to disk. */
public class HsFullSpillingStrategy implements HsSpillingStrategy {
private final float numBuffersTriggerSpillingRatio;
@@ -118,6 +114,11 @@ public class HsFullSpillingStrategy implements HsSpillingStrategy {
}
}
+ /**
+ * Release subpartition's spilled buffer from head. Each subpartition fairly retains a fixed
+ * number of buffers, and all the remaining buffers are released. If this subpartition does not
+ * have so many qualified buffers, all of them will be retained.
+ */
private void checkRelease(
HsSpillingInfoProvider spillingInfoProvider, int poolSize, Decision.Builder builder) {
if (spillingInfoProvider.getNumTotalRequestedBuffers() < poolSize * releaseThreshold) {
@@ -125,47 +126,28 @@ public class HsFullSpillingStrategy implements HsSpillingStrategy {
return;
}
- int releaseNum = (int) (spillingInfoProvider.getPoolSize() * releaseBufferRatio);
+ int survivedNum = (int) (poolSize - poolSize * releaseBufferRatio);
+ int numSubpartitions = spillingInfoProvider.getNumSubpartitions();
+ int subpartitionSurvivedNum = survivedNum / numSubpartitions;
- // first, release all consumed buffers
- TreeMap<Integer, Deque<BufferIndexAndChannel>> consumedBuffersToRelease = new TreeMap<>();
- int numConsumedBuffers = 0;
- for (int subpartitionId = 0;
- subpartitionId < spillingInfoProvider.getNumSubpartitions();
- subpartitionId++) {
+ TreeMap<Integer, Deque<BufferIndexAndChannel>> bufferToRelease = new TreeMap<>();
- Deque<BufferIndexAndChannel> consumedSpillSubpartitionBuffers =
+ for (int subpartitionId = 0; subpartitionId < numSubpartitions; subpartitionId++) {
+ Deque<BufferIndexAndChannel> buffersInOrder =
spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.SPILL, ConsumeStatus.CONSUMED);
- numConsumedBuffers += consumedSpillSubpartitionBuffers.size();
- consumedBuffersToRelease.put(subpartitionId, consumedSpillSubpartitionBuffers);
- }
-
- // make up the releaseNum with unconsumed buffers, if needed, w.r.t. the consuming priority
- TreeMap<Integer, List<BufferIndexAndChannel>> unconsumedBufferToRelease = new TreeMap<>();
- if (releaseNum > numConsumedBuffers) {
- TreeMap<Integer, Deque<BufferIndexAndChannel>> unconsumedBuffers = new TreeMap<>();
- for (int subpartitionId = 0;
- subpartitionId < spillingInfoProvider.getNumSubpartitions();
- subpartitionId++) {
- unconsumedBuffers.put(
- subpartitionId,
- spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.SPILL, ConsumeStatus.NOT_CONSUMED));
+ subpartitionId, SpillStatus.SPILL, ConsumeStatus.ALL);
+ // if the number of subpartition buffers less than survived buffers, reserved all of
+ // them.
+ int releaseNum = Math.max(0, buffersInOrder.size() - subpartitionSurvivedNum);
+ while (releaseNum-- != 0) {
+ buffersInOrder.pollLast();
}
- unconsumedBufferToRelease.putAll(
- getBuffersByConsumptionPriorityInOrder(
- spillingInfoProvider.getNextBufferIndexToConsume(),
- unconsumedBuffers,
- releaseNum - numConsumedBuffers));
+ bufferToRelease.put(subpartitionId, buffersInOrder);
}
// collect results in order
- for (int i = 0; i < spillingInfoProvider.getNumSubpartitions(); i++) {
- List<BufferIndexAndChannel> toRelease = new ArrayList<>();
- toRelease.addAll(consumedBuffersToRelease.getOrDefault(i, new ArrayDeque<>()));
- toRelease.addAll(unconsumedBufferToRelease.getOrDefault(i, new ArrayList<>()));
- builder.addBufferToRelease(i, toRelease);
+ for (int i = 0; i < numSubpartitions; i++) {
+ builder.addBufferToRelease(i, bufferToRelease.getOrDefault(i, new ArrayDeque<>()));
}
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategyTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategyTest.java
index 3602b15b2d1..e88918b2c0d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategyTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategyTest.java
@@ -32,7 +32,6 @@ import java.util.Optional;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createBufferIndexAndChannelsList;
import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.entry;
/** Tests for {@link HsFullSpillingStrategy}. */
class HsFullSpillingStrategyTest {
@@ -129,9 +128,7 @@ class HsFullSpillingStrategyTest {
.addSubpartitionBuffers(subpartition1, subpartitionBuffers1)
.addSubpartitionBuffers(subpartition2, subpartitionBuffers2)
.addSpillBuffers(subpartition1, Arrays.asList(0, 1, 2, 3))
- .addConsumedBuffers(subpartition1, Arrays.asList(0, 1))
.addSpillBuffers(subpartition2, Arrays.asList(1, 2, 3))
- .addConsumedBuffers(subpartition2, Arrays.asList(0, 1))
.setGetNumTotalUnSpillBuffersSupplier(
() -> (int) (10 * NUM_BUFFERS_TRIGGER_SPILLING_RATIO))
.setGetNumTotalRequestedBuffersSupplier(() -> 10)
@@ -151,47 +148,13 @@ class HsFullSpillingStrategyTest {
assertThat(decision.getBufferToSpill()).isEqualTo(expectedSpillBuffers);
Map<Integer, List<BufferIndexAndChannel>> expectedReleaseBuffers = new HashMap<>();
- // all consumed spill buffers should release.
expectedReleaseBuffers.put(
subpartition1, new ArrayList<>(subpartitionBuffers1.subList(0, 2)));
- // priority higher buffers should release.
- expectedReleaseBuffers.get(subpartition1).addAll(subpartitionBuffers1.subList(3, 4));
- // all consumed spill buffers should release.
expectedReleaseBuffers.put(
- subpartition2, new ArrayList<>(subpartitionBuffers2.subList(1, 2)));
- // priority higher buffers should release.
- expectedReleaseBuffers.get(subpartition2).addAll(subpartitionBuffers2.subList(2, 4));
+ subpartition2, new ArrayList<>(subpartitionBuffers2.subList(1, 3)));
assertThat(decision.getBufferToRelease()).isEqualTo(expectedReleaseBuffers);
}
- /** All consumed buffers that already spill should release regardless of the release ratio. */
- @Test
- void testDecideActionWithGlobalInfoAllConsumedSpillBufferShouldRelease() {
- final int subpartitionId = 0;
- List<BufferIndexAndChannel> subpartitionBuffers =
- createBufferIndexAndChannelsList(subpartitionId, 0, 1, 2, 3, 4);
-
- final int poolSize = 5;
- TestingSpillingInfoProvider spillInfoProvider =
- TestingSpillingInfoProvider.builder()
- .setGetNumSubpartitionsSupplier(() -> 1)
- .addSubpartitionBuffers(subpartitionId, subpartitionBuffers)
- .addSpillBuffers(subpartitionId, Arrays.asList(0, 1, 2, 3, 4))
- .addConsumedBuffers(subpartitionId, Arrays.asList(0, 1, 2, 3))
- .setGetNumTotalUnSpillBuffersSupplier(() -> 0)
- .setGetNumTotalRequestedBuffersSupplier(() -> poolSize)
- .setGetPoolSizeSupplier(() -> poolSize)
- .build();
-
- int numReleaseBuffer = (int) (poolSize * FULL_SPILL_RELEASE_RATIO);
- Decision decision = spillStrategy.decideActionWithGlobalInfo(spillInfoProvider);
- assertThat(decision.getBufferToSpill()).isEmpty();
- assertThat(decision.getBufferToRelease())
- .containsOnly(entry(subpartitionId, subpartitionBuffers.subList(0, 4)))
- .extractingByKey(subpartitionId)
- .satisfies((buffers) -> assertThat(buffers).hasSizeGreaterThan(numReleaseBuffer));
- }
-
@Test
void testOnResultPartitionClosed() {
final int subpartition1 = 0;
[flink] 07/08: [FLINK-28889] fail fast when multiple consumer is not allowed but register many times
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit c04985e6ac968af2181ecaa7f0356d54fdd47848
Author: Weijie Guo <re...@163.com>
AuthorDate: Wed Oct 26 11:57:22 2022 +0800
[FLINK-28889] fail fast when multiple consumer is not allowed but register many times
---
.../partition/hybrid/HsResultPartition.java | 16 +++++++-
.../partition/hybrid/HsResultPartitionTest.java | 46 ++++++++++++++++++++++
2 files changed, 61 insertions(+), 1 deletion(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
index 06a15a507c5..26aee9450ad 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
@@ -36,6 +36,7 @@ import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
import org.apache.flink.runtime.io.network.partition.ResultSubpartitionView;
+import org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleConfiguration.SpillingStrategyType;
import org.apache.flink.runtime.metrics.groups.TaskIOMetricGroup;
import org.apache.flink.util.function.SupplierWithException;
@@ -192,9 +193,11 @@ public class HsResultPartition extends ResultPartition {
HsSubpartitionConsumer subpartitionConsumer =
new HsSubpartitionConsumer(availabilityListener);
+ HsConsumerId lastConsumerId = lastConsumerIds[subpartitionId];
+ checkMultipleConsumerIsAllowed(lastConsumerId, hybridShuffleConfiguration);
// assign a unique id for each consumer, now it is guaranteed by the value that is one
// higher than the last consumerId's id field.
- HsConsumerId consumerId = HsConsumerId.newId(lastConsumerIds[subpartitionId]);
+ HsConsumerId consumerId = HsConsumerId.newId(lastConsumerId);
lastConsumerIds[subpartitionId] = consumerId;
HsDataView diskDataView =
fileDataManager.registerNewConsumer(
@@ -295,4 +298,15 @@ public class HsResultPartition extends ResultPartition {
throw new IllegalConfigurationException("Illegal spilling strategy.");
}
}
+
+ private void checkMultipleConsumerIsAllowed(
+ HsConsumerId lastConsumerId, HybridShuffleConfiguration hybridShuffleConfiguration) {
+ if (hybridShuffleConfiguration.getSpillingStrategyType()
+ == SpillingStrategyType.SELECTIVE) {
+ checkState(
+ lastConsumerId == null,
+ "Multiple consumer is not allowed for %s spilling strategy mode",
+ hybridShuffleConfiguration.getSpillingStrategyType());
+ }
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
index 18425d5ea01..296834db22b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
@@ -67,6 +67,7 @@ import java.util.function.BiConsumer;
import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatNoException;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link HsResultPartition}. */
@@ -381,6 +382,51 @@ class HsResultPartitionTest {
}
}
+ @Test
+ void testSelectiveSpillingStrategyRegisterMultipleConsumer() throws Exception {
+ final int numSubpartitions = 2;
+ BufferPool bufferPool = globalPool.createBufferPool(2, 2);
+ try (HsResultPartition partition =
+ createHsResultPartition(
+ 2,
+ bufferPool,
+ HybridShuffleConfiguration.builder(
+ numSubpartitions, readBufferPool.getNumBuffersPerRequest())
+ .setSpillingStrategyType(
+ HybridShuffleConfiguration.SpillingStrategyType.SELECTIVE)
+ .build())) {
+ partition.createSubpartitionView(0, new NoOpBufferAvailablityListener());
+ assertThatThrownBy(
+ () ->
+ partition.createSubpartitionView(
+ 0, new NoOpBufferAvailablityListener()))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("Multiple consumer is not allowed");
+ }
+ }
+
+ @Test
+ void testFullSpillingStrategyRegisterMultipleConsumer() throws Exception {
+ final int numSubpartitions = 2;
+ BufferPool bufferPool = globalPool.createBufferPool(2, 2);
+ try (HsResultPartition partition =
+ createHsResultPartition(
+ 2,
+ bufferPool,
+ HybridShuffleConfiguration.builder(
+ numSubpartitions, readBufferPool.getNumBuffersPerRequest())
+ .setSpillingStrategyType(
+ HybridShuffleConfiguration.SpillingStrategyType.FULL)
+ .build())) {
+ partition.createSubpartitionView(0, new NoOpBufferAvailablityListener());
+ assertThatNoException()
+ .isThrownBy(
+ () ->
+ partition.createSubpartitionView(
+ 0, new NoOpBufferAvailablityListener()));
+ }
+ }
+
private static void recordDataWritten(
ByteBuffer record,
Queue<Tuple2<ByteBuffer, Buffer.DataType>>[] dataWritten,
[flink] 06/08: [FLINK-28889] Assign an unique identifier for each downstream consumer.
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit d65722f29a92440b63a98543b7ca62bd5ccd145e
Author: Weijie Guo <re...@163.com>
AuthorDate: Wed Oct 19 17:28:58 2022 +0800
[FLINK-28889] Assign an unique identifier for each downstream consumer.
---
.../partition/hybrid/HsResultPartition.java | 24 +++++---
.../partition/hybrid/HsResultPartitionTest.java | 72 ++++++++++++++++++++++
2 files changed, 87 insertions(+), 9 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
index 678fb697b29..06a15a507c5 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
@@ -69,6 +69,9 @@ public class HsResultPartition extends ResultPartition {
private final HybridShuffleConfiguration hybridShuffleConfiguration;
+ /** Record the last assigned consumerId for each subpartition. */
+ private final HsConsumerId[] lastConsumerIds;
+
private boolean hasNotifiedEndOfUserRecords;
@Nullable private HsMemoryDataManager memoryDataManager;
@@ -110,6 +113,7 @@ public class HsResultPartition extends ResultPartition {
dataFilePath,
HsSubpartitionFileReaderImpl.Factory.INSTANCE,
hybridShuffleConfiguration);
+ this.lastConsumerIds = new HsConsumerId[numSubpartitions];
}
// Called by task thread.
@@ -186,21 +190,23 @@ public class HsResultPartition extends ResultPartition {
throw new PartitionNotFoundException(getPartitionId());
}
- HsSubpartitionConsumer subpartitionView = new HsSubpartitionConsumer(availabilityListener);
+ HsSubpartitionConsumer subpartitionConsumer =
+ new HsSubpartitionConsumer(availabilityListener);
+ // assign a unique id for each consumer, now it is guaranteed by the value that is one
+ // higher than the last consumerId's id field.
+ HsConsumerId consumerId = HsConsumerId.newId(lastConsumerIds[subpartitionId]);
+ lastConsumerIds[subpartitionId] = consumerId;
HsDataView diskDataView =
- // TODO pass real consumerId in the next commit.
fileDataManager.registerNewConsumer(
- subpartitionId, HsConsumerId.DEFAULT, subpartitionView);
+ subpartitionId, consumerId, subpartitionConsumer);
HsDataView memoryDataView =
checkNotNull(memoryDataManager)
- // TODO pass real consumerId in the next commit.
- .registerNewConsumer(
- subpartitionId, HsConsumerId.DEFAULT, subpartitionView);
+ .registerNewConsumer(subpartitionId, consumerId, subpartitionConsumer);
- subpartitionView.setDiskDataView(diskDataView);
- subpartitionView.setMemoryDataView(memoryDataView);
- return subpartitionView;
+ subpartitionConsumer.setDiskDataView(diskDataView);
+ subpartitionConsumer.setMemoryDataView(memoryDataView);
+ return subpartitionConsumer;
}
@Override
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
index 70378a453d0..18425d5ea01 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
@@ -221,6 +221,63 @@ class HsResultPartitionTest {
}
}
+ /** Test write and read data from single subpartition with multiple consumer. */
+ @Test
+ void testMultipleConsumer() throws Exception {
+ final int numBuffers = 10;
+ final int numRecords = 10;
+ final int numConsumers = 2;
+ final int targetChannel = 0;
+ final Random random = new Random();
+
+ BufferPool bufferPool = globalPool.createBufferPool(numBuffers, numBuffers);
+ try (HsResultPartition resultPartition = createHsResultPartition(2, bufferPool)) {
+ List<ByteBuffer> dataWritten = new ArrayList<>();
+ for (int i = 0; i < numRecords; i++) {
+ ByteBuffer record = generateRandomData(bufferSize, random);
+ resultPartition.emitRecord(record, targetChannel);
+ dataWritten.add(record);
+ }
+ resultPartition.finish();
+
+ Tuple2[] viewAndListeners =
+ createMultipleConsumerView(resultPartition, targetChannel, 2);
+
+ List<List<Buffer>> dataRead = new ArrayList<>();
+ for (int i = 0; i < numConsumers; i++) {
+ dataRead.add(new ArrayList<>());
+ }
+ readData(
+ viewAndListeners,
+ (buffer, subpartition) -> {
+ int numBytes = buffer.readableBytes();
+ if (buffer.isBuffer()) {
+ MemorySegment segment =
+ MemorySegmentFactory.allocateUnpooledSegment(numBytes);
+ segment.put(0, buffer.getNioBufferReadable(), numBytes);
+ dataRead.get(subpartition)
+ .add(
+ new NetworkBuffer(
+ segment,
+ (buf) -> {},
+ buffer.getDataType(),
+ numBytes));
+ }
+ });
+
+ for (int i = 0; i < numConsumers; i++) {
+ assertThat(dataWritten).hasSameSizeAs(dataRead.get(i));
+ List<Buffer> readBufferList = dataRead.get(i);
+ for (int j = 0; j < dataWritten.size(); j++) {
+ ByteBuffer bufferWritten = dataWritten.get(j);
+ bufferWritten.rewind();
+ Buffer bufferRead = readBufferList.get(j);
+ assertThat(bufferRead.getNioBufferReadable()).isEqualTo(bufferWritten);
+ }
+ }
+ }
+ }
+
@Test
void testClose() throws Exception {
final int numBuffers = 1;
@@ -484,6 +541,21 @@ class HsResultPartitionTest {
return viewAndListeners;
}
+ /** Create multiple consumer and bufferAvailabilityListener for single subpartition. */
+ private Tuple2<ResultSubpartitionView, TestingBufferAvailabilityListener>[]
+ createMultipleConsumerView(
+ HsResultPartition partition, int subpartitionId, int numConsumers)
+ throws Exception {
+ Tuple2<ResultSubpartitionView, TestingBufferAvailabilityListener>[] viewAndListeners =
+ new Tuple2[numConsumers];
+ for (int consumer = 0; consumer < numConsumers; ++consumer) {
+ TestingBufferAvailabilityListener listener = new TestingBufferAvailabilityListener();
+ viewAndListeners[consumer] =
+ Tuple2.of(partition.createSubpartitionView(subpartitionId, listener), listener);
+ }
+ return viewAndListeners;
+ }
+
private static final class TestingBufferAvailabilityListener
implements BufferAvailabilityListener {
[flink] 03/08: [FLINK-28889] HsBufferContext supports multiple consumer.
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 14392cf4ff291d1dbe316f4687e79760591127f3
Author: Weijie Guo <re...@163.com>
AuthorDate: Mon Sep 26 17:02:22 2022 +0800
[FLINK-28889] HsBufferContext supports multiple consumer.
---
.../network/partition/hybrid/HsBufferContext.java | 16 +++---
.../io/network/partition/hybrid/HsConsumerId.java | 63 ++++++++++++++++++++++
.../partition/hybrid/HsFullSpillingStrategy.java | 12 +++--
.../partition/hybrid/HsMemoryDataManager.java | 5 +-
.../hybrid/HsSelectiveSpillingStrategy.java | 13 +++--
.../partition/hybrid/HsSpillingInfoProvider.java | 24 ++++++++-
.../hybrid/HsSubpartitionMemoryDataManager.java | 22 ++++----
.../partition/hybrid/HsBufferContextTest.java | 27 +++++++---
.../network/partition/hybrid/HsConsumerIdTest.java | 43 +++++++++++++++
.../HsSubpartitionMemoryDataManagerTest.java | 28 ++++++----
.../hybrid/TestingSpillingInfoProvider.java | 8 +--
11 files changed, 212 insertions(+), 49 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContext.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContext.java
index 4a5f5fd9651..06970c97e54 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContext.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContext.java
@@ -22,8 +22,11 @@ import org.apache.flink.runtime.io.network.buffer.Buffer;
import javax.annotation.Nullable;
+import java.util.Collections;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
import static org.apache.flink.util.Preconditions.checkState;
@@ -37,7 +40,7 @@ import static org.apache.flink.util.Preconditions.checkState;
* can no longer be spilled or consumed.
* <li>{@link #spillStarted} indicates that spilling of the buffer has started, either completed
* or not.
- * <li>{@link #consumed} indicates that buffer has been consumed by the downstream.
+ * <li>{@link #consumed} indicates that buffer has been consumed by these consumers.
* </ul>
*
* <p>Reference count of the buffer is maintained as follows: *
@@ -64,7 +67,7 @@ public class HsBufferContext {
private boolean spillStarted;
- private boolean consumed;
+ private final Set<HsConsumerId> consumed = Collections.newSetFromMap(new ConcurrentHashMap<>());
@Nullable private CompletableFuture<Void> spilledFuture;
@@ -89,8 +92,8 @@ public class HsBufferContext {
return spillStarted;
}
- public boolean isConsumed() {
- return consumed;
+ public boolean isConsumed(HsConsumerId consumerId) {
+ return consumed.contains(consumerId);
}
public Optional<CompletableFuture<Void>> getSpilledFuture() {
@@ -127,10 +130,9 @@ public class HsBufferContext {
return true;
}
- public void consumed() {
+ public void consumed(HsConsumerId consumerId) {
checkState(!released, "Buffer is already released.");
- checkState(!consumed, "Consume buffer repeatedly is unexpected.");
- consumed = true;
+ checkState(consumed.add(consumerId), "Consume buffer repeatedly is unexpected.");
// increase ref count when buffer is consumed, will be decreased when downstream finish
// consuming.
buffer.retainBuffer();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerId.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerId.java
new file mode 100644
index 00000000000..132da40f6d7
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerId.java
@@ -0,0 +1,63 @@
+/*
+ * 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;
+
+import javax.annotation.Nullable;
+
+import java.util.Objects;
+
+/** This class represents the identifier of hybrid shuffle's consumer. */
+public class HsConsumerId {
+ /**
+ * This consumer id is used in the scenarios that information related to specific consumer needs
+ * to be ignored.
+ */
+ public static final HsConsumerId ANY = new HsConsumerId(-1);
+
+ /** This consumer id is used for the first consumer of a single subpartition. */
+ public static final HsConsumerId DEFAULT = new HsConsumerId(0);
+
+ /** This is a unique field for each consumer of a single subpartition. */
+ private final int id;
+
+ private HsConsumerId(int id) {
+ this.id = id;
+ }
+
+ public static HsConsumerId newId(@Nullable HsConsumerId lastId) {
+ return lastId == null ? DEFAULT : new HsConsumerId(lastId.id + 1);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ HsConsumerId that = (HsConsumerId) o;
+ return id == that.id;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(id);
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java
index ac2d86fcc30..94a067fbaeb 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFullSpillingStrategy.java
@@ -18,7 +18,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid;
-import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatus;
+import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
import java.util.ArrayDeque;
@@ -88,12 +88,14 @@ public class HsFullSpillingStrategy implements HsSpillingStrategy {
subpartitionId,
// get all not start spilling buffers.
spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.NOT_SPILL, ConsumeStatus.ALL))
+ subpartitionId,
+ SpillStatus.NOT_SPILL,
+ ConsumeStatusWithId.ALL_ANY))
.addBufferToRelease(
subpartitionId,
// get all not released buffers.
spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.ALL, ConsumeStatus.ALL));
+ subpartitionId, SpillStatus.ALL, ConsumeStatusWithId.ALL_ANY));
}
return builder.build();
}
@@ -110,7 +112,7 @@ public class HsFullSpillingStrategy implements HsSpillingStrategy {
builder.addBufferToSpill(
i,
spillingInfoProvider.getBuffersInOrder(
- i, SpillStatus.NOT_SPILL, ConsumeStatus.ALL));
+ i, SpillStatus.NOT_SPILL, ConsumeStatusWithId.ALL_ANY));
}
}
@@ -135,7 +137,7 @@ public class HsFullSpillingStrategy implements HsSpillingStrategy {
for (int subpartitionId = 0; subpartitionId < numSubpartitions; subpartitionId++) {
Deque<BufferIndexAndChannel> buffersInOrder =
spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.SPILL, ConsumeStatus.ALL);
+ subpartitionId, SpillStatus.SPILL, ConsumeStatusWithId.ALL_ANY);
// if the number of subpartition buffers less than survived buffers, reserved all of
// them.
int releaseNum = Math.max(0, buffersInOrder.size() - subpartitionSurvivedNum);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
index f19cb58a795..a8d7d89e7e4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
@@ -223,10 +223,11 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
// Write lock should be acquired before invoke this method.
@Override
public Deque<BufferIndexAndChannel> getBuffersInOrder(
- int subpartitionId, SpillStatus spillStatus, ConsumeStatus consumeStatus) {
+ int subpartitionId, SpillStatus spillStatus, ConsumeStatusWithId consumeStatusWithId) {
HsSubpartitionMemoryDataManager targetSubpartitionDataManager =
getSubpartitionMemoryDataManager(subpartitionId);
- return targetSubpartitionDataManager.getBuffersSatisfyStatus(spillStatus, consumeStatus);
+ return targetSubpartitionDataManager.getBuffersSatisfyStatus(
+ spillStatus, consumeStatusWithId);
}
// Write lock should be acquired before invoke this method.
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java
index dc356cb20d5..8c7217a4b93 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java
@@ -19,6 +19,7 @@
package org.apache.flink.runtime.io.network.partition.hybrid;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatus;
+import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
import java.util.Deque;
@@ -85,7 +86,11 @@ public class HsSelectiveSpillingStrategy implements HsSpillingStrategy {
subpartitionToBuffers.put(
channel,
spillingInfoProvider.getBuffersInOrder(
- channel, SpillStatus.NOT_SPILL, ConsumeStatus.NOT_CONSUMED));
+ channel,
+ SpillStatus.NOT_SPILL,
+ // selective spilling strategy does not support multiple consumer.
+ ConsumeStatusWithId.fromStatusAndConsumerId(
+ ConsumeStatus.NOT_CONSUMED, HsConsumerId.DEFAULT)));
}
TreeMap<Integer, List<BufferIndexAndChannel>> subpartitionToHighPriorityBuffers =
@@ -113,12 +118,14 @@ public class HsSelectiveSpillingStrategy implements HsSpillingStrategy {
subpartitionId,
// get all not start spilling buffers.
spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.NOT_SPILL, ConsumeStatus.ALL))
+ subpartitionId,
+ SpillStatus.NOT_SPILL,
+ ConsumeStatusWithId.ALL_ANY))
.addBufferToRelease(
subpartitionId,
// get all not released buffers.
spillingInfoProvider.getBuffersInOrder(
- subpartitionId, SpillStatus.ALL, ConsumeStatus.ALL));
+ subpartitionId, SpillStatus.ALL, ConsumeStatusWithId.ALL_ANY));
}
return builder.build();
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java
index 6168c637de4..deb32e7058e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java
@@ -43,13 +43,13 @@ public interface HsSpillingInfoProvider {
*
* @param subpartitionId target buffers belong to.
* @param spillStatus expected buffer spill status.
- * @param consumeStatus expected buffer consume status.
+ * @param consumeStatusWithId expected buffer consume status and consumer id.
* @return all buffers satisfy specific status of this subpartition, This queue must be sorted
* according to bufferIndex from small to large, in other words, head is the buffer with the
* minimum bufferIndex in the current subpartition.
*/
Deque<BufferIndexAndChannel> getBuffersInOrder(
- int subpartitionId, SpillStatus spillStatus, ConsumeStatus consumeStatus);
+ int subpartitionId, SpillStatus spillStatus, ConsumeStatusWithId consumeStatusWithId);
/** Get total number of not decided to spill buffers. */
int getNumTotalUnSpillBuffers();
@@ -79,4 +79,24 @@ public interface HsSpillingInfoProvider {
/** The buffer is either consumed or not consumed. */
ALL
}
+
+ /** This class represents a pair of {@link ConsumeStatus} and consumer id. */
+ class ConsumeStatusWithId {
+ public static final ConsumeStatusWithId ALL_ANY =
+ new ConsumeStatusWithId(ConsumeStatus.ALL, HsConsumerId.ANY);
+
+ ConsumeStatus status;
+
+ HsConsumerId consumerId;
+
+ private ConsumeStatusWithId(ConsumeStatus status, HsConsumerId consumerId) {
+ this.status = status;
+ this.consumerId = consumerId;
+ }
+
+ public static ConsumeStatusWithId fromStatusAndConsumerId(
+ ConsumeStatus consumeStatus, HsConsumerId consumerId) {
+ return new ConsumeStatusWithId(consumeStatus, consumerId);
+ }
+ }
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
index 619bc01d05a..f6804a62b1b 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
@@ -30,6 +30,7 @@ import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAndBacklog;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatus;
+import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
import org.apache.flink.util.function.SupplierWithException;
import org.apache.flink.util.function.ThrowingRunnable;
@@ -144,7 +145,8 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
HsBufferContext bufferContext =
checkNotNull(unConsumedBuffers.pollFirst());
- bufferContext.consumed();
+ // TODO move this logical to consumer and pass real consumerId.
+ bufferContext.consumed(HsConsumerId.DEFAULT);
DataType nextDataType =
peekNextToConsumeDataTypeInternal(toConsumeIndex + 1);
return Optional.of(Tuple2.of(bufferContext, nextDataType));
@@ -200,14 +202,14 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
* ConsumeStatus}.
*
* @param spillStatus the status of spilling expected.
- * @param consumeStatus the status of consuming expected.
+ * @param consumeStatusWithId the status and consumerId expected.
* @return buffers satisfy expected status in order.
*/
@SuppressWarnings("FieldAccessNotGuarded")
// Note that: callWithLock ensure that code block guarded by resultPartitionReadLock and
// subpartitionLock.
public Deque<BufferIndexAndChannel> getBuffersSatisfyStatus(
- SpillStatus spillStatus, ConsumeStatus consumeStatus) {
+ SpillStatus spillStatus, ConsumeStatusWithId consumeStatusWithId) {
return callWithLock(
() -> {
// TODO return iterator to avoid completely traversing the queue for each call.
@@ -216,7 +218,7 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
allBuffers.forEach(
(bufferContext -> {
if (isBufferSatisfyStatus(
- bufferContext, spillStatus, consumeStatus)) {
+ bufferContext, spillStatus, consumeStatusWithId)) {
targetBuffers.add(bufferContext.getBufferIndexAndChannel());
}
}));
@@ -460,7 +462,7 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
@GuardedBy("subpartitionLock")
private void checkAndMarkBufferReadable(HsBufferContext bufferContext) {
// only spill buffer needs to be marked as released.
- if (isBufferSatisfyStatus(bufferContext, SpillStatus.SPILL, ConsumeStatus.ALL)) {
+ if (isBufferSatisfyStatus(bufferContext, SpillStatus.SPILL, ConsumeStatusWithId.ALL_ANY)) {
bufferContext
.getSpilledFuture()
.orElseThrow(
@@ -480,7 +482,9 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
@GuardedBy("subpartitionLock")
private boolean isBufferSatisfyStatus(
- HsBufferContext bufferContext, SpillStatus spillStatus, ConsumeStatus consumeStatus) {
+ HsBufferContext bufferContext,
+ SpillStatus spillStatus,
+ ConsumeStatusWithId consumeStatusWithId) {
// released buffer is not needed.
if (bufferContext.isReleased()) {
return false;
@@ -494,12 +498,12 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
match = bufferContext.isSpillStarted();
break;
}
- switch (consumeStatus) {
+ switch (consumeStatusWithId.status) {
case NOT_CONSUMED:
- match &= !bufferContext.isConsumed();
+ match &= !bufferContext.isConsumed(consumeStatusWithId.consumerId);
break;
case CONSUMED:
- match &= bufferContext.isConsumed();
+ match &= bufferContext.isConsumed(consumeStatusWithId.consumerId);
break;
}
return match;
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContextTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContextTest.java
index 960357d4127..f0419d114df 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContextTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsBufferContextTest.java
@@ -81,25 +81,40 @@ class HsBufferContextTest {
@Test
void testBufferConsumed() {
+ final HsConsumerId consumerId = HsConsumerId.DEFAULT;
Buffer buffer = bufferContext.getBuffer();
- bufferContext.consumed();
- assertThat(bufferContext.isConsumed()).isTrue();
+ bufferContext.consumed(consumerId);
+ assertThat(bufferContext.isConsumed(consumerId)).isTrue();
assertThat(buffer.refCnt()).isEqualTo(2);
}
@Test
void testBufferConsumedRepeatedly() {
- bufferContext.consumed();
- assertThatThrownBy(() -> bufferContext.consumed())
+ final HsConsumerId consumerId = HsConsumerId.DEFAULT;
+ bufferContext.consumed(consumerId);
+ assertThatThrownBy(() -> bufferContext.consumed(consumerId))
.isInstanceOf(IllegalStateException.class)
.hasMessageContaining("Consume buffer repeatedly is unexpected.");
}
+ @Test
+ void testBufferConsumedMultipleConsumer() {
+ HsConsumerId consumer0 = HsConsumerId.newId(null);
+ HsConsumerId consumer1 = HsConsumerId.newId(consumer0);
+ bufferContext.consumed(consumer0);
+ bufferContext.consumed(consumer1);
+
+ assertThat(bufferContext.isConsumed(consumer0)).isTrue();
+ assertThat(bufferContext.isConsumed(consumer1)).isTrue();
+
+ assertThat(bufferContext.isConsumed(HsConsumerId.newId(consumer1))).isFalse();
+ }
+
@Test
void testBufferStartSpillOrConsumedAfterReleased() {
bufferContext.release();
assertThat(bufferContext.startSpilling(new CompletableFuture<>())).isFalse();
- assertThatThrownBy(() -> bufferContext.consumed())
+ assertThatThrownBy(() -> bufferContext.consumed(HsConsumerId.DEFAULT))
.isInstanceOf(IllegalStateException.class)
.hasMessageContaining("Buffer is already released.");
}
@@ -117,7 +132,7 @@ class HsBufferContextTest {
@Test
void testBufferConsumedThenRelease() {
Buffer buffer = bufferContext.getBuffer();
- bufferContext.consumed();
+ bufferContext.consumed(HsConsumerId.DEFAULT);
bufferContext.release();
assertThat(buffer.refCnt()).isEqualTo(1);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerIdTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerIdTest.java
new file mode 100644
index 00000000000..591d6b0ccfd
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsConsumerIdTest.java
@@ -0,0 +1,43 @@
+/*
+ * 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;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link HsConsumerId}. */
+class HsConsumerIdTest {
+ @Test
+ void testNewIdFromNull() {
+ HsConsumerId consumerId = HsConsumerId.newId(null);
+ assertThat(consumerId).isNotNull().isEqualTo(HsConsumerId.DEFAULT);
+ }
+
+ @Test
+ void testConsumerIdEquals() {
+ HsConsumerId consumerId = HsConsumerId.newId(null);
+ HsConsumerId consumerId1 = HsConsumerId.newId(consumerId);
+ HsConsumerId consumerId2 = HsConsumerId.newId(consumerId);
+ assertThat(consumerId1.hashCode()).isEqualTo(consumerId2.hashCode());
+ assertThat(consumerId1).isEqualTo(consumerId2);
+
+ assertThat(HsConsumerId.newId(consumerId2)).isNotEqualTo(consumerId2);
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
index ba53f2a9efb..174c021990d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
@@ -53,6 +53,8 @@ import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
+import static org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId.ALL_ANY;
+import static org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId.fromStatusAndConsumerId;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createBufferBuilder;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createTestingOutputMetrics;
import static org.assertj.core.api.Assertions.assertThat;
@@ -290,40 +292,44 @@ class HsSubpartitionMemoryDataManagerTest {
subpartitionMemoryDataManager.consumeBuffer(1);
checkBufferIndex(
- subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.ALL, ConsumeStatus.ALL),
+ subpartitionMemoryDataManager.getBuffersSatisfyStatus(SpillStatus.ALL, ALL_ANY),
Arrays.asList(0, 1, 2, 3));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.ALL, ConsumeStatus.CONSUMED),
+ SpillStatus.ALL,
+ fromStatusAndConsumerId(ConsumeStatus.CONSUMED, HsConsumerId.DEFAULT)),
Arrays.asList(0, 1));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.ALL, ConsumeStatus.NOT_CONSUMED),
+ SpillStatus.ALL,
+ fromStatusAndConsumerId(ConsumeStatus.NOT_CONSUMED, HsConsumerId.DEFAULT)),
Arrays.asList(2, 3));
checkBufferIndex(
- subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.SPILL, ConsumeStatus.ALL),
+ subpartitionMemoryDataManager.getBuffersSatisfyStatus(SpillStatus.SPILL, ALL_ANY),
Arrays.asList(1, 2));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.NOT_SPILL, ConsumeStatus.ALL),
+ SpillStatus.NOT_SPILL, ALL_ANY),
Arrays.asList(0, 3));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.SPILL, ConsumeStatus.NOT_CONSUMED),
+ SpillStatus.SPILL,
+ fromStatusAndConsumerId(ConsumeStatus.NOT_CONSUMED, HsConsumerId.DEFAULT)),
Collections.singletonList(2));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.SPILL, ConsumeStatus.CONSUMED),
+ SpillStatus.SPILL,
+ fromStatusAndConsumerId(ConsumeStatus.CONSUMED, HsConsumerId.DEFAULT)),
Collections.singletonList(1));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.NOT_SPILL, ConsumeStatus.CONSUMED),
+ SpillStatus.NOT_SPILL,
+ fromStatusAndConsumerId(ConsumeStatus.CONSUMED, HsConsumerId.DEFAULT)),
Collections.singletonList(0));
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(
- SpillStatus.NOT_SPILL, ConsumeStatus.NOT_CONSUMED),
+ SpillStatus.NOT_SPILL,
+ fromStatusAndConsumerId(ConsumeStatus.NOT_CONSUMED, HsConsumerId.DEFAULT)),
Collections.singletonList(3));
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java
index f8b303479f0..67b34d0e774 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java
@@ -79,7 +79,7 @@ public class TestingSpillingInfoProvider implements HsSpillingInfoProvider {
@Override
public Deque<BufferIndexAndChannel> getBuffersInOrder(
- int subpartitionId, SpillStatus spillStatus, ConsumeStatus consumeStatus) {
+ int subpartitionId, SpillStatus spillStatus, ConsumeStatusWithId consumeStatusWithId) {
Deque<BufferIndexAndChannel> buffersInOrder = new ArrayDeque<>();
List<BufferIndexAndChannel> subpartitionBuffers = allBuffers.get(subpartitionId);
@@ -90,7 +90,7 @@ public class TestingSpillingInfoProvider implements HsSpillingInfoProvider {
for (int i = 0; i < subpartitionBuffers.size(); i++) {
if (isBufferSatisfyStatus(
spillStatus,
- consumeStatus,
+ consumeStatusWithId,
spillBufferIndexes
.getOrDefault(subpartitionId, Collections.emptySet())
.contains(i),
@@ -124,7 +124,7 @@ public class TestingSpillingInfoProvider implements HsSpillingInfoProvider {
private static boolean isBufferSatisfyStatus(
SpillStatus spillStatus,
- ConsumeStatus consumeStatus,
+ ConsumeStatusWithId consumeStatusWithId,
boolean isSpill,
boolean isConsumed) {
boolean isNeeded = true;
@@ -136,7 +136,7 @@ public class TestingSpillingInfoProvider implements HsSpillingInfoProvider {
isNeeded = isSpill;
break;
}
- switch (consumeStatus) {
+ switch (consumeStatusWithId.status) {
case NOT_CONSUMED:
isNeeded &= !isConsumed;
break;
[flink] 04/08: [FLINK-28889] Introduce HsSubpartitionConsumerMemoryDataManager and let HsMemoryDataManager supports multiple consumer.
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit a2b402d66de9ecde75881796a6d7ee09fc4d6963
Author: Weijie Guo <re...@163.com>
AuthorDate: Mon Oct 24 18:16:01 2022 +0800
[FLINK-28889] Introduce HsSubpartitionConsumerMemoryDataManager and let HsMemoryDataManager supports multiple consumer.
---
.../io/network/partition/hybrid/HsDataView.java | 4 +-
.../partition/hybrid/HsFileDataManager.java | 3 +-
.../partition/hybrid/HsMemoryDataManager.java | 60 ++++---
.../hybrid/HsMemoryDataManagerOperation.java | 15 +-
.../partition/hybrid/HsResultPartition.java | 6 +-
.../hybrid/HsSelectiveSpillingStrategy.java | 3 +-
.../partition/hybrid/HsSpillingInfoProvider.java | 10 +-
...titionView.java => HsSubpartitionConsumer.java} | 20 ++-
... HsSubpartitionConsumerInternalOperations.java} | 6 +-
.../HsSubpartitionConsumerMemoryDataManager.java | 185 +++++++++++++++++++
.../partition/hybrid/HsSubpartitionFileReader.java | 2 +-
.../hybrid/HsSubpartitionFileReaderImpl.java | 6 +-
.../hybrid/HsSubpartitionMemoryDataManager.java | 173 ++++++------------
.../partition/hybrid/HsFileDataManagerTest.java | 10 +-
.../partition/hybrid/HsMemoryDataManagerTest.java | 20 +++
...sSubpartitionConsumerMemoryDataManagerTest.java | 197 +++++++++++++++++++++
.../hybrid/HsSubpartitionFileReaderImplTest.java | 44 ++---
.../HsSubpartitionMemoryDataManagerTest.java | 132 ++++----------
.../partition/hybrid/HsSubpartitionViewTest.java | 60 ++++---
.../hybrid/TestingMemoryDataManagerOperation.java | 26 ++-
.../hybrid/TestingSpillingInfoProvider.java | 2 +-
...tingSubpartitionConsumerInternalOperation.java} | 6 +-
22 files changed, 669 insertions(+), 321 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsDataView.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsDataView.java
index c3d4cb9ed83..b8888cfcb38 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsDataView.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsDataView.java
@@ -24,8 +24,8 @@ import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAn
import java.util.Optional;
/**
- * A view for {@link HsSubpartitionView} to find out what data exists in memory or disk and polling
- * the data.
+ * A view for {@link HsSubpartitionConsumer} to find out what data exists in memory or disk and
+ * polling the data.
*/
public interface HsDataView {
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java
index 75d322f6452..82db961be9c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java
@@ -153,7 +153,8 @@ public class HsFileDataManager implements Runnable, BufferRecycler {
/** This method only called by result partition to create subpartitionFileReader. */
public HsDataView registerNewSubpartition(
- int subpartitionId, HsSubpartitionViewInternalOperations operation) throws IOException {
+ int subpartitionId, HsSubpartitionConsumerInternalOperations operation)
+ throws IOException {
synchronized (lock) {
checkState(!isReleased, "HsFileDataManager is already released.");
lazyInitialize();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
index a8d7d89e7e4..da0c3962ad2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
@@ -24,6 +24,7 @@ import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
import org.apache.flink.runtime.io.network.buffer.BufferPool;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingStrategy.Decision;
+import org.apache.flink.util.Preconditions;
import org.apache.flink.util.concurrent.ExecutorThreadFactory;
import org.apache.flink.util.concurrent.FutureUtils;
import org.apache.flink.util.function.SupplierWithException;
@@ -35,6 +36,7 @@ import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.Path;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.Deque;
import java.util.List;
import java.util.Map;
@@ -72,8 +74,12 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
private final AtomicInteger numUnSpillBuffers = new AtomicInteger(0);
- private final Map<Integer, HsSubpartitionViewInternalOperations> subpartitionViewOperationsMap =
- new ConcurrentHashMap<>();
+ /**
+ * Each element of the list is all views of the subpartition corresponding to its index, which
+ * are stored in the form of a map that maps consumer id to its subpartition view.
+ */
+ private final List<Map<HsConsumerId, HsSubpartitionConsumerInternalOperations>>
+ subpartitionViewOperationsMap;
/**
* Currently, it is only used to regularly check the actual size of local buffer pool (the size
@@ -106,6 +112,7 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
ReadWriteLock readWriteLock = new ReentrantReadWriteLock(true);
this.lock = readWriteLock.writeLock();
+ this.subpartitionViewOperationsMap = new ArrayList<>(numSubpartitions);
for (int subpartitionId = 0; subpartitionId < numSubpartitions; ++subpartitionId) {
subpartitionMemoryDataManagers[subpartitionId] =
new HsSubpartitionMemoryDataManager(
@@ -114,6 +121,7 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
readWriteLock.readLock(),
bufferCompressor,
this);
+ subpartitionViewOperationsMap.add(new ConcurrentHashMap<>());
}
poolSize = new AtomicInteger(this.bufferPool.getNumBuffers());
@@ -156,20 +164,19 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
}
/**
- * Register {@link HsSubpartitionViewInternalOperations} to {@link
+ * Register {@link HsSubpartitionConsumerInternalOperations} to {@link
* #subpartitionViewOperationsMap}. It is used to obtain the consumption progress of the
* subpartition.
*/
- public HsDataView registerSubpartitionView(
- int subpartitionId, HsSubpartitionViewInternalOperations viewOperations) {
- HsSubpartitionViewInternalOperations oldView =
- subpartitionViewOperationsMap.put(subpartitionId, viewOperations);
- if (oldView != null) {
- LOG.debug(
- "subpartition : {} register subpartition view will replace old view. ",
- subpartitionId);
- }
- return getSubpartitionMemoryDataManager(subpartitionId);
+ public HsDataView registerNewConsumer(
+ int subpartitionId,
+ HsConsumerId consumerId,
+ HsSubpartitionConsumerInternalOperations viewOperations) {
+ HsSubpartitionConsumerInternalOperations oldView =
+ subpartitionViewOperationsMap.get(subpartitionId).put(consumerId, viewOperations);
+ Preconditions.checkState(
+ oldView == null, "Each subpartition view should have unique consumerId.");
+ return getSubpartitionMemoryDataManager(subpartitionId).registerNewConsumer(consumerId);
}
/** Close this {@link HsMemoryDataManager}, it means no data can append to memory. */
@@ -232,11 +239,11 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
// Write lock should be acquired before invoke this method.
@Override
- public List<Integer> getNextBufferIndexToConsume() {
+ public List<Integer> getNextBufferIndexToConsume(HsConsumerId consumerId) {
ArrayList<Integer> consumeIndexes = new ArrayList<>(numSubpartitions);
for (int channel = 0; channel < numSubpartitions; channel++) {
- HsSubpartitionViewInternalOperations viewOperation =
- subpartitionViewOperationsMap.get(channel);
+ HsSubpartitionConsumerInternalOperations viewOperation =
+ subpartitionViewOperationsMap.get(channel).get(consumerId);
// Access consuming offset without lock to prevent deadlock.
// A consuming thread may being blocked on the memory data manager lock, while holding
// the viewOperation lock.
@@ -280,12 +287,23 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
}
@Override
- public void onDataAvailable(int subpartitionId) {
- HsSubpartitionViewInternalOperations subpartitionViewInternalOperations =
+ public void onDataAvailable(int subpartitionId, Collection<HsConsumerId> consumerIds) {
+ Map<HsConsumerId, HsSubpartitionConsumerInternalOperations> consumerViewMap =
subpartitionViewOperationsMap.get(subpartitionId);
- if (subpartitionViewInternalOperations != null) {
- subpartitionViewInternalOperations.notifyDataAvailable();
- }
+ consumerIds.forEach(
+ consumerId -> {
+ HsSubpartitionConsumerInternalOperations consumerView =
+ consumerViewMap.get(consumerId);
+ if (consumerView != null) {
+ consumerView.notifyDataAvailable();
+ }
+ });
+ }
+
+ @Override
+ public void onConsumerReleased(int subpartitionId, HsConsumerId consumerId) {
+ subpartitionViewOperationsMap.get(subpartitionId).remove(consumerId);
+ getSubpartitionMemoryDataManager(subpartitionId).releaseConsumer(consumerId);
}
// ------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerOperation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerOperation.java
index 37df7709d63..794e95b98ea 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerOperation.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerOperation.java
@@ -20,6 +20,8 @@ package org.apache.flink.runtime.io.network.partition.hybrid;
import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import java.util.Collection;
+
/**
* This interface is used by {@link HsSubpartitionMemoryDataManager} to operate {@link
* HsMemoryDataManager}. Spilling decision may be made and handled inside these operations.
@@ -53,7 +55,16 @@ public interface HsMemoryDataManagerOperation {
/**
* This method is called when subpartition data become available.
*
- * @param subpartitionId the subpartition need notify data available.
+ * @param subpartitionId the subpartition's identifier that this consumer belongs to.
+ * @param consumerIds the consumer's identifier which need notify data available.
+ */
+ void onDataAvailable(int subpartitionId, Collection<HsConsumerId> consumerIds);
+
+ /**
+ * This method is called when consumer is decided to released.
+ *
+ * @param subpartitionId the subpartition's identifier that this consumer belongs to.
+ * @param consumerId the consumer's identifier which decided to be released.
*/
- void onDataAvailable(int subpartitionId);
+ void onConsumerReleased(int subpartitionId, HsConsumerId consumerId);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
index 9b75f1e3724..5e10b773397 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
@@ -186,13 +186,15 @@ public class HsResultPartition extends ResultPartition {
throw new PartitionNotFoundException(getPartitionId());
}
- HsSubpartitionView subpartitionView = new HsSubpartitionView(availabilityListener);
+ HsSubpartitionConsumer subpartitionView = new HsSubpartitionConsumer(availabilityListener);
HsDataView diskDataView =
fileDataManager.registerNewSubpartition(subpartitionId, subpartitionView);
HsDataView memoryDataView =
checkNotNull(memoryDataManager)
- .registerSubpartitionView(subpartitionId, subpartitionView);
+ // TODO pass real consumerId in the next commit.
+ .registerNewConsumer(
+ subpartitionId, HsConsumerId.DEFAULT, subpartitionView);
subpartitionView.setDiskDataView(diskDataView);
subpartitionView.setMemoryDataView(memoryDataView);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java
index 8c7217a4b93..56cf0a39f4a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSelectiveSpillingStrategy.java
@@ -95,7 +95,8 @@ public class HsSelectiveSpillingStrategy implements HsSpillingStrategy {
TreeMap<Integer, List<BufferIndexAndChannel>> subpartitionToHighPriorityBuffers =
getBuffersByConsumptionPriorityInOrder(
- spillingInfoProvider.getNextBufferIndexToConsume(),
+ // selective spilling strategy does not support multiple consumer.
+ spillingInfoProvider.getNextBufferIndexToConsume(HsConsumerId.DEFAULT),
subpartitionToBuffers,
spillNum);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java
index deb32e7058e..7c6b0ccd70d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSpillingInfoProvider.java
@@ -31,12 +31,14 @@ public interface HsSpillingInfoProvider {
int getNumSubpartitions();
/**
- * Get all downstream next buffer index to consume.
+ * Get all subpartition's next buffer index to consume of specific consumer.
*
- * @return A list containing all downstream next buffer index to consume, if the downstream
- * subpartition view has not been registered, the corresponding return value is -1.
+ * @param consumerId of the target downstream consumer.
+ * @return A list containing all subpartition's next buffer index to consume of specific
+ * consumer, if the downstream subpartition view has not been registered, the corresponding
+ * return value is -1.
*/
- List<Integer> getNextBufferIndexToConsume();
+ List<Integer> getNextBufferIndexToConsume(HsConsumerId consumerId);
/**
* Get all buffers with the expected status from the subpartition.
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionView.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumer.java
similarity index 94%
rename from flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionView.java
rename to flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumer.java
index bc9ec5a1744..7878d68bc33 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionView.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumer.java
@@ -33,8 +33,8 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.apache.flink.util.Preconditions.checkState;
/** The read view of HsResultPartition, data can be read from memory or disk. */
-public class HsSubpartitionView
- implements ResultSubpartitionView, HsSubpartitionViewInternalOperations {
+public class HsSubpartitionConsumer
+ implements ResultSubpartitionView, HsSubpartitionConsumerInternalOperations {
private final BufferAvailabilityListener availabilityListener;
private final Object lock = new Object();
@@ -66,7 +66,7 @@ public class HsSubpartitionView
// memoryDataView can be null only before initialization.
private HsDataView memoryDataView;
- public HsSubpartitionView(BufferAvailabilityListener availabilityListener) {
+ public HsSubpartitionConsumer(BufferAvailabilityListener availabilityListener) {
this.availabilityListener = availabilityListener;
}
@@ -270,21 +270,25 @@ public class HsSubpartitionView
}
private void releaseInternal(@Nullable Throwable throwable) {
- boolean releaseSubpartitionReader = false;
+ boolean releaseDiskView;
+ boolean releaseMemoryView;
synchronized (lock) {
if (isReleased) {
return;
}
isReleased = true;
failureCause = throwable;
- if (diskDataView != null) {
- releaseSubpartitionReader = true;
- }
+ releaseDiskView = diskDataView != null;
+ releaseMemoryView = memoryDataView != null;
}
// release subpartition reader outside of lock to avoid deadlock.
- if (releaseSubpartitionReader) {
+ if (releaseDiskView) {
//noinspection FieldAccessNotGuarded
diskDataView.releaseDataView();
}
+ if (releaseMemoryView) {
+ //noinspection FieldAccessNotGuarded
+ memoryDataView.releaseDataView();
+ }
}
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewInternalOperations.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerInternalOperations.java
similarity index 87%
rename from flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewInternalOperations.java
rename to flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerInternalOperations.java
index ab967e3d450..3f2e58507f3 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewInternalOperations.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerInternalOperations.java
@@ -19,10 +19,10 @@
package org.apache.flink.runtime.io.network.partition.hybrid;
/**
- * Operations provided by HsSubpartitionView that are used by other internal components of hybrid
- * result partition.
+ * Operations provided by {@link HsSubpartitionConsumer} that are used by other internal components
+ * of hybrid result partition.
*/
-public interface HsSubpartitionViewInternalOperations {
+public interface HsSubpartitionConsumerInternalOperations {
/** Callback for new data become available. */
void notifyDataAvailable();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManager.java
new file mode 100644
index 00000000000..606ffa32bbf
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManager.java
@@ -0,0 +1,185 @@
+/*
+ * 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;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.partition.ResultSubpartition;
+import org.apache.flink.util.function.SupplierWithException;
+
+import javax.annotation.concurrent.GuardedBy;
+
+import java.util.Deque;
+import java.util.LinkedList;
+import java.util.Optional;
+import java.util.concurrent.locks.Lock;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * This class is responsible for managing the data of a single consumer. {@link
+ * HsSubpartitionMemoryDataManager} will create a new {@link
+ * HsSubpartitionConsumerMemoryDataManager} when a consumer is registered.
+ */
+public class HsSubpartitionConsumerMemoryDataManager implements HsDataView {
+
+ @GuardedBy("consumerLock")
+ private final Deque<HsBufferContext> unConsumedBuffers = new LinkedList<>();
+
+ private final Lock consumerLock;
+
+ private final Lock resultPartitionLock;
+
+ private final HsConsumerId consumerId;
+
+ private final int subpartitionId;
+
+ private final HsMemoryDataManagerOperation memoryDataManagerOperation;
+
+ public HsSubpartitionConsumerMemoryDataManager(
+ Lock resultPartitionLock,
+ Lock consumerLock,
+ int subpartitionId,
+ HsConsumerId consumerId,
+ HsMemoryDataManagerOperation memoryDataManagerOperation) {
+ this.resultPartitionLock = resultPartitionLock;
+ this.consumerLock = consumerLock;
+ this.subpartitionId = subpartitionId;
+ this.consumerId = consumerId;
+ this.memoryDataManagerOperation = memoryDataManagerOperation;
+ }
+
+ @GuardedBy("consumerLock")
+ // this method only called from subpartitionMemoryDataManager with write lock.
+ public void addInitialBuffers(Deque<HsBufferContext> buffers) {
+ unConsumedBuffers.addAll(buffers);
+ }
+
+ @GuardedBy("consumerLock")
+ // this method only called from subpartitionMemoryDataManager with write lock.
+ public boolean addBuffer(HsBufferContext bufferContext) {
+ unConsumedBuffers.add(bufferContext);
+ trimHeadingReleasedBuffers();
+ return unConsumedBuffers.size() <= 1;
+ }
+
+ /**
+ * Check whether the head of {@link #unConsumedBuffers} is the buffer to be consumed. If so,
+ * return the buffer and backlog.
+ *
+ * @param toConsumeIndex index of buffer to be consumed.
+ * @return If the head of {@link #unConsumedBuffers} is target, return optional of the buffer
+ * and backlog. Otherwise, return {@link Optional#empty()}.
+ */
+ @SuppressWarnings("FieldAccessNotGuarded")
+ // Note that: callWithLock ensure that code block guarded by resultPartitionReadLock and
+ // subpartitionLock.
+ @Override
+ public Optional<ResultSubpartition.BufferAndBacklog> consumeBuffer(int toConsumeIndex) {
+ Optional<Tuple2<HsBufferContext, Buffer.DataType>> bufferAndNextDataType =
+ callWithLock(
+ () -> {
+ if (!checkFirstUnConsumedBufferIndex(toConsumeIndex)) {
+ return Optional.empty();
+ }
+
+ HsBufferContext bufferContext =
+ checkNotNull(unConsumedBuffers.pollFirst());
+ bufferContext.consumed(consumerId);
+ Buffer.DataType nextDataType =
+ peekNextToConsumeDataTypeInternal(toConsumeIndex + 1);
+ return Optional.of(Tuple2.of(bufferContext, nextDataType));
+ });
+
+ bufferAndNextDataType.ifPresent(
+ tuple ->
+ memoryDataManagerOperation.onBufferConsumed(
+ tuple.f0.getBufferIndexAndChannel()));
+ return bufferAndNextDataType.map(
+ tuple ->
+ new ResultSubpartition.BufferAndBacklog(
+ tuple.f0.getBuffer().readOnlySlice(),
+ getBacklog(),
+ tuple.f1,
+ toConsumeIndex));
+ }
+
+ /**
+ * Check whether the head of {@link #unConsumedBuffers} is the buffer to be consumed next time.
+ * If so, return the next buffer's data type.
+ *
+ * @param nextToConsumeIndex index of the buffer to be consumed next time.
+ * @return If the head of {@link #unConsumedBuffers} is target, return the buffer's data type.
+ * Otherwise, return {@link Buffer.DataType#NONE}.
+ */
+ @SuppressWarnings("FieldAccessNotGuarded")
+ // Note that: callWithLock ensure that code block guarded by resultPartitionReadLock and
+ // consumerLock.
+ @Override
+ public Buffer.DataType peekNextToConsumeDataType(int nextToConsumeIndex) {
+ return callWithLock(() -> peekNextToConsumeDataTypeInternal(nextToConsumeIndex));
+ }
+
+ @GuardedBy("consumerLock")
+ private Buffer.DataType peekNextToConsumeDataTypeInternal(int nextToConsumeIndex) {
+ return checkFirstUnConsumedBufferIndex(nextToConsumeIndex)
+ ? checkNotNull(unConsumedBuffers.peekFirst()).getBuffer().getDataType()
+ : Buffer.DataType.NONE;
+ }
+
+ @GuardedBy("consumerLock")
+ private boolean checkFirstUnConsumedBufferIndex(int expectedBufferIndex) {
+ trimHeadingReleasedBuffers();
+ return !unConsumedBuffers.isEmpty()
+ && unConsumedBuffers.peekFirst().getBufferIndexAndChannel().getBufferIndex()
+ == expectedBufferIndex;
+ }
+
+ @SuppressWarnings("FieldAccessNotGuarded")
+ // Un-synchronized get unConsumedBuffers size to provide memory data backlog,this will make the
+ // result greater than or equal to the actual backlog, but obtaining an accurate backlog will
+ // bring too much extra overhead.
+ @Override
+ public int getBacklog() {
+ return unConsumedBuffers.size();
+ }
+
+ @Override
+ public void releaseDataView() {
+ memoryDataManagerOperation.onConsumerReleased(subpartitionId, consumerId);
+ }
+
+ @GuardedBy("consumerLock")
+ private void trimHeadingReleasedBuffers() {
+ while (!unConsumedBuffers.isEmpty() && unConsumedBuffers.peekFirst().isReleased()) {
+ unConsumedBuffers.removeFirst();
+ }
+ }
+
+ private <R, E extends Exception> R callWithLock(SupplierWithException<R, E> callable) throws E {
+ try {
+ resultPartitionLock.lock();
+ consumerLock.lock();
+ return callable.get();
+ } finally {
+ consumerLock.unlock();
+ resultPartitionLock.unlock();
+ }
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java
index 2ec2a632281..c87bb237660 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java
@@ -58,7 +58,7 @@ public interface HsSubpartitionFileReader extends Comparable<HsSubpartitionFileR
HsSubpartitionFileReader createFileReader(
int subpartitionId,
FileChannel dataFileChannel,
- HsSubpartitionViewInternalOperations operation,
+ HsSubpartitionConsumerInternalOperations operation,
HsFileDataIndex dataIndex,
int maxBuffersReadAhead,
Consumer<HsSubpartitionFileReader> fileReaderReleaser,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java
index e6dc7122c5e..de7bd599e14 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java
@@ -57,7 +57,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
private final FileChannel dataFileChannel;
- private final HsSubpartitionViewInternalOperations operations;
+ private final HsSubpartitionConsumerInternalOperations operations;
private final CachedRegionManager cachedRegionManager;
@@ -72,7 +72,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
public HsSubpartitionFileReaderImpl(
int subpartitionId,
FileChannel dataFileChannel,
- HsSubpartitionViewInternalOperations operations,
+ HsSubpartitionConsumerInternalOperations operations,
HsFileDataIndex dataIndex,
int maxBufferReadAhead,
Consumer<HsSubpartitionFileReader> fileReaderReleaser,
@@ -487,7 +487,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
public HsSubpartitionFileReader createFileReader(
int subpartitionId,
FileChannel dataFileChannel,
- HsSubpartitionViewInternalOperations operation,
+ HsSubpartitionConsumerInternalOperations operation,
HsFileDataIndex dataIndex,
int maxBuffersReadAhead,
Consumer<HsSubpartitionFileReader> fileReaderReleaser,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
index f6804a62b1b..d01cb0d2c78 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
@@ -18,7 +18,6 @@
package org.apache.flink.runtime.io.network.partition.hybrid;
-import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.network.buffer.Buffer;
@@ -28,7 +27,6 @@ import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
-import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAndBacklog;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatus;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
@@ -40,6 +38,7 @@ import javax.annotation.concurrent.GuardedBy;
import java.nio.ByteBuffer;
import java.util.ArrayDeque;
+import java.util.ArrayList;
import java.util.Deque;
import java.util.HashMap;
import java.util.LinkedList;
@@ -49,16 +48,18 @@ import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Collectors;
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
/**
* This class is responsible for managing the data in a single subpartition. One {@link
* HsMemoryDataManager} will hold multiple {@link HsSubpartitionMemoryDataManager}.
*/
-public class HsSubpartitionMemoryDataManager implements HsDataView {
+public class HsSubpartitionMemoryDataManager {
private final int targetChannel;
private final int bufferSize;
@@ -74,9 +75,6 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
@GuardedBy("subpartitionLock")
private final Deque<HsBufferContext> allBuffers = new LinkedList<>();
- @GuardedBy("subpartitionLock")
- private final Deque<HsBufferContext> unConsumedBuffers = new LinkedList<>();
-
@GuardedBy("subpartitionLock")
private final Map<Integer, HsBufferContext> bufferIndexToContexts = new HashMap<>();
@@ -84,7 +82,10 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
private final Lock resultPartitionLock;
/** DO NOT USE DIRECTLY. Use {@link #runWithLock} or {@link #callWithLock} instead. */
- private final Object subpartitionLock = new Object();
+ private final ReentrantReadWriteLock subpartitionLock = new ReentrantReadWriteLock();
+
+ @GuardedBy("subpartitionLock")
+ private final Map<HsConsumerId, HsSubpartitionConsumerMemoryDataManager> consumerMap;
@Nullable private final BufferCompressor bufferCompressor;
@@ -101,82 +102,7 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
this.resultPartitionLock = resultPartitionLock;
this.memoryDataManagerOperation = memoryDataManagerOperation;
this.bufferCompressor = bufferCompressor;
- }
-
- // ------------------------------------------------------------------------
- // Called by Consumer
- // ------------------------------------------------------------------------
-
- /**
- * Check whether the head of {@link #unConsumedBuffers} is the buffer to be consumed next time.
- * If so, return the next buffer's data type.
- *
- * @param nextToConsumeIndex index of the buffer to be consumed next time.
- * @return If the head of {@link #unConsumedBuffers} is target, return the buffer's data type.
- * Otherwise, return {@link DataType#NONE}.
- */
- @SuppressWarnings("FieldAccessNotGuarded")
- // Note that: callWithLock ensure that code block guarded by resultPartitionReadLock and
- // subpartitionLock.
- @Override
- public DataType peekNextToConsumeDataType(int nextToConsumeIndex) {
- return callWithLock(() -> peekNextToConsumeDataTypeInternal(nextToConsumeIndex));
- }
-
- /**
- * Check whether the head of {@link #unConsumedBuffers} is the buffer to be consumed. If so,
- * return the buffer and backlog.
- *
- * @param toConsumeIndex index of buffer to be consumed.
- * @return If the head of {@link #unConsumedBuffers} is target, return optional of the buffer
- * and backlog. Otherwise, return {@link Optional#empty()}.
- */
- @SuppressWarnings("FieldAccessNotGuarded")
- // Note that: callWithLock ensure that code block guarded by resultPartitionReadLock and
- // subpartitionLock.
- @Override
- public Optional<BufferAndBacklog> consumeBuffer(int toConsumeIndex) {
- Optional<Tuple2<HsBufferContext, DataType>> bufferAndNextDataType =
- callWithLock(
- () -> {
- if (!checkFirstUnConsumedBufferIndex(toConsumeIndex)) {
- return Optional.empty();
- }
-
- HsBufferContext bufferContext =
- checkNotNull(unConsumedBuffers.pollFirst());
- // TODO move this logical to consumer and pass real consumerId.
- bufferContext.consumed(HsConsumerId.DEFAULT);
- DataType nextDataType =
- peekNextToConsumeDataTypeInternal(toConsumeIndex + 1);
- return Optional.of(Tuple2.of(bufferContext, nextDataType));
- });
-
- bufferAndNextDataType.ifPresent(
- tuple ->
- memoryDataManagerOperation.onBufferConsumed(
- tuple.f0.getBufferIndexAndChannel()));
- return bufferAndNextDataType.map(
- tuple ->
- new BufferAndBacklog(
- tuple.f0.getBuffer().readOnlySlice(),
- getBacklog(),
- tuple.f1,
- toConsumeIndex));
- }
-
- @SuppressWarnings("FieldAccessNotGuarded")
- // Un-synchronized get unConsumedBuffers size to provide memory data backlog,this will make the
- // result greater than or equal to the actual backlog, but obtaining an accurate backlog will
- // bring too much extra overhead.
- @Override
- public int getBacklog() {
- return unConsumedBuffers.size();
- }
-
- @Override
- public void releaseDataView() {
- // nothing to do for memory data.
+ this.consumerMap = new HashMap<>();
}
// ------------------------------------------------------------------------
@@ -285,6 +211,29 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
this.outputMetrics = checkNotNull(outputMetrics);
}
+ @SuppressWarnings("FieldAccessNotGuarded")
+ public HsSubpartitionConsumerMemoryDataManager registerNewConsumer(HsConsumerId consumerId) {
+ return callWithLock(
+ () -> {
+ checkState(!consumerMap.containsKey(consumerId));
+ HsSubpartitionConsumerMemoryDataManager newConsumer =
+ new HsSubpartitionConsumerMemoryDataManager(
+ resultPartitionLock,
+ subpartitionLock.readLock(),
+ targetChannel,
+ consumerId,
+ memoryDataManagerOperation);
+ newConsumer.addInitialBuffers(allBuffers);
+ consumerMap.put(consumerId, newConsumer);
+ return newConsumer;
+ });
+ }
+
+ @SuppressWarnings("FieldAccessNotGuarded")
+ public void releaseConsumer(HsConsumerId consumerId) {
+ runWithLock(() -> checkNotNull(consumerMap.remove(consumerId)));
+ }
+
// ------------------------------------------------------------------------
// Internal Methods
// ------------------------------------------------------------------------
@@ -393,36 +342,22 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
// subpartitionLock.
private void addFinishedBuffer(HsBufferContext bufferContext) {
finishedBufferIndex++;
- boolean needNotify =
- callWithLock(
- () -> {
- allBuffers.add(bufferContext);
- unConsumedBuffers.add(bufferContext);
- bufferIndexToContexts.put(
- bufferContext.getBufferIndexAndChannel().getBufferIndex(),
- bufferContext);
- trimHeadingReleasedBuffers(unConsumedBuffers);
- updateStatistics(bufferContext.getBuffer());
- return unConsumedBuffers.size() <= 1;
- });
- if (needNotify) {
- memoryDataManagerOperation.onDataAvailable(targetChannel);
- }
- }
-
- @GuardedBy("subpartitionLock")
- private DataType peekNextToConsumeDataTypeInternal(int nextToConsumeIndex) {
- return checkFirstUnConsumedBufferIndex(nextToConsumeIndex)
- ? checkNotNull(unConsumedBuffers.peekFirst()).getBuffer().getDataType()
- : DataType.NONE;
- }
-
- @GuardedBy("subpartitionLock")
- private boolean checkFirstUnConsumedBufferIndex(int expectedBufferIndex) {
- trimHeadingReleasedBuffers(unConsumedBuffers);
- return !unConsumedBuffers.isEmpty()
- && unConsumedBuffers.peekFirst().getBufferIndexAndChannel().getBufferIndex()
- == expectedBufferIndex;
+ List<HsConsumerId> needNotify = new ArrayList<>(consumerMap.size());
+ runWithLock(
+ () -> {
+ allBuffers.add(bufferContext);
+ bufferIndexToContexts.put(
+ bufferContext.getBufferIndexAndChannel().getBufferIndex(),
+ bufferContext);
+ for (Map.Entry<HsConsumerId, HsSubpartitionConsumerMemoryDataManager>
+ consumerEntry : consumerMap.entrySet()) {
+ if (consumerEntry.getValue().addBuffer(bufferContext)) {
+ needNotify.add(consumerEntry.getKey());
+ }
+ }
+ updateStatistics(bufferContext.getBuffer());
+ });
+ memoryDataManagerOperation.onDataAvailable(targetChannel, needNotify);
}
/**
@@ -517,10 +452,10 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
private <E extends Exception> void runWithLock(ThrowingRunnable<E> runnable) throws E {
try {
resultPartitionLock.lock();
- synchronized (subpartitionLock) {
- runnable.run();
- }
+ subpartitionLock.writeLock().lock();
+ runnable.run();
} finally {
+ subpartitionLock.writeLock().unlock();
resultPartitionLock.unlock();
}
}
@@ -528,10 +463,10 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
private <R, E extends Exception> R callWithLock(SupplierWithException<R, E> callable) throws E {
try {
resultPartitionLock.lock();
- synchronized (subpartitionLock) {
- return callable.get();
- }
+ subpartitionLock.writeLock().lock();
+ return callable.get();
} finally {
+ subpartitionLock.writeLock().unlock();
resultPartitionLock.unlock();
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java
index 4c168c13fd6..0980688b7f2 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java
@@ -78,7 +78,7 @@ class HsFileDataManagerTest {
private HsFileDataManager fileDataManager;
- private TestingSubpartitionViewInternalOperation subpartitionViewOperation;
+ private TestingSubpartitionConsumerInternalOperation subpartitionViewOperation;
private TestingHsSubpartitionFileReader.Factory factory;
@@ -101,7 +101,7 @@ class HsFileDataManagerTest {
HybridShuffleConfiguration.builder(
NUM_SUBPARTITIONS, bufferPool.getNumBuffersPerRequest())
.build());
- subpartitionViewOperation = new TestingSubpartitionViewInternalOperation();
+ subpartitionViewOperation = new TestingSubpartitionConsumerInternalOperation();
}
@AfterEach
@@ -352,8 +352,8 @@ class HsFileDataManagerTest {
void testConsumeWhileReleaseNoDeadlock() throws Exception {
CompletableFuture<Void> consumerStart = new CompletableFuture<>();
CompletableFuture<Void> readerFail = new CompletableFuture<>();
- HsSubpartitionView subpartitionView =
- new HsSubpartitionView(new NoOpBufferAvailablityListener());
+ HsSubpartitionConsumer subpartitionView =
+ new HsSubpartitionConsumer(new NoOpBufferAvailablityListener());
HsSubpartitionFileReaderImpl subpartitionFileReader =
new HsSubpartitionFileReaderImpl(
@@ -497,7 +497,7 @@ class HsFileDataManagerTest {
public HsSubpartitionFileReader createFileReader(
int subpartitionId,
FileChannel dataFileChannel,
- HsSubpartitionViewInternalOperations operation,
+ HsSubpartitionConsumerInternalOperations operation,
HsFileDataIndex dataIndex,
int maxBuffersReadAhead,
Consumer<HsSubpartitionFileReader> fileReaderReleaser,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerTest.java
index 773b847ecc1..e3c3663af06 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManagerTest.java
@@ -41,6 +41,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createTestingOutputMetrics;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link HsMemoryDataManager}. */
class HsMemoryDataManagerTest {
@@ -192,6 +193,25 @@ class HsMemoryDataManagerTest {
assertThat(resultPartitionReleaseFuture).isCompleted();
}
+ @Test
+ void testSubpartitionConsumerRelease() throws Exception {
+ HsSpillingStrategy spillingStrategy = TestingSpillingStrategy.builder().build();
+ HsMemoryDataManager memoryDataManager = createMemoryDataManager(spillingStrategy);
+ memoryDataManager.registerNewConsumer(
+ 0, HsConsumerId.DEFAULT, new TestingSubpartitionConsumerInternalOperation());
+ assertThatThrownBy(
+ () ->
+ memoryDataManager.registerNewConsumer(
+ 0,
+ HsConsumerId.DEFAULT,
+ new TestingSubpartitionConsumerInternalOperation()))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("Each subpartition view should have unique consumerId.");
+ memoryDataManager.onConsumerReleased(0, HsConsumerId.DEFAULT);
+ memoryDataManager.registerNewConsumer(
+ 0, HsConsumerId.DEFAULT, new TestingSubpartitionConsumerInternalOperation());
+ }
+
@Test
void testPoolSizeCheck() throws Exception {
final int requiredBuffers = 10;
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManagerTest.java
new file mode 100644
index 00000000000..9cb8d8b1751
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionConsumerMemoryDataManagerTest.java
@@ -0,0 +1,197 @@
+/*
+ * 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;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.ReadOnlySlicedNetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.ResultSubpartition;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayDeque;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.locks.ReentrantLock;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createBuffer;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link HsSubpartitionConsumerMemoryDataManager}. */
+@SuppressWarnings("FieldAccessNotGuarded")
+class HsSubpartitionConsumerMemoryDataManagerTest {
+
+ private static final int BUFFER_SIZE = Long.BYTES;
+
+ private static final int SUBPARTITION_ID = 0;
+
+ @Test
+ void testPeekNextToConsumeDataTypeNotMeetBufferIndexToConsume() throws Exception {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(memoryDataManagerOperation);
+
+ subpartitionConsumerMemoryDataManager.addBuffer(createBufferContext(0, false));
+ assertThat(subpartitionConsumerMemoryDataManager.peekNextToConsumeDataType(1))
+ .isEqualTo(Buffer.DataType.NONE);
+ }
+
+ @Test
+ void testPeekNextToConsumeDataTypeTrimHeadingReleasedBuffers() throws Exception {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(memoryDataManagerOperation);
+
+ HsBufferContext buffer1 = createBufferContext(0, false);
+ HsBufferContext buffer2 = createBufferContext(1, false);
+ subpartitionConsumerMemoryDataManager.addBuffer(buffer1);
+ subpartitionConsumerMemoryDataManager.addBuffer(buffer2);
+ subpartitionConsumerMemoryDataManager.addBuffer(createBufferContext(2, true));
+
+ buffer1.release();
+ buffer2.release();
+
+ assertThat(subpartitionConsumerMemoryDataManager.peekNextToConsumeDataType(2))
+ .isEqualTo(Buffer.DataType.EVENT_BUFFER);
+ }
+
+ @Test
+ void testConsumeBufferFirstUnConsumedBufferIndexNotMeetNextToConsume() throws Exception {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(memoryDataManagerOperation);
+
+ subpartitionConsumerMemoryDataManager.addBuffer(createBufferContext(0, false));
+ assertThat(subpartitionConsumerMemoryDataManager.consumeBuffer(1)).isNotPresent();
+ }
+
+ @Test
+ void testConsumeBufferTrimHeadingReleasedBuffers() throws Exception {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(memoryDataManagerOperation);
+
+ HsBufferContext buffer1 = createBufferContext(0, false);
+ HsBufferContext buffer2 = createBufferContext(1, false);
+ subpartitionConsumerMemoryDataManager.addBuffer(buffer1);
+ subpartitionConsumerMemoryDataManager.addBuffer(buffer2);
+ subpartitionConsumerMemoryDataManager.addBuffer(createBufferContext(2, true));
+
+ buffer1.release();
+ buffer2.release();
+
+ assertThat(subpartitionConsumerMemoryDataManager.consumeBuffer(2)).isPresent();
+ }
+
+ @Test
+ void testConsumeBufferReturnSlice() {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(memoryDataManagerOperation);
+
+ subpartitionConsumerMemoryDataManager.addBuffer(createBufferContext(0, false));
+
+ Optional<ResultSubpartition.BufferAndBacklog> bufferOpt =
+ subpartitionConsumerMemoryDataManager.consumeBuffer(0);
+ assertThat(bufferOpt)
+ .hasValueSatisfying(
+ (bufferAndBacklog ->
+ assertThat(bufferAndBacklog.buffer())
+ .isInstanceOf(ReadOnlySlicedNetworkBuffer.class)));
+ }
+
+ @Test
+ void testAddBuffer() {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(memoryDataManagerOperation);
+ ArrayDeque<HsBufferContext> initialBuffers = new ArrayDeque<>();
+ initialBuffers.add(createBufferContext(0, false));
+ initialBuffers.add(createBufferContext(1, false));
+ subpartitionConsumerMemoryDataManager.addInitialBuffers(initialBuffers);
+ subpartitionConsumerMemoryDataManager.addBuffer(createBufferContext(2, true));
+
+ assertThat(subpartitionConsumerMemoryDataManager.consumeBuffer(0))
+ .hasValueSatisfying(
+ bufferAndBacklog -> {
+ assertThat(bufferAndBacklog.getSequenceNumber()).isEqualTo(0);
+ assertThat(bufferAndBacklog.buffer().getDataType())
+ .isEqualTo(Buffer.DataType.DATA_BUFFER);
+ });
+ assertThat(subpartitionConsumerMemoryDataManager.consumeBuffer(1))
+ .hasValueSatisfying(
+ bufferAndBacklog -> {
+ assertThat(bufferAndBacklog.getSequenceNumber()).isEqualTo(1);
+ assertThat(bufferAndBacklog.buffer().getDataType())
+ .isEqualTo(Buffer.DataType.DATA_BUFFER);
+ });
+ assertThat(subpartitionConsumerMemoryDataManager.consumeBuffer(2))
+ .hasValueSatisfying(
+ bufferAndBacklog -> {
+ assertThat(bufferAndBacklog.getSequenceNumber()).isEqualTo(2);
+ assertThat(bufferAndBacklog.buffer().getDataType())
+ .isEqualTo(Buffer.DataType.EVENT_BUFFER);
+ });
+ }
+
+ @Test
+ void testRelease() {
+ CompletableFuture<HsConsumerId> consumerReleasedFuture = new CompletableFuture<>();
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder()
+ .setOnConsumerReleasedBiConsumer(
+ (subpartitionId, consumerId) -> {
+ consumerReleasedFuture.complete(consumerId);
+ })
+ .build();
+ HsConsumerId consumerId = HsConsumerId.newId(null);
+ HsSubpartitionConsumerMemoryDataManager subpartitionConsumerMemoryDataManager =
+ createSubpartitionConsumerMemoryDataManager(consumerId, memoryDataManagerOperation);
+ subpartitionConsumerMemoryDataManager.releaseDataView();
+ assertThat(consumerReleasedFuture).isCompletedWithValue(consumerId);
+ }
+
+ private static HsBufferContext createBufferContext(int bufferIndex, boolean isEvent) {
+ return new HsBufferContext(
+ createBuffer(BUFFER_SIZE, isEvent), bufferIndex, SUBPARTITION_ID);
+ }
+
+ private HsSubpartitionConsumerMemoryDataManager createSubpartitionConsumerMemoryDataManager(
+ HsMemoryDataManagerOperation memoryDataManagerOperation) {
+ return createSubpartitionConsumerMemoryDataManager(
+ HsConsumerId.DEFAULT, memoryDataManagerOperation);
+ }
+
+ private HsSubpartitionConsumerMemoryDataManager createSubpartitionConsumerMemoryDataManager(
+ HsConsumerId consumerId, HsMemoryDataManagerOperation memoryDataManagerOperation) {
+ return new HsSubpartitionConsumerMemoryDataManager(
+ new ReentrantLock(),
+ new ReentrantLock(),
+ // consumerMemoryDataManager is a member of subpartitionMemoryDataManager, using a
+ // fixed subpartition id is enough.
+ SUBPARTITION_ID,
+ consumerId,
+ memoryDataManagerOperation);
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java
index 03a560589d6..e60eff5dc09 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java
@@ -75,7 +75,7 @@ class HsSubpartitionFileReaderImplTest {
private HsFileDataIndex diskIndex;
- private TestingSubpartitionViewInternalOperation subpartitionOperation;
+ private TestingSubpartitionConsumerInternalOperation subpartitionOperation;
private FileChannel dataFileChannel;
@@ -87,7 +87,7 @@ class HsSubpartitionFileReaderImplTest {
Path dataFilePath = Files.createFile(tempPath.resolve(UUID.randomUUID().toString()));
dataFileChannel = openFileChannel(dataFilePath);
diskIndex = new HsFileDataIndexImpl(1);
- subpartitionOperation = new TestingSubpartitionViewInternalOperation();
+ subpartitionOperation = new TestingSubpartitionConsumerInternalOperation();
currentFileOffset = 0L;
}
@@ -99,10 +99,10 @@ class HsSubpartitionFileReaderImplTest {
@Test
void testReadBuffer() throws Exception {
diskIndex = new HsFileDataIndexImpl(2);
- TestingSubpartitionViewInternalOperation viewNotifier1 =
- new TestingSubpartitionViewInternalOperation();
- TestingSubpartitionViewInternalOperation viewNotifier2 =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier1 =
+ new TestingSubpartitionConsumerInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier2 =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl fileReader1 = createSubpartitionFileReader(0, viewNotifier1);
HsSubpartitionFileReaderImpl fileReader2 = createSubpartitionFileReader(1, viewNotifier2);
@@ -142,8 +142,8 @@ class HsSubpartitionFileReaderImplTest {
new BufferDecompressor(bufferSize, compressionFactoryName);
diskIndex = new HsFileDataIndexImpl(1);
- TestingSubpartitionViewInternalOperation viewNotifier =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl fileReader1 = createSubpartitionFileReader(0, viewNotifier);
writeDataToFile(0, 0, 1, 3, bufferCompressor);
@@ -362,10 +362,10 @@ class HsSubpartitionFileReaderImplTest {
@Test
void testCompareTo() throws Exception {
diskIndex = new HsFileDataIndexImpl(2);
- TestingSubpartitionViewInternalOperation viewNotifier1 =
- new TestingSubpartitionViewInternalOperation();
- TestingSubpartitionViewInternalOperation viewNotifier2 =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier1 =
+ new TestingSubpartitionConsumerInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier2 =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl fileReader1 = createSubpartitionFileReader(0, viewNotifier1);
HsSubpartitionFileReaderImpl fileReader2 = createSubpartitionFileReader(1, viewNotifier2);
assertThat(fileReader1).isEqualByComparingTo(fileReader2);
@@ -390,8 +390,8 @@ class HsSubpartitionFileReaderImplTest {
@Test
void testConsumeBuffer() throws Throwable {
- TestingSubpartitionViewInternalOperation viewNotifier =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl subpartitionFileReader =
createSubpartitionFileReader(0, viewNotifier);
@@ -428,8 +428,8 @@ class HsSubpartitionFileReaderImplTest {
@Test
void testPeekNextToConsumeDataTypeOrConsumeBufferThrowException() {
- TestingSubpartitionViewInternalOperation viewNotifier =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl subpartitionFileReader =
createSubpartitionFileReader(0, viewNotifier);
@@ -446,8 +446,8 @@ class HsSubpartitionFileReaderImplTest {
@Test
void testPeekNextToConsumeDataType() throws Throwable {
- TestingSubpartitionViewInternalOperation viewNotifier =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl subpartitionFileReader =
createSubpartitionFileReader(0, viewNotifier);
@@ -477,8 +477,8 @@ class HsSubpartitionFileReaderImplTest {
*/
@Test
void testSubpartitionReaderRegisterMultipleTimes() throws Exception {
- TestingSubpartitionViewInternalOperation viewNotifier =
- new TestingSubpartitionViewInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier =
+ new TestingSubpartitionConsumerInternalOperation();
HsSubpartitionFileReaderImpl subpartitionFileReader =
createSubpartitionFileReader(0, viewNotifier);
// mock the scenario that buffer 0 is already read form memory.
@@ -491,7 +491,7 @@ class HsSubpartitionFileReaderImplTest {
checkData(subpartitionFileReader, 2, 3);
// after failover, new view and subpartitionFileReader will be created.
- viewNotifier = new TestingSubpartitionViewInternalOperation();
+ viewNotifier = new TestingSubpartitionConsumerInternalOperation();
subpartitionFileReader = createSubpartitionFileReader(0, viewNotifier);
subpartitionFileReader.prepareForScheduling();
memorySegments = createsMemorySegments(3);
@@ -529,7 +529,7 @@ class HsSubpartitionFileReaderImplTest {
}
private HsSubpartitionFileReaderImpl createSubpartitionFileReader(
- int targetChannel, HsSubpartitionViewInternalOperations operations) {
+ int targetChannel, HsSubpartitionConsumerInternalOperations operations) {
return new HsSubpartitionFileReaderImpl(
targetChannel,
dataFileChannel,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
index 174c021990d..998c87397f1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
@@ -28,7 +28,6 @@ import org.apache.flink.runtime.io.network.buffer.Buffer.DataType;
import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
-import org.apache.flink.runtime.io.network.buffer.ReadOnlySlicedNetworkBuffer;
import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAndBacklog;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatus;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
@@ -53,11 +52,14 @@ import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
+import static org.apache.flink.runtime.io.network.partition.hybrid.HsConsumerId.DEFAULT;
import static org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId.ALL_ANY;
import static org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatusWithId.fromStatusAndConsumerId;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createBufferBuilder;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createTestingOutputMetrics;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatNoException;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link HsSubpartitionMemoryDataManager}. */
class HsSubpartitionMemoryDataManagerTest {
@@ -121,98 +123,9 @@ class HsSubpartitionMemoryDataManagerTest {
assertThat(finishedBuffers).hasValue(2);
}
- @Test
- void testPeekNextToConsumeDataTypeNotMeetBufferIndexToConsume() throws Exception {
- TestingMemoryDataManagerOperation memoryDataManagerOperation =
- TestingMemoryDataManagerOperation.builder()
- .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(RECORD_SIZE))
- .build();
- HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
- createSubpartitionMemoryDataManager(memoryDataManagerOperation);
-
- subpartitionMemoryDataManager.append(createRecord(0), DataType.DATA_BUFFER);
-
- assertThat(subpartitionMemoryDataManager.peekNextToConsumeDataType(1))
- .isEqualTo(DataType.NONE);
- }
-
- @Test
- void testPeekNextToConsumeDataTypeTrimHeadingReleasedBuffers() throws Exception {
- TestingMemoryDataManagerOperation memoryDataManagerOperation =
- TestingMemoryDataManagerOperation.builder()
- .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(RECORD_SIZE))
- .build();
- HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
- createSubpartitionMemoryDataManager(memoryDataManagerOperation);
-
- subpartitionMemoryDataManager.append(createRecord(0), DataType.DATA_BUFFER);
- subpartitionMemoryDataManager.append(createRecord(1), DataType.DATA_BUFFER);
- subpartitionMemoryDataManager.append(createRecord(2), DataType.EVENT_BUFFER);
-
- List<BufferIndexAndChannel> toRelease =
- HybridShuffleTestUtils.createBufferIndexAndChannelsList(0, 0, 1);
- subpartitionMemoryDataManager.releaseSubpartitionBuffers(toRelease);
-
- assertThat(subpartitionMemoryDataManager.peekNextToConsumeDataType(2))
- .isEqualTo(DataType.EVENT_BUFFER);
- }
-
- @Test
- void testConsumeBufferFirstUnConsumedBufferIndexNotMeetNextToConsume() throws Exception {
- TestingMemoryDataManagerOperation memoryDataManagerOperation =
- TestingMemoryDataManagerOperation.builder()
- .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(RECORD_SIZE))
- .build();
- HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
- createSubpartitionMemoryDataManager(memoryDataManagerOperation);
-
- subpartitionMemoryDataManager.append(createRecord(0), DataType.DATA_BUFFER);
-
- assertThat(subpartitionMemoryDataManager.consumeBuffer(1)).isNotPresent();
- }
-
- @Test
- void testConsumeBufferTrimHeadingReleasedBuffers() throws Exception {
- TestingMemoryDataManagerOperation memoryDataManagerOperation =
- TestingMemoryDataManagerOperation.builder()
- .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(RECORD_SIZE))
- .build();
- HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
- createSubpartitionMemoryDataManager(memoryDataManagerOperation);
-
- subpartitionMemoryDataManager.append(createRecord(0), DataType.DATA_BUFFER);
- subpartitionMemoryDataManager.append(createRecord(1), DataType.DATA_BUFFER);
- subpartitionMemoryDataManager.append(createRecord(2), DataType.EVENT_BUFFER);
-
- List<BufferIndexAndChannel> toRelease =
- HybridShuffleTestUtils.createBufferIndexAndChannelsList(0, 0, 1);
- subpartitionMemoryDataManager.releaseSubpartitionBuffers(toRelease);
-
- assertThat(subpartitionMemoryDataManager.consumeBuffer(2)).isPresent();
- }
-
- @Test
- void testConsumeBufferReturnSlice() throws Exception {
- TestingMemoryDataManagerOperation memoryDataManagerOperation =
- TestingMemoryDataManagerOperation.builder()
- .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(RECORD_SIZE))
- .build();
- HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
- createSubpartitionMemoryDataManager(memoryDataManagerOperation);
-
- subpartitionMemoryDataManager.append(createRecord(0), DataType.DATA_BUFFER);
-
- Optional<BufferAndBacklog> bufferOpt = subpartitionMemoryDataManager.consumeBuffer(0);
- assertThat(bufferOpt)
- .hasValueSatisfying(
- (bufferAndBacklog ->
- assertThat(bufferAndBacklog.buffer())
- .isInstanceOf(ReadOnlySlicedNetworkBuffer.class)));
- }
-
@ParameterizedTest
@ValueSource(strings = {"LZ4", "LZO", "ZSTD", "NULL"})
- void testConsumeBuffer(String compressionFactoryName) throws Exception {
+ void testCompressBufferAndConsume(String compressionFactoryName) throws Exception {
final int numDataBuffers = 10;
final int numRecordsPerBuffer = 10;
// write numRecordsPerBuffer long record to one buffer, as a single long is
@@ -248,9 +161,11 @@ class HsSubpartitionMemoryDataManagerTest {
subpartitionMemoryDataManager.append(createRecord(recordValue), DataType.EVENT_BUFFER);
expectedRecords.add(Tuple2.of(recordValue, DataType.EVENT_BUFFER));
+ HsSubpartitionConsumerMemoryDataManager consumer =
+ subpartitionMemoryDataManager.registerNewConsumer(DEFAULT);
ArrayList<Optional<BufferAndBacklog>> bufferAndBacklogOpts = new ArrayList<>();
for (int i = 0; i < numDataBuffers + 1; i++) {
- bufferAndBacklogOpts.add(subpartitionMemoryDataManager.consumeBuffer(i));
+ bufferAndBacklogOpts.add(consumer.consumeBuffer(i));
}
checkConsumedBufferAndNextDataType(
numRecordsPerBuffer, bufferDecompressor, expectedRecords, bufferAndBacklogOpts);
@@ -276,6 +191,8 @@ class HsSubpartitionMemoryDataManagerTest {
.build();
HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
createSubpartitionMemoryDataManager(memoryDataManagerOperation);
+ HsSubpartitionConsumerMemoryDataManager consumer =
+ subpartitionMemoryDataManager.registerNewConsumer(DEFAULT);
final int numBuffers = 4;
for (int i = 0; i < numBuffers; i++) {
subpartitionMemoryDataManager.append(createRecord(i), DataType.DATA_BUFFER);
@@ -288,8 +205,8 @@ class HsSubpartitionMemoryDataManagerTest {
subpartitionMemoryDataManager.spillSubpartitionBuffers(toStartSpilling, spilledDoneFuture);
// consume buffer 0, 1
- subpartitionMemoryDataManager.consumeBuffer(0);
- subpartitionMemoryDataManager.consumeBuffer(1);
+ consumer.consumeBuffer(0);
+ consumer.consumeBuffer(1);
checkBufferIndex(
subpartitionMemoryDataManager.getBuffersSatisfyStatus(SpillStatus.ALL, ALL_ANY),
@@ -432,6 +349,33 @@ class HsSubpartitionMemoryDataManagerTest {
assertThat(metrics.getNumBytesOut().getCount()).isEqualTo(recordSize + eventSize);
}
+ @Test
+ void testConsumerRegisterRepeatedly() {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
+ createSubpartitionMemoryDataManager(memoryDataManagerOperation);
+
+ HsConsumerId consumerId = HsConsumerId.newId(null);
+ subpartitionMemoryDataManager.registerNewConsumer(consumerId);
+ assertThatThrownBy(() -> subpartitionMemoryDataManager.registerNewConsumer(consumerId))
+ .isInstanceOf(IllegalStateException.class);
+ }
+
+ @Test
+ void testRegisterAndReleaseConsumer() {
+ TestingMemoryDataManagerOperation memoryDataManagerOperation =
+ TestingMemoryDataManagerOperation.builder().build();
+ HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
+ createSubpartitionMemoryDataManager(memoryDataManagerOperation);
+
+ HsConsumerId consumerId = HsConsumerId.newId(null);
+ subpartitionMemoryDataManager.registerNewConsumer(consumerId);
+ subpartitionMemoryDataManager.releaseConsumer(consumerId);
+ assertThatNoException()
+ .isThrownBy(() -> subpartitionMemoryDataManager.registerNewConsumer(consumerId));
+ }
+
private static void checkBufferIndex(
Deque<BufferIndexAndChannel> bufferWithIdentities, List<Integer> expectedIndexes) {
List<Integer> bufferIndexes =
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewTest.java
index 451ab8980da..11ea5c3766b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionViewTest.java
@@ -42,11 +42,11 @@ import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffle
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
-/** Tests for {@link HsSubpartitionView}. */
+/** Tests for {@link HsSubpartitionConsumer}. */
class HsSubpartitionViewTest {
@Test
void testGetNextBufferFromDisk() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
BufferAndBacklog bufferAndBacklog = createBufferAndBacklog(1, DataType.DATA_BUFFER, 0);
CompletableFuture<Void> consumeBufferFromMemoryFuture = new CompletableFuture<>();
@@ -77,7 +77,7 @@ class HsSubpartitionViewTest {
final int bufferSize = 16;
NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, bufferSize);
BufferPool bufferPool = networkBufferPool.createBufferPool(10, 10);
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
CompletableFuture<Void> acquireWriteLock = new CompletableFuture<>();
@@ -102,7 +102,8 @@ class HsSubpartitionViewTest {
} catch (Exception e) {
throw new RuntimeException(e);
}
- spillingInfoProvider.getNextBufferIndexToConsume();
+ spillingInfoProvider.getNextBufferIndexToConsume(
+ HsConsumerId.DEFAULT);
return HsSpillingStrategy.Decision.NO_ACTION;
})
.build();
@@ -117,7 +118,8 @@ class HsSubpartitionViewTest {
null,
0);
memoryDataManager.setOutputMetrics(createTestingOutputMetrics());
- HsDataView hsDataView = memoryDataManager.registerSubpartitionView(0, subpartitionView);
+ HsDataView hsDataView =
+ memoryDataManager.registerNewConsumer(0, HsConsumerId.DEFAULT, subpartitionView);
subpartitionView.setMemoryDataView(hsDataView);
subpartitionView.setDiskDataView(TestingHsDataView.NO_OP);
@@ -128,7 +130,7 @@ class HsSubpartitionViewTest {
@Test
void testGetNextBufferFromDiskNextDataTypeIsNone() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
BufferAndBacklog bufferAndBacklog = createBufferAndBacklog(0, DataType.NONE, 0);
TestingHsDataView diskDataView =
@@ -158,7 +160,7 @@ class HsSubpartitionViewTest {
@Test
void testGetNextBufferFromMemory() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
BufferAndBacklog bufferAndBacklog = createBufferAndBacklog(1, DataType.DATA_BUFFER, 0);
TestingHsDataView memoryDataView =
@@ -179,7 +181,7 @@ class HsSubpartitionViewTest {
@Test
void testGetNextBufferThrowException() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
TestingHsDataView diskDataView =
TestingHsDataView.builder()
@@ -200,7 +202,7 @@ class HsSubpartitionViewTest {
@Test
void testGetNextBufferZeroBacklog() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
final int diskBacklog = 0;
final int memoryBacklog = 10;
@@ -236,7 +238,7 @@ class HsSubpartitionViewTest {
@Test
void testNotifyDataAvailableNeedNotify() {
CompletableFuture<Void> notifyAvailableFuture = new CompletableFuture<>();
- HsSubpartitionView subpartitionView =
+ HsSubpartitionConsumer subpartitionView =
createSubpartitionView(() -> notifyAvailableFuture.complete(null));
TestingHsDataView memoryDataView =
@@ -256,7 +258,7 @@ class HsSubpartitionViewTest {
@Test
void testNotifyDataAvailableNotNeedNotify() {
CompletableFuture<Void> notifyAvailableFuture = new CompletableFuture<>();
- HsSubpartitionView subpartitionView =
+ HsSubpartitionConsumer subpartitionView =
createSubpartitionView(() -> notifyAvailableFuture.complete(null));
TestingHsDataView memoryDataView =
@@ -277,7 +279,7 @@ class HsSubpartitionViewTest {
@Test
void testGetZeroBacklogNeedNotify() {
CompletableFuture<Void> notifyAvailableFuture = new CompletableFuture<>();
- HsSubpartitionView subpartitionView =
+ HsSubpartitionConsumer subpartitionView =
createSubpartitionView(() -> notifyAvailableFuture.complete(null));
subpartitionView.setMemoryDataView(TestingHsDataView.NO_OP);
subpartitionView.setDiskDataView(
@@ -294,7 +296,7 @@ class HsSubpartitionViewTest {
@Test
void testGetAvailabilityAndBacklogPositiveCredit() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
subpartitionView.setMemoryDataView(TestingHsDataView.NO_OP);
final int backlog = 2;
@@ -311,7 +313,7 @@ class HsSubpartitionViewTest {
void testGetAvailabilityAndBacklogNonPositiveCreditNextIsData() {
final int backlog = 2;
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
subpartitionView.setMemoryDataView(
TestingHsDataView.builder()
.setConsumeBufferFunction(
@@ -336,7 +338,7 @@ class HsSubpartitionViewTest {
void testGetAvailabilityAndBacklogNonPositiveCreditNextIsEvent() {
final int backlog = 2;
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
subpartitionView.setMemoryDataView(
TestingHsDataView.builder()
.setConsumeBufferFunction(
@@ -359,23 +361,29 @@ class HsSubpartitionViewTest {
@Test
void testRelease() throws Exception {
- HsSubpartitionView subpartitionView = createSubpartitionView();
- CompletableFuture<Void> releaseDataViewFuture = new CompletableFuture<>();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
+ CompletableFuture<Void> releaseDiskViewFuture = new CompletableFuture<>();
+ CompletableFuture<Void> releaseMemoryViewFuture = new CompletableFuture<>();
TestingHsDataView diskDataView =
TestingHsDataView.builder()
- .setReleaseDataViewRunnable(() -> releaseDataViewFuture.complete(null))
+ .setReleaseDataViewRunnable(() -> releaseDiskViewFuture.complete(null))
+ .build();
+ TestingHsDataView memoryDataView =
+ TestingHsDataView.builder()
+ .setReleaseDataViewRunnable(() -> releaseMemoryViewFuture.complete(null))
.build();
subpartitionView.setDiskDataView(diskDataView);
- subpartitionView.setMemoryDataView(TestingHsDataView.NO_OP);
+ subpartitionView.setMemoryDataView(memoryDataView);
subpartitionView.releaseAllResources();
assertThat(subpartitionView.isReleased()).isTrue();
- assertThat(releaseDataViewFuture).isCompleted();
+ assertThat(releaseDiskViewFuture).isCompleted();
+ assertThat(releaseMemoryViewFuture).isCompleted();
}
@Test
void testGetConsumingOffset() {
AtomicInteger nextBufferIndex = new AtomicInteger(0);
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
TestingHsDataView diskDataView =
TestingHsDataView.builder()
.setConsumeBufferFunction(
@@ -398,7 +406,7 @@ class HsSubpartitionViewTest {
@Test
void testSetDataViewRepeatedly() {
- HsSubpartitionView subpartitionView = createSubpartitionView();
+ HsSubpartitionConsumer subpartitionView = createSubpartitionView();
subpartitionView.setMemoryDataView(TestingHsDataView.NO_OP);
assertThatThrownBy(() -> subpartitionView.setMemoryDataView(TestingHsDataView.NO_OP))
@@ -411,13 +419,13 @@ class HsSubpartitionViewTest {
.hasMessageContaining("repeatedly set disk data view is not allowed.");
}
- private static HsSubpartitionView createSubpartitionView() {
- return new HsSubpartitionView(new NoOpBufferAvailablityListener());
+ private static HsSubpartitionConsumer createSubpartitionView() {
+ return new HsSubpartitionConsumer(new NoOpBufferAvailablityListener());
}
- private static HsSubpartitionView createSubpartitionView(
+ private static HsSubpartitionConsumer createSubpartitionView(
BufferAvailabilityListener bufferAvailabilityListener) {
- return new HsSubpartitionView(bufferAvailabilityListener);
+ return new HsSubpartitionConsumer(bufferAvailabilityListener);
}
private static BufferAndBacklog createBufferAndBacklog(
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingMemoryDataManagerOperation.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingMemoryDataManagerOperation.java
index 2b16d5d4d4f..5ac37c2a1d9 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingMemoryDataManagerOperation.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingMemoryDataManagerOperation.java
@@ -21,6 +21,7 @@ package org.apache.flink.runtime.io.network.partition.hybrid;
import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
import org.apache.flink.util.function.SupplierWithException;
+import java.util.Collection;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
@@ -37,18 +38,22 @@ public class TestingMemoryDataManagerOperation implements HsMemoryDataManagerOpe
private final Runnable onDataAvailableRunnable;
+ private final BiConsumer<Integer, HsConsumerId> onConsumerReleasedBiConsumer;
+
private TestingMemoryDataManagerOperation(
SupplierWithException<BufferBuilder, InterruptedException>
requestBufferFromPoolSupplier,
BiConsumer<Integer, Integer> markBufferReadableConsumer,
Consumer<BufferIndexAndChannel> onBufferConsumedConsumer,
Runnable onBufferFinishedRunnable,
- Runnable onDataAvailableRunnable) {
+ Runnable onDataAvailableRunnable,
+ BiConsumer<Integer, HsConsumerId> onConsumerReleasedBiConsumer) {
this.requestBufferFromPoolSupplier = requestBufferFromPoolSupplier;
this.markBufferReadableConsumer = markBufferReadableConsumer;
this.onBufferConsumedConsumer = onBufferConsumedConsumer;
this.onBufferFinishedRunnable = onBufferFinishedRunnable;
this.onDataAvailableRunnable = onDataAvailableRunnable;
+ this.onConsumerReleasedBiConsumer = onConsumerReleasedBiConsumer;
}
@Override
@@ -72,10 +77,15 @@ public class TestingMemoryDataManagerOperation implements HsMemoryDataManagerOpe
}
@Override
- public void onDataAvailable(int subpartitionId) {
+ public void onDataAvailable(int subpartitionId, Collection<HsConsumerId> consumerIds) {
onDataAvailableRunnable.run();
}
+ @Override
+ public void onConsumerReleased(int subpartitionId, HsConsumerId consumerId) {
+ onConsumerReleasedBiConsumer.accept(subpartitionId, consumerId);
+ }
+
public static Builder builder() {
return new Builder();
}
@@ -93,6 +103,9 @@ public class TestingMemoryDataManagerOperation implements HsMemoryDataManagerOpe
private Runnable onDataAvailableRunnable = () -> {};
+ private BiConsumer<Integer, HsConsumerId> onConsumerReleasedBiConsumer =
+ (ignore1, ignore2) -> {};
+
public Builder setRequestBufferFromPoolSupplier(
SupplierWithException<BufferBuilder, InterruptedException>
requestBufferFromPoolSupplier) {
@@ -122,6 +135,12 @@ public class TestingMemoryDataManagerOperation implements HsMemoryDataManagerOpe
return this;
}
+ public Builder setOnConsumerReleasedBiConsumer(
+ BiConsumer<Integer, HsConsumerId> onConsumerReleasedBiConsumer) {
+ this.onConsumerReleasedBiConsumer = onConsumerReleasedBiConsumer;
+ return this;
+ }
+
private Builder() {}
public TestingMemoryDataManagerOperation build() {
@@ -130,7 +149,8 @@ public class TestingMemoryDataManagerOperation implements HsMemoryDataManagerOpe
markBufferReadableConsumer,
onBufferConsumedConsumer,
onBufferFinishedRunnable,
- onDataAvailableRunnable);
+ onDataAvailableRunnable,
+ onConsumerReleasedBiConsumer);
}
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java
index 67b34d0e774..9980f4176fa 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSpillingInfoProvider.java
@@ -73,7 +73,7 @@ public class TestingSpillingInfoProvider implements HsSpillingInfoProvider {
}
@Override
- public List<Integer> getNextBufferIndexToConsume() {
+ public List<Integer> getNextBufferIndexToConsume(HsConsumerId consumerId) {
return getNextBufferIndexToConsumeSupplier.get();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSubpartitionViewInternalOperation.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSubpartitionConsumerInternalOperation.java
similarity index 88%
rename from flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSubpartitionViewInternalOperation.java
rename to flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSubpartitionConsumerInternalOperation.java
index ff2f8250b2a..1f32189da53 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSubpartitionViewInternalOperation.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/TestingSubpartitionConsumerInternalOperation.java
@@ -18,9 +18,9 @@
package org.apache.flink.runtime.io.network.partition.hybrid;
-/** Mock {@link HsSubpartitionViewInternalOperations} for test. */
-public class TestingSubpartitionViewInternalOperation
- implements HsSubpartitionViewInternalOperations {
+/** Mock {@link HsSubpartitionConsumerInternalOperations} for test. */
+public class TestingSubpartitionConsumerInternalOperation
+ implements HsSubpartitionConsumerInternalOperations {
// -1 indicates downstream just start consuming offset.
private int consumingOffset = -1;
[flink] 08/08: [FLINK-28889] HsResultPartition support broadcast optimize
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit d11940c4a78c71548b5a06af50da2e5f9cb68918
Author: Weijie Guo <re...@163.com>
AuthorDate: Mon Sep 26 16:46:34 2022 +0800
[FLINK-28889] HsResultPartition support broadcast optimize
This closes #21122
---
.../ResultPartitionDeploymentDescriptor.java | 5 ++
.../network/partition/ResultPartitionFactory.java | 12 +++
.../partition/hybrid/HsResultPartition.java | 22 +++++-
.../network/partition/ResultPartitionBuilder.java | 8 ++
.../partition/ResultPartitionFactoryTest.java | 18 +++++
.../partition/hybrid/HsResultPartitionTest.java | 86 ++++++++++++++++++++++
.../shuffle/PartitionDescriptorBuilder.java | 9 ++-
7 files changed, 155 insertions(+), 5 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java
index 3bf2c802817..f510bb08acb 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java
@@ -63,6 +63,11 @@ public class ResultPartitionDeploymentDescriptor implements Serializable {
return partitionDescriptor.getPartitionId();
}
+ /** Whether the resultPartition is a broadcast edge. */
+ public boolean isBroadcast() {
+ return partitionDescriptor.isBroadcast();
+ }
+
public ResultPartitionType getPartitionType() {
return partitionDescriptor.getPartitionType();
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactory.java
index a95f5e76770..7858f4c6709 100755
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactory.java
@@ -126,6 +126,7 @@ public class ResultPartitionFactory {
desc.getPartitionType(),
desc.getNumberOfSubpartitions(),
desc.getMaxParallelism(),
+ desc.isBroadcast(),
createBufferPoolFactory(desc.getNumberOfSubpartitions(), desc.getPartitionType()));
}
@@ -137,6 +138,7 @@ public class ResultPartitionFactory {
ResultPartitionType type,
int numberOfSubpartitions,
int maxParallelism,
+ boolean isBroadcast,
SupplierWithException<BufferPool, IOException> bufferPoolFactory) {
BufferCompressor bufferCompressor = null;
if (type.supportCompression() && batchShuffleCompressionEnabled) {
@@ -216,6 +218,15 @@ public class ResultPartitionFactory {
}
} else if (type == ResultPartitionType.HYBRID_FULL
|| type == ResultPartitionType.HYBRID_SELECTIVE) {
+ if (type == ResultPartitionType.HYBRID_SELECTIVE && isBroadcast) {
+ // for broadcast result partition, it can be optimized to always use full spilling
+ // strategy to significantly reduce shuffle data writing cost.
+ LOG.info(
+ "{} result partition has been replaced by {} result partition to reduce shuffle data writing cost.",
+ type,
+ ResultPartitionType.HYBRID_FULL);
+ type = ResultPartitionType.HYBRID_FULL;
+ }
partition =
new HsResultPartition(
taskNameWithSubtaskAndId,
@@ -240,6 +251,7 @@ public class ResultPartitionFactory {
.SpillingStrategyType.SELECTIVE)
.build(),
bufferCompressor,
+ isBroadcast,
bufferPoolFactory);
} else {
throw new IllegalArgumentException("Unrecognized ResultPartitionType: " + type);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
index 26aee9450ad..d4f2f632aff 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
@@ -60,6 +60,8 @@ import static org.apache.flink.util.Preconditions.checkState;
public class HsResultPartition extends ResultPartition {
public static final String DATA_FILE_SUFFIX = ".hybrid.data";
+ public static final int BROADCAST_CHANNEL = 0;
+
private final HsFileDataIndex dataIndex;
private final HsFileDataManager fileDataManager;
@@ -77,6 +79,9 @@ public class HsResultPartition extends ResultPartition {
@Nullable private HsMemoryDataManager memoryDataManager;
+ /** Whether this result partition broadcasts all data and event. */
+ private final boolean isBroadcastOnly;
+
public HsResultPartition(
String owningTaskName,
int partitionIndex,
@@ -91,6 +96,7 @@ public class HsResultPartition extends ResultPartition {
int networkBufferSize,
HybridShuffleConfiguration hybridShuffleConfiguration,
@Nullable BufferCompressor bufferCompressor,
+ boolean isBroadcastOnly,
SupplierWithException<BufferPool, IOException> bufferPoolFactory) {
super(
owningTaskName,
@@ -103,9 +109,10 @@ public class HsResultPartition extends ResultPartition {
bufferCompressor,
bufferPoolFactory);
this.networkBufferSize = networkBufferSize;
- this.dataIndex = new HsFileDataIndexImpl(numSubpartitions);
+ this.dataIndex = new HsFileDataIndexImpl(isBroadcastOnly ? 1 : numSubpartitions);
this.dataFilePath = new File(dataFileBashPath + DATA_FILE_SUFFIX).toPath();
this.hybridShuffleConfiguration = hybridShuffleConfiguration;
+ this.isBroadcastOnly = isBroadcastOnly;
this.fileDataManager =
new HsFileDataManager(
readBufferPool,
@@ -126,7 +133,7 @@ public class HsResultPartition extends ResultPartition {
this.fileDataManager.setup();
this.memoryDataManager =
new HsMemoryDataManager(
- numSubpartitions,
+ isBroadcastOnly ? 1 : numSubpartitions,
networkBufferSize,
bufferPool,
getSpillingStrategy(hybridShuffleConfiguration),
@@ -167,8 +174,12 @@ public class HsResultPartition extends ResultPartition {
private void broadcast(ByteBuffer record, Buffer.DataType dataType) throws IOException {
numBytesProduced.inc(record.remaining());
- for (int i = 0; i < numSubpartitions; i++) {
- emit(record.duplicate(), i, dataType);
+ if (isBroadcastOnly) {
+ emit(record, BROADCAST_CHANNEL, dataType);
+ } else {
+ for (int i = 0; i < numSubpartitions; i++) {
+ emit(record.duplicate(), i, dataType);
+ }
}
}
@@ -190,6 +201,9 @@ public class HsResultPartition extends ResultPartition {
if (!Files.isReadable(dataFilePath)) {
throw new PartitionNotFoundException(getPartitionId());
}
+ // if broadcastOptimize is enabled, map every subpartitionId to the special broadcast
+ // channel.
+ subpartitionId = isBroadcastOnly ? BROADCAST_CHANNEL : subpartitionId;
HsSubpartitionConsumer subpartitionConsumer =
new HsSubpartitionConsumer(availabilityListener);
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionBuilder.java
index 5ef48002b03..7ad6cd08b84 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionBuilder.java
@@ -47,6 +47,8 @@ public class ResultPartitionBuilder {
private int numTargetKeyGroups = 1;
+ private boolean isBroadcast = false;
+
private ResultPartitionManager partitionManager = new ResultPartitionManager();
private FileChannelManager channelManager = NoOpFileChannelManager.INSTANCE;
@@ -211,6 +213,11 @@ public class ResultPartitionBuilder {
return this;
}
+ public ResultPartitionBuilder setBroadcast(boolean broadcast) {
+ isBroadcast = broadcast;
+ return this;
+ }
+
public ResultPartition build() {
ResultPartitionFactory resultPartitionFactory =
new ResultPartitionFactory(
@@ -244,6 +251,7 @@ public class ResultPartitionBuilder {
partitionType,
numberOfSubpartitions,
numTargetKeyGroups,
+ isBroadcast,
factory);
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java
index 62420fabdcd..79a0aa89463 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java
@@ -132,12 +132,29 @@ class ResultPartitionFactoryTest {
assertThat(resultPartition.isReleased()).isFalse();
}
+ @Test
+ public void testHybridBroadcastEdgeAlwaysUseFullResultPartition() {
+ ResultPartition resultPartition =
+ createResultPartition(
+ ResultPartitionType.HYBRID_SELECTIVE, Integer.MAX_VALUE, true);
+ assertThat(resultPartition.partitionType).isEqualTo(ResultPartitionType.HYBRID_FULL);
+
+ resultPartition =
+ createResultPartition(ResultPartitionType.HYBRID_FULL, Integer.MAX_VALUE, true);
+ assertThat(resultPartition.partitionType).isEqualTo(ResultPartitionType.HYBRID_FULL);
+ }
+
private static ResultPartition createResultPartition(ResultPartitionType partitionType) {
return createResultPartition(partitionType, Integer.MAX_VALUE);
}
private static ResultPartition createResultPartition(
ResultPartitionType partitionType, int sortShuffleMinParallelism) {
+ return createResultPartition(partitionType, sortShuffleMinParallelism, false);
+ }
+
+ private static ResultPartition createResultPartition(
+ ResultPartitionType partitionType, int sortShuffleMinParallelism, boolean isBroadcast) {
final ResultPartitionManager manager = new ResultPartitionManager();
final ResultPartitionFactory factory =
@@ -163,6 +180,7 @@ class ResultPartitionFactoryTest {
new ResultPartitionDeploymentDescriptor(
PartitionDescriptorBuilder.newBuilder()
.setPartitionType(partitionType)
+ .setIsBroadcast(isBroadcast)
.build(),
NettyShuffleDescriptorBuilder.newBuilder().buildLocal(),
1);
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
index 296834db22b..4d5fb55c203 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartitionTest.java
@@ -279,6 +279,60 @@ class HsResultPartitionTest {
}
}
+ @Test
+ void testBroadcastResultPartition() throws Exception {
+ final int numBuffers = 10;
+ final int numRecords = 10;
+ final int numConsumers = 2;
+ final Random random = new Random();
+
+ BufferPool bufferPool = globalPool.createBufferPool(numBuffers, numBuffers);
+ try (HsResultPartition resultPartition = createHsResultPartition(2, bufferPool, true)) {
+ List<ByteBuffer> dataWritten = new ArrayList<>();
+ for (int i = 0; i < numRecords; i++) {
+ ByteBuffer record = generateRandomData(bufferSize, random);
+ resultPartition.broadcastRecord(record);
+ dataWritten.add(record);
+ }
+ resultPartition.finish();
+
+ Tuple2[] viewAndListeners = createSubpartitionViews(resultPartition, 2);
+
+ List<List<Buffer>> dataRead = new ArrayList<>();
+ for (int i = 0; i < numConsumers; i++) {
+ dataRead.add(new ArrayList<>());
+ }
+ readData(
+ viewAndListeners,
+ (buffer, subpartition) -> {
+ int numBytes = buffer.readableBytes();
+ if (buffer.isBuffer()) {
+ MemorySegment segment =
+ MemorySegmentFactory.allocateUnpooledSegment(numBytes);
+ segment.put(0, buffer.getNioBufferReadable(), numBytes);
+ dataRead.get(subpartition)
+ .add(
+ new NetworkBuffer(
+ segment,
+ (buf) -> {},
+ buffer.getDataType(),
+ numBytes));
+ }
+ });
+
+ for (int i = 0; i < numConsumers; i++) {
+ assertThat(dataWritten).hasSameSizeAs(dataRead.get(i));
+ List<Buffer> readBufferList = dataRead.get(i);
+ for (int j = 0; j < dataWritten.size(); j++) {
+ ByteBuffer bufferWritten = dataWritten.get(j);
+ bufferWritten.rewind();
+ Buffer bufferRead = readBufferList.get(j);
+ assertThat(bufferRead.getNioBufferReadable()).isEqualTo(bufferWritten);
+ }
+ }
+ }
+ }
+
@Test
void testClose() throws Exception {
final int numBuffers = 1;
@@ -427,6 +481,20 @@ class HsResultPartitionTest {
}
}
+ @Test
+ void testMetricsUpdateForBroadcastOnlyResultPartition() throws Exception {
+ BufferPool bufferPool = globalPool.createBufferPool(3, 3);
+ try (HsResultPartition partition = createHsResultPartition(2, bufferPool, true)) {
+ partition.broadcastRecord(ByteBuffer.allocate(bufferSize));
+ assertThat(taskIOMetricGroup.getNumBuffersOutCounter().getCount()).isEqualTo(1);
+ assertThat(taskIOMetricGroup.getNumBytesOutCounter().getCount()).isEqualTo(bufferSize);
+ IOMetrics ioMetrics = taskIOMetricGroup.createSnapshot();
+ assertThat(ioMetrics.getNumBytesProducedOfPartitions())
+ .hasSize(1)
+ .containsValue((long) bufferSize);
+ }
+ }
+
private static void recordDataWritten(
ByteBuffer record,
Queue<Tuple2<ByteBuffer, Buffer.DataType>>[] dataWritten,
@@ -493,9 +561,25 @@ class HsResultPartitionTest {
private HsResultPartition createHsResultPartition(int numSubpartitions, BufferPool bufferPool)
throws IOException {
+ return createHsResultPartition(numSubpartitions, bufferPool, false);
+ }
+
+ private HsResultPartition createHsResultPartition(
+ int numSubpartitions,
+ BufferPool bufferPool,
+ HybridShuffleConfiguration hybridShuffleConfiguration)
+ throws IOException {
+ return createHsResultPartition(
+ numSubpartitions, bufferPool, false, hybridShuffleConfiguration);
+ }
+
+ private HsResultPartition createHsResultPartition(
+ int numSubpartitions, BufferPool bufferPool, boolean isBroadcastOnly)
+ throws IOException {
return createHsResultPartition(
numSubpartitions,
bufferPool,
+ isBroadcastOnly,
HybridShuffleConfiguration.builder(
numSubpartitions, readBufferPool.getNumBuffersPerRequest())
.build());
@@ -504,6 +588,7 @@ class HsResultPartitionTest {
private HsResultPartition createHsResultPartition(
int numSubpartitions,
BufferPool bufferPool,
+ boolean isBroadcastOnly,
HybridShuffleConfiguration hybridShuffleConfiguration)
throws IOException {
HsResultPartition hsResultPartition =
@@ -521,6 +606,7 @@ class HsResultPartitionTest {
bufferSize,
hybridShuffleConfiguration,
null,
+ isBroadcastOnly,
() -> bufferPool);
taskIOMetricGroup =
UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup();
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/PartitionDescriptorBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/PartitionDescriptorBuilder.java
index de85cd10adc..d2305eb1b31 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/PartitionDescriptorBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/PartitionDescriptorBuilder.java
@@ -28,6 +28,8 @@ public class PartitionDescriptorBuilder {
private ResultPartitionType partitionType;
private int totalNumberOfPartitions = 1;
+ private boolean isBroadcast = false;
+
private PartitionDescriptorBuilder() {
this.partitionId = new IntermediateResultPartitionID();
this.partitionType = ResultPartitionType.PIPELINED;
@@ -48,6 +50,11 @@ public class PartitionDescriptorBuilder {
return this;
}
+ public PartitionDescriptorBuilder setIsBroadcast(boolean isBroadcast) {
+ this.isBroadcast = isBroadcast;
+ return this;
+ }
+
public PartitionDescriptor build() {
return new PartitionDescriptor(
new IntermediateDataSetID(),
@@ -56,7 +63,7 @@ public class PartitionDescriptorBuilder {
partitionType,
1,
0,
- false,
+ isBroadcast,
true);
}
[flink] 05/08: [FLINK-28889] HsFileDataManager supports multiple consumer.
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit f95c12e9a4a71930fb20feda611bd3c1584380ce
Author: Weijie Guo <re...@163.com>
AuthorDate: Wed Oct 19 15:20:02 2022 +0800
[FLINK-28889] HsFileDataManager supports multiple consumer.
---
.../partition/hybrid/HsFileDataManager.java | 7 ++--
.../partition/hybrid/HsResultPartition.java | 4 ++-
.../partition/hybrid/HsSubpartitionFileReader.java | 1 +
.../hybrid/HsSubpartitionFileReaderImpl.java | 10 ++++--
.../partition/hybrid/HsFileDataManagerTest.java | 24 +++++++------
.../hybrid/HsSubpartitionFileReaderImplTest.java | 40 ++++++++++++++++++++++
6 files changed, 71 insertions(+), 15 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java
index 82db961be9c..a54456cabcc 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManager.java
@@ -152,8 +152,10 @@ public class HsFileDataManager implements Runnable, BufferRecycler {
}
/** This method only called by result partition to create subpartitionFileReader. */
- public HsDataView registerNewSubpartition(
- int subpartitionId, HsSubpartitionConsumerInternalOperations operation)
+ public HsDataView registerNewConsumer(
+ int subpartitionId,
+ HsConsumerId consumerId,
+ HsSubpartitionConsumerInternalOperations operation)
throws IOException {
synchronized (lock) {
checkState(!isReleased, "HsFileDataManager is already released.");
@@ -162,6 +164,7 @@ public class HsFileDataManager implements Runnable, BufferRecycler {
HsSubpartitionFileReader subpartitionReader =
fileReaderFactory.createFileReader(
subpartitionId,
+ consumerId,
dataFileChannel,
operation,
dataIndex,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
index 5e10b773397..678fb697b29 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsResultPartition.java
@@ -188,7 +188,9 @@ public class HsResultPartition extends ResultPartition {
HsSubpartitionConsumer subpartitionView = new HsSubpartitionConsumer(availabilityListener);
HsDataView diskDataView =
- fileDataManager.registerNewSubpartition(subpartitionId, subpartitionView);
+ // TODO pass real consumerId in the next commit.
+ fileDataManager.registerNewConsumer(
+ subpartitionId, HsConsumerId.DEFAULT, subpartitionView);
HsDataView memoryDataView =
checkNotNull(memoryDataManager)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java
index c87bb237660..5b4bc0158bd 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReader.java
@@ -57,6 +57,7 @@ public interface HsSubpartitionFileReader extends Comparable<HsSubpartitionFileR
interface Factory {
HsSubpartitionFileReader createFileReader(
int subpartitionId,
+ HsConsumerId consumerId,
FileChannel dataFileChannel,
HsSubpartitionConsumerInternalOperations operation,
HsFileDataIndex dataIndex,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java
index de7bd599e14..e40e917cfa1 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImpl.java
@@ -55,6 +55,8 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
private final int subpartitionId;
+ private final HsConsumerId consumerId;
+
private final FileChannel dataFileChannel;
private final HsSubpartitionConsumerInternalOperations operations;
@@ -71,6 +73,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
public HsSubpartitionFileReaderImpl(
int subpartitionId,
+ HsConsumerId consumerId,
FileChannel dataFileChannel,
HsSubpartitionConsumerInternalOperations operations,
HsFileDataIndex dataIndex,
@@ -78,6 +81,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
Consumer<HsSubpartitionFileReader> fileReaderReleaser,
ByteBuffer headerBuf) {
this.subpartitionId = subpartitionId;
+ this.consumerId = consumerId;
this.dataFileChannel = dataFileChannel;
this.operations = operations;
this.headerBuf = headerBuf;
@@ -95,12 +99,12 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
return false;
}
HsSubpartitionFileReaderImpl that = (HsSubpartitionFileReaderImpl) o;
- return subpartitionId == that.subpartitionId;
+ return subpartitionId == that.subpartitionId && Objects.equals(consumerId, that.consumerId);
}
@Override
public int hashCode() {
- return Objects.hash(subpartitionId);
+ return Objects.hash(subpartitionId, consumerId);
}
/**
@@ -486,6 +490,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
@Override
public HsSubpartitionFileReader createFileReader(
int subpartitionId,
+ HsConsumerId consumerId,
FileChannel dataFileChannel,
HsSubpartitionConsumerInternalOperations operation,
HsFileDataIndex dataIndex,
@@ -494,6 +499,7 @@ public class HsSubpartitionFileReaderImpl implements HsSubpartitionFileReader {
ByteBuffer headerBuffer) {
return new HsSubpartitionFileReaderImpl(
subpartitionId,
+ consumerId,
dataFileChannel,
operation,
dataIndex,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java
index 0980688b7f2..401f95ff823 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsFileDataManagerTest.java
@@ -52,6 +52,7 @@ import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;
+import static org.apache.flink.runtime.io.network.partition.hybrid.HsConsumerId.DEFAULT;
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.assertj.core.api.Assertions.assertThat;
@@ -123,7 +124,7 @@ class HsFileDataManagerTest {
assertThat(reader.readBuffers).isEmpty();
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
@@ -142,7 +143,7 @@ class HsFileDataManagerTest {
factory.allReaders.add(reader);
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
@@ -174,7 +175,7 @@ class HsFileDataManagerTest {
});
factory.allReaders.add(reader);
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
@@ -207,8 +208,8 @@ class HsFileDataManagerTest {
factory.allReaders.add(reader1);
factory.allReaders.add(reader2);
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
- fileDataManager.registerNewSubpartition(1, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(1, DEFAULT, subpartitionViewOperation);
// trigger run.
ioExecutor.trigger();
@@ -243,7 +244,7 @@ class HsFileDataManagerTest {
reader.setFailConsumer((cause::complete));
factory.allReaders.add(reader);
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
@@ -269,7 +270,7 @@ class HsFileDataManagerTest {
});
factory.allReaders.add(reader);
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
@@ -314,7 +315,7 @@ class HsFileDataManagerTest {
};
releaseThread.start();
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
@@ -335,7 +336,8 @@ class HsFileDataManagerTest {
fileDataManager.release();
assertThatThrownBy(
() -> {
- fileDataManager.registerNewSubpartition(0, subpartitionViewOperation);
+ fileDataManager.registerNewConsumer(
+ 0, DEFAULT, subpartitionViewOperation);
ioExecutor.trigger();
})
.isInstanceOf(IllegalStateException.class)
@@ -358,6 +360,7 @@ class HsFileDataManagerTest {
HsSubpartitionFileReaderImpl subpartitionFileReader =
new HsSubpartitionFileReaderImpl(
0,
+ DEFAULT,
dataFileChannel,
subpartitionView,
new HsFileDataIndexImpl(NUM_SUBPARTITIONS),
@@ -376,7 +379,7 @@ class HsFileDataManagerTest {
}
};
factory.allReaders.add(subpartitionFileReader);
- HsDataView diskDataView = fileDataManager.registerNewSubpartition(0, subpartitionView);
+ HsDataView diskDataView = fileDataManager.registerNewConsumer(0, DEFAULT, subpartitionView);
subpartitionView.setDiskDataView(diskDataView);
TestingHsDataView memoryDataView =
TestingHsDataView.builder()
@@ -496,6 +499,7 @@ class HsFileDataManagerTest {
@Override
public HsSubpartitionFileReader createFileReader(
int subpartitionId,
+ HsConsumerId consumerId,
FileChannel dataFileChannel,
HsSubpartitionConsumerInternalOperations operation,
HsFileDataIndex dataIndex,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java
index e60eff5dc09..cccf9f5d93e 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionFileReaderImplTest.java
@@ -501,6 +501,38 @@ class HsSubpartitionFileReaderImplTest {
checkData(subpartitionFileReader, 1, 2, 3);
}
+ @Test
+ void testMultipleFileReaderOfSingleSubpartition() throws Exception {
+ TestingSubpartitionConsumerInternalOperation viewNotifier1 =
+ new TestingSubpartitionConsumerInternalOperation();
+ TestingSubpartitionConsumerInternalOperation viewNotifier2 =
+ new TestingSubpartitionConsumerInternalOperation();
+
+ HsConsumerId consumer0 = HsConsumerId.newId(null);
+ HsSubpartitionFileReaderImpl consumer1 =
+ createSubpartitionFileReader(0, consumer0, viewNotifier1);
+ HsSubpartitionFileReaderImpl consumer2 =
+ createSubpartitionFileReader(0, HsConsumerId.newId(consumer0), viewNotifier2);
+
+ assertThat(consumer1).isNotEqualTo(consumer2);
+
+ // write data to a single subpartition, then read these buffers by two consumers
+ // respectively.
+ writeDataToFile(0, 0, 1, 3);
+
+ consumer1.prepareForScheduling();
+ Queue<MemorySegment> memorySegments = createsMemorySegments(3);
+ consumer1.readBuffers(memorySegments, (ignore) -> {});
+ assertThat(memorySegments).isEmpty();
+ checkData(consumer1, 1, 2, 3);
+
+ consumer2.prepareForScheduling();
+ memorySegments = createsMemorySegments(3);
+ consumer2.readBuffers(memorySegments, (ignore) -> {});
+ assertThat(memorySegments).isEmpty();
+ checkData(consumer2, 1, 2, 3);
+ }
+
private static void checkData(
HsSubpartitionFileReaderImpl fileReader,
BufferDecompressor bufferDecompressor,
@@ -530,8 +562,16 @@ class HsSubpartitionFileReaderImplTest {
private HsSubpartitionFileReaderImpl createSubpartitionFileReader(
int targetChannel, HsSubpartitionConsumerInternalOperations operations) {
+ return createSubpartitionFileReader(targetChannel, HsConsumerId.DEFAULT, operations);
+ }
+
+ private HsSubpartitionFileReaderImpl createSubpartitionFileReader(
+ int targetChannel,
+ HsConsumerId consumerId,
+ HsSubpartitionConsumerInternalOperations operations) {
return new HsSubpartitionFileReaderImpl(
targetChannel,
+ consumerId,
dataFileChannel,
operations,
diskIndex,
[flink] 01/08: [hotfix] Migrate ResultPartitionFactoryTest to JUnit5 and AssertJ.
Posted by xt...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 840a5198df7c548ab92b050282a3b09bfb3b08d3
Author: Weijie Guo <re...@163.com>
AuthorDate: Wed Oct 26 12:26:58 2022 +0800
[hotfix] Migrate ResultPartitionFactoryTest to JUnit5 and AssertJ.
---
.../partition/ResultPartitionFactoryTest.java | 69 +++++++++++-----------
1 file changed, 33 insertions(+), 36 deletions(-)
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java
index a7d1b04dacb..62420fabdcd 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java
@@ -26,113 +26,110 @@ import org.apache.flink.runtime.io.network.partition.hybrid.HsResultPartition;
import org.apache.flink.runtime.shuffle.PartitionDescriptorBuilder;
import org.apache.flink.runtime.util.EnvironmentInformation;
import org.apache.flink.runtime.util.NettyShuffleDescriptorBuilder;
-import org.apache.flink.util.TestLogger;
+import org.apache.flink.util.TestLoggerExtension;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
-import org.junit.Test;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
-import java.util.Arrays;
import java.util.concurrent.Executors;
-import static org.hamcrest.CoreMatchers.instanceOf;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.assertj.core.api.Assertions.assertThat;
/** Tests for the {@link ResultPartitionFactory}. */
@SuppressWarnings("StaticVariableUsedBeforeInitialization")
-public class ResultPartitionFactoryTest extends TestLogger {
+@ExtendWith(TestLoggerExtension.class)
+class ResultPartitionFactoryTest {
private static final String tempDir = EnvironmentInformation.getTemporaryFileDirectory();
private static final int SEGMENT_SIZE = 64;
private static FileChannelManager fileChannelManager;
- @BeforeClass
- public static void setUp() {
+ @BeforeAll
+ static void setUp() {
fileChannelManager = new FileChannelManagerImpl(new String[] {tempDir}, "testing");
}
- @AfterClass
- public static void shutdown() throws Exception {
+ @AfterAll
+ static void shutdown() throws Exception {
fileChannelManager.close();
}
@Test
- public void testBoundedBlockingSubpartitionsCreated() {
+ void testBoundedBlockingSubpartitionsCreated() {
final BoundedBlockingResultPartition resultPartition =
(BoundedBlockingResultPartition)
createResultPartition(ResultPartitionType.BLOCKING);
- Arrays.stream(resultPartition.subpartitions)
- .forEach(sp -> assertThat(sp, instanceOf(BoundedBlockingSubpartition.class)));
+ assertThat(resultPartition.subpartitions)
+ .allSatisfy((sp) -> assertThat(sp).isInstanceOf(BoundedBlockingSubpartition.class));
}
@Test
- public void testPipelinedSubpartitionsCreated() {
+ void testPipelinedSubpartitionsCreated() {
final PipelinedResultPartition resultPartition =
(PipelinedResultPartition) createResultPartition(ResultPartitionType.PIPELINED);
- Arrays.stream(resultPartition.subpartitions)
- .forEach(sp -> assertThat(sp, instanceOf(PipelinedSubpartition.class)));
+ assertThat(resultPartition.subpartitions)
+ .allSatisfy((sp) -> assertThat(sp).isInstanceOf(PipelinedSubpartition.class));
}
@Test
- public void testSortMergePartitionCreated() {
+ void testSortMergePartitionCreated() {
ResultPartition resultPartition = createResultPartition(ResultPartitionType.BLOCKING, 1);
- assertTrue(resultPartition instanceof SortMergeResultPartition);
+ assertThat(resultPartition).isInstanceOf(SortMergeResultPartition.class);
}
@Test
- public void testHybridFullResultPartitionCreated() {
+ void testHybridFullResultPartitionCreated() {
ResultPartition resultPartition = createResultPartition(ResultPartitionType.HYBRID_FULL);
- assertTrue(resultPartition instanceof HsResultPartition);
+ assertThat(resultPartition).isInstanceOf(HsResultPartition.class);
}
@Test
- public void testHybridSelectiveResultPartitionCreated() {
+ void testHybridSelectiveResultPartitionCreated() {
ResultPartition resultPartition =
createResultPartition(ResultPartitionType.HYBRID_SELECTIVE);
- assertTrue(resultPartition instanceof HsResultPartition);
+ assertThat(resultPartition).isInstanceOf(HsResultPartition.class);
}
@Test
- public void testNoReleaseOnConsumptionForBoundedBlockingPartition() {
+ void testNoReleaseOnConsumptionForBoundedBlockingPartition() {
final ResultPartition resultPartition = createResultPartition(ResultPartitionType.BLOCKING);
resultPartition.onConsumedSubpartition(0);
- assertFalse(resultPartition.isReleased());
+ assertThat(resultPartition.isReleased()).isFalse();
}
@Test
- public void testNoReleaseOnConsumptionForSortMergePartition() {
+ void testNoReleaseOnConsumptionForSortMergePartition() {
final ResultPartition resultPartition =
createResultPartition(ResultPartitionType.BLOCKING, 1);
resultPartition.onConsumedSubpartition(0);
- assertFalse(resultPartition.isReleased());
+ assertThat(resultPartition.isReleased()).isFalse();
}
@Test
- public void testNoReleaseOnConsumptionForHybridFullPartition() {
+ void testNoReleaseOnConsumptionForHybridFullPartition() {
final ResultPartition resultPartition =
createResultPartition(ResultPartitionType.HYBRID_FULL);
resultPartition.onConsumedSubpartition(0);
- assertFalse(resultPartition.isReleased());
+ assertThat(resultPartition.isReleased()).isFalse();
}
@Test
- public void testNoReleaseOnConsumptionForHybridSelectivePartition() {
+ void testNoReleaseOnConsumptionForHybridSelectivePartition() {
final ResultPartition resultPartition =
createResultPartition(ResultPartitionType.HYBRID_SELECTIVE);
resultPartition.onConsumedSubpartition(0);
- assertFalse(resultPartition.isReleased());
+ assertThat(resultPartition.isReleased()).isFalse();
}
private static ResultPartition createResultPartition(ResultPartitionType partitionType) {
@@ -171,7 +168,7 @@ public class ResultPartitionFactoryTest extends TestLogger {
1);
// guard our test assumptions
- assertEquals(1, descriptor.getNumberOfSubpartitions());
+ assertThat(descriptor.getNumberOfSubpartitions()).isEqualTo(1);
final ResultPartition partition = factory.create("test", 0, descriptor);
manager.registerResultPartition(partition);