[FLINK-36934][network] Enrich numSubpartitions to TierMasterAgent#add…

…PartitionAndGetShuffleDescriptor
liuyongvs · Dec 20, 2024 · 1523f2c · 1523f2c
1 parent a0d2d87
commit 1523f2c
Show file tree

Hide file tree

Showing 7 changed files with 11 additions and 9 deletions.
diff --git a/...flink/runtime/io/network/partition/hybrid/tiered/shuffle/TieredInternalShuffleMaster.java b/...flink/runtime/io/network/partition/hybrid/tiered/shuffle/TieredInternalShuffleMaster.java
@@ -81,9 +81,9 @@ public void unregisterJob(JobID jobID) {
     }
 
     public List<TierShuffleDescriptor> addPartitionAndGetShuffleDescriptor(
-            JobID jobID, ResultPartitionID resultPartitionID) {
+            JobID jobID, int numSubpartitions, ResultPartitionID resultPartitionID) {
         return tieredStorageMasterClient.addPartitionAndGetShuffleDescriptor(
-                jobID, resultPartitionID);
+                jobID, numSubpartitions, resultPartitionID);
     }
 
     public void releasePartition(ShuffleDescriptor shuffleDescriptor) {

diff --git a/...e/flink/runtime/io/network/partition/hybrid/tiered/storage/TieredStorageMasterClient.java b/...e/flink/runtime/io/network/partition/hybrid/tiered/storage/TieredStorageMasterClient.java
@@ -49,12 +49,12 @@ public void unregisterJob(JobID jobID) {
     }
 
     public List<TierShuffleDescriptor> addPartitionAndGetShuffleDescriptor(
-            JobID jobID, ResultPartitionID resultPartitionID) {
+            JobID jobID, int numSubpartitions, ResultPartitionID resultPartitionID) {
         return tiers.stream()
                 .map(
                         tierMasterAgent ->
                                 tierMasterAgent.addPartitionAndGetShuffleDescriptor(
-                                        jobID, resultPartitionID))
+                                        jobID, numSubpartitions, resultPartitionID))
                 .collect(Collectors.toList());
     }
 

diff --git a/...ava/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/NoOpMasterAgent.java b/...ava/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/NoOpMasterAgent.java
@@ -38,7 +38,7 @@ public void unregisterJob(JobID jobID) {
 
     @Override
     public TierShuffleDescriptor addPartitionAndGetShuffleDescriptor(
-            JobID jobID, ResultPartitionID resultPartitionID) {
+            JobID jobID, int numSubpartitions, ResultPartitionID resultPartitionID) {
         // noop
         return NoOpTierShuffleDescriptor.INSTANCE;
     }

diff --git a/...ava/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/TierMasterAgent.java b/...ava/org/apache/flink/runtime/io/network/partition/hybrid/tiered/tier/TierMasterAgent.java
@@ -32,7 +32,7 @@ public interface TierMasterAgent {
 
     /** Add a new tiered storage partition and get the {@link TierShuffleDescriptor}. */
     TierShuffleDescriptor addPartitionAndGetShuffleDescriptor(
-            JobID jobID, ResultPartitionID resultPartitionID);
+            JobID jobID, int numSubpartitions, ResultPartitionID resultPartitionID);
 
     /**
      * Release a tiered storage partition.

diff --git a/...e/flink/runtime/io/network/partition/hybrid/tiered/tier/remote/RemoteTierMasterAgent.java b/...e/flink/runtime/io/network/partition/hybrid/tiered/tier/remote/RemoteTierMasterAgent.java
@@ -56,7 +56,7 @@ public void unregisterJob(JobID jobID) {
 
     @Override
     public TierShuffleDescriptor addPartitionAndGetShuffleDescriptor(
-            JobID jobID, ResultPartitionID resultPartitionID) {
+            JobID jobID, int numSubpartitions, ResultPartitionID resultPartitionID) {
         TieredStoragePartitionId partitionId = convertId(resultPartitionID);
         resourceRegistry.registerResource(
                 partitionId,

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/NettyShuffleMaster.java b/flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/NettyShuffleMaster.java
@@ -108,7 +108,9 @@ public CompletableFuture<NettyShuffleDescriptor> registerPartitionWithProducer(
         if (tieredInternalShuffleMaster != null) {
             tierShuffleDescriptors =
                     tieredInternalShuffleMaster.addPartitionAndGetShuffleDescriptor(
-                            jobID, resultPartitionID);
+                            jobID,
+                            partitionDescriptor.getNumberOfSubpartitions(),
+                            resultPartitionID);
         }
 
         NettyShuffleDescriptor shuffleDeploymentDescriptor =

diff --git a/...ink/runtime/io/network/partition/hybrid/tiered/tier/remote/RemoteTierMasterAgentTest.java b/...ink/runtime/io/network/partition/hybrid/tiered/tier/remote/RemoteTierMasterAgentTest.java
@@ -52,7 +52,7 @@ void testAddAndReleasePartition() throws IOException {
         RemoteTierMasterAgent masterAgent =
                 new RemoteTierMasterAgent(tempFolder.getAbsolutePath(), resourceRegistry);
         TierShuffleDescriptor tierShuffleDescriptor =
-                masterAgent.addPartitionAndGetShuffleDescriptor(new JobID(), resultPartitionID);
+                masterAgent.addPartitionAndGetShuffleDescriptor(new JobID(), 1, resultPartitionID);
         assertThat(partitionFile.exists()).isTrue();
         masterAgent.releasePartition(tierShuffleDescriptor);