[FLINK-36934][network] Enrich numSubpartitions to TierMasterAgent#addPartitionAndGetShuffleDescriptor

Author: reswqa

flink-runtime/src/main/java/org/apache/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) {

flink-runtime/src/main/java/org/apache/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());
     }
 

flink-runtime/src/main/java/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;
     }

flink-runtime/src/main/java/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.

flink-runtime/src/main/java/org/apache/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,

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 =

flink-runtime/src/test/java/org/apache/flink/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);