[FLINK-33985][runtime] Support obtain all partitions existing in cluster through ShuffleMaster.

This closes apache#24553.

Author: JunRuiLee

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NettyShuffleEnvironment.java

@@ -45,6 +45,7 @@
 import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
 import org.apache.flink.runtime.shuffle.ShuffleEnvironment;
 import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
+import org.apache.flink.runtime.shuffle.ShuffleMetrics;
 import org.apache.flink.runtime.taskmanager.NettyShuffleEnvironmentConfiguration;
 import org.apache.flink.util.Preconditions;
 
@@ -198,6 +199,12 @@ public Collection<ResultPartitionID> getPartitionsOccupyingLocalResources() {
         return resultPartitionManager.getUnreleasedPartitions();
     }
 
+    @Override
+    public Optional<ShuffleMetrics> getMetricsIfPartitionOccupyingLocalResource(
+            ResultPartitionID partitionId) {
+        return resultPartitionManager.getMetricsOfPartition(partitionId);
+    }
+
     // --------------------------------------------------------------------------------------------
     //  Create Output Writers and Input Readers
     // --------------------------------------------------------------------------------------------

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/ResultPartition.java

@@ -215,6 +215,10 @@ public ResultPartitionType getPartitionType() {
         return partitionType;
     }
 
+    public ResultPartitionBytesCounter getResultPartitionBytes() {
+        return resultPartitionBytes;
+    }
+
     // ------------------------------------------------------------------------
 
     @Override

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/ResultPartitionManager.java

@@ -19,6 +19,8 @@
 package org.apache.flink.runtime.io.network.partition;
 
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.runtime.shuffle.DefaultShuffleMetrics;
+import org.apache.flink.runtime.shuffle.ShuffleMetrics;
 import org.apache.flink.util.CollectionUtil;
 import org.apache.flink.util.concurrent.ScheduledExecutor;
 
@@ -293,4 +295,18 @@ public Collection<ResultPartitionID> getUnreleasedPartitions() {
             return registeredPartitions.keySet();
         }
     }
+
+    public Optional<ShuffleMetrics> getMetricsOfPartition(ResultPartitionID partitionId) {
+        synchronized (registeredPartitions) {
+            final ResultPartition partition = registeredPartitions.get(partitionId);
+
+            if (partition == null) {
+                return Optional.empty();
+            }
+
+            return Optional.of(
+                    new DefaultShuffleMetrics(
+                            partition.getResultPartitionBytes().createSnapshot()));
+        }
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/TaskExecutorPartitionTracker.java

@@ -44,6 +44,14 @@ public interface TaskExecutorPartitionTracker
      */
     void stopTrackingAndReleaseJobPartitionsFor(JobID producingJobId);
 
+    /**
+     * Get all partitions tracked for the given job.
+     *
+     * @param producingJobId the job id
+     * @return the tracked partitions
+     */
+    Collection<TaskExecutorPartitionInfo> getTrackedPartitionsFor(JobID producingJobId);
+
     /** Promotes the given partitions. */
     void promoteJobPartitions(Collection<ResultPartitionID> partitionsToPromote);
 

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/TaskExecutorPartitionTrackerImpl.java

@@ -35,6 +35,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static java.util.stream.Collectors.toList;
+
 /**
  * Utility for tracking partitions and issuing release calls to task executors and shuffle masters.
  */
@@ -83,6 +85,19 @@ public void stopTrackingAndReleaseJobPartitionsFor(JobID producingJobId) {
         shuffleEnvironment.releasePartitionsLocally(partitionsForJob);
     }
 
+    @Override
+    public Collection<TaskExecutorPartitionInfo> getTrackedPartitionsFor(JobID producingJobId) {
+        return partitionTable.getTrackedPartitions(producingJobId).stream()
+                .map(
+                        partitionId -> {
+                            final PartitionInfo<JobID, TaskExecutorPartitionInfo> partitionInfo =
+                                    partitionInfos.get(partitionId);
+                            Preconditions.checkNotNull(partitionInfo);
+                            return partitionInfo.getMetaInfo();
+                        })
+                .collect(toList());
+    }
+
     @Override
     public void promoteJobPartitions(Collection<ResultPartitionID> partitionsToPromote) {
         LOG.debug("Promoting Job Partitions {}", partitionsToPromote);

flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java

@@ -90,6 +90,7 @@
 import org.apache.flink.runtime.scheduler.SchedulerNG;
 import org.apache.flink.runtime.shuffle.JobShuffleContext;
 import org.apache.flink.runtime.shuffle.JobShuffleContextImpl;
+import org.apache.flink.runtime.shuffle.PartitionWithMetrics;
 import org.apache.flink.runtime.shuffle.ShuffleMaster;
 import org.apache.flink.runtime.slots.ResourceRequirement;
 import org.apache.flink.runtime.state.KeyGroupRange;
@@ -115,9 +116,11 @@
 
 import java.io.IOException;
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -957,6 +960,27 @@ public CompletableFuture<?> stopTrackingAndReleasePartitions(
         return future;
     }
 
+    @Override
+    public CompletableFuture<Collection<PartitionWithMetrics>>
+            getAllPartitionWithMetricsOnTaskManagers() {
+        final List<CompletableFuture<Collection<PartitionWithMetrics>>> allFutures =
+                new ArrayList<>();
+        registeredTaskManagers
+                .values()
+                .forEach(
+                        taskManager ->
+                                allFutures.add(
+                                        taskManager
+                                                .getTaskExecutorGateway()
+                                                .getPartitionWithMetrics(jobGraph.getJobID())));
+        return FutureUtils.combineAll(allFutures)
+                .thenApply(
+                        partitions ->
+                                partitions.stream()
+                                        .flatMap(Collection::stream)
+                                        .collect(Collectors.toList()));
+    }
+
     @Override
     public CompletableFuture<Acknowledge> notifyNewBlockedNodes(Collection<BlockedNode> newNodes) {
         blocklistHandler.addNewBlockedNodes(newNodes);

flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMasterGateway.java

@@ -46,6 +46,7 @@
 import org.apache.flink.runtime.rpc.FencedRpcGateway;
 import org.apache.flink.runtime.rpc.RpcTimeout;
 import org.apache.flink.runtime.scheduler.ExecutionGraphInfo;
+import org.apache.flink.runtime.shuffle.PartitionWithMetrics;
 import org.apache.flink.runtime.slots.ResourceRequirement;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorToJobManagerHeartbeatPayload;
 import org.apache.flink.runtime.taskexecutor.slot.SlotOffer;
@@ -55,6 +56,7 @@
 import javax.annotation.Nullable;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
 
 /** {@link JobMaster} rpc gateway interface. */
@@ -301,6 +303,11 @@ CompletableFuture<CoordinationResponse> deliverCoordinationRequestToCoordinator(
     CompletableFuture<?> stopTrackingAndReleasePartitions(
             Collection<ResultPartitionID> partitionIds);
 
+    default CompletableFuture<Collection<PartitionWithMetrics>>
+            getAllPartitionWithMetricsOnTaskManagers() {
+        return CompletableFuture.completedFuture(Collections.emptyList());
+    }
+
     /**
      * Read current {@link JobResourceRequirements job resource requirements}.
      *

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/DefaultPartitionWithMetrics.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.shuffle;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/** Default {@link PartitionWithMetrics} implementation. */
+public class DefaultPartitionWithMetrics implements PartitionWithMetrics {
+    private final ShuffleDescriptor shuffleDescriptor;
+    private final ShuffleMetrics partitionMetrics;
+
+    public DefaultPartitionWithMetrics(
+            ShuffleDescriptor shuffleDescriptor, ShuffleMetrics partitionMetrics) {
+        this.shuffleDescriptor = checkNotNull(shuffleDescriptor);
+        this.partitionMetrics = checkNotNull(partitionMetrics);
+    }
+
+    @Override
+    public ShuffleMetrics getPartitionMetrics() {
+        return partitionMetrics;
+    }
+
+    @Override
+    public ShuffleDescriptor getPartition() {
+        return shuffleDescriptor;
+    }
+}

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/DefaultShuffleMetrics.java

@@ -0,0 +1,37 @@
+/*
+ * 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.shuffle;
+
+import org.apache.flink.runtime.executiongraph.ResultPartitionBytes;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/** Default {@link ShuffleMetrics} implementation. */
+public class DefaultShuffleMetrics implements ShuffleMetrics {
+    private final ResultPartitionBytes partitionBytes;
+
+    public DefaultShuffleMetrics(ResultPartitionBytes partitionBytes) {
+        this.partitionBytes = checkNotNull(partitionBytes);
+    }
+
+    @Override
+    public ResultPartitionBytes getPartitionBytes() {
+        return partitionBytes;
+    }
+}

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/JobShuffleContext.java

@@ -39,4 +39,10 @@ public interface JobShuffleContext {
      */
     CompletableFuture<?> stopTrackingAndReleasePartitions(
             Collection<ResultPartitionID> partitionIds);
+
+    /**
+     * Retrieves a collection containing descriptions and metrics of existing result partitions from
+     * all TaskManagers.
+     */
+    CompletableFuture<Collection<PartitionWithMetrics>> getAllPartitionWithMetricsOnTaskManagers();
 }

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/JobShuffleContextImpl.java

@@ -49,4 +49,10 @@ public CompletableFuture<?> stopTrackingAndReleasePartitions(
             Collection<ResultPartitionID> partitionIds) {
         return jobMasterGateway.stopTrackingAndReleasePartitions(partitionIds);
     }
+
+    @Override
+    public CompletableFuture<Collection<PartitionWithMetrics>>
+            getAllPartitionWithMetricsOnTaskManagers() {
+        return jobMasterGateway.getAllPartitionWithMetricsOnTaskManagers();
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/NettyShuffleMaster.java

@@ -31,6 +31,9 @@
 
 import javax.annotation.Nullable;
 
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -58,6 +61,8 @@ public class NettyShuffleMaster implements ShuffleMaster<NettyShuffleDescriptor>
 
     @Nullable private final TieredInternalShuffleMaster tieredInternalShuffleMaster;
 
+    private final Map<JobID, JobShuffleContext> jobShuffleContexts = new HashMap<>();
+
     public NettyShuffleMaster(Configuration conf) {
         checkNotNull(conf);
         buffersPerInputChannel =
@@ -165,4 +170,21 @@ private boolean isHybridShuffleNewModeEnabled(Configuration conf) {
                         || conf.get(BATCH_SHUFFLE_MODE) == ALL_EXCHANGES_HYBRID_SELECTIVE)
                 && conf.get(NETWORK_HYBRID_SHUFFLE_ENABLE_NEW_MODE);
     }
+
+    @Override
+    public CompletableFuture<Collection<PartitionWithMetrics>> getAllPartitionWithMetrics(
+            JobID jobId) {
+        return checkNotNull(jobShuffleContexts.get(jobId))
+                .getAllPartitionWithMetricsOnTaskManagers();
+    }
+
+    @Override
+    public void registerJob(JobShuffleContext context) {
+        jobShuffleContexts.put(context.getJobId(), context);
+    }
+
+    @Override
+    public void unregisterJob(JobID jobId) {
+        jobShuffleContexts.remove(jobId);
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/PartitionWithMetrics.java

@@ -0,0 +1,28 @@
+/*
+ * 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.shuffle;
+
+import java.io.Serializable;
+
+/** Interface representing the description and metrics of a result partition. */
+public interface PartitionWithMetrics extends Serializable {
+    ShuffleMetrics getPartitionMetrics();
+
+    ShuffleDescriptor getPartition();
+}

flink-runtime/src/main/java/org/apache/flink/runtime/shuffle/ShuffleEnvironment.java

@@ -33,6 +33,7 @@
 import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
+import java.util.Optional;
 
 /**
  * Interface for the implementation of shuffle service local environment.
@@ -162,6 +163,18 @@ List<P> createResultPartitionWriters(
      */
     Collection<ResultPartitionID> getPartitionsOccupyingLocalResources();
 
+    /**
+     * Get metrics of the partition if it still occupies some resources locally and have not been
+     * released yet.
+     *
+     * @param partitionId the partition id
+     * @return An Optional of {@link ShuffleMetrics}, if found, of the given partition
+     */
+    default Optional<ShuffleMetrics> getMetricsIfPartitionOccupyingLocalResource(
+            ResultPartitionID partitionId) {
+        return Optional.empty();
+    }
+
     /**
      * Factory method for the {@link InputGate InputGates} to consume result partitions.
      *