@Override
    public CompletableFuture<Collection<PartitionWithMetrics>> getPartitionWithMetrics(
            Duration timeout, Set<ResultPartitionID> expectedPartitions) {
        this.partitionsToFetch = expectedPartitions;
        this.fetchPartitionsFuture = new CompletableFuture<>();

        // check already fetched partitions
        checkPartitionOnTaskManagerReportFinished();

        return FutureUtils.orTimeout(
                        fetchPartitionsFuture, timeout.toMillis(), TimeUnit.MILLISECONDS, null)
                .handleAsync(
                        (metrics, throwable) -> {
                            stopFetchAndRetainPartitionWithMetricsOnTaskManager();

                            if (throwable != null) {
                                if (throwable instanceof TimeoutException) {
                                    log.warn(
                                            "Timeout occurred after {} ms "
                                                    + "while fetching partition(s) ({}) from task managers.",
                                            timeout.toMillis(),
                                            expectedPartitions);

                                    return new ArrayList<>(fetchedPartitionsWithMetrics.values());
                                }
                                throw new CompletionException(throwable);
                            }

                            return new ArrayList<>(fetchedPartitionsWithMetrics.values());
                        },
                        getMainThreadExecutor());
    }
@Test
    void testGetPartitionWithMetrics() throws Exception {
        JobVertex jobVertex = new JobVertex("jobVertex");
        jobVertex.setInvokableClass(NoOpInvokable.class);
        jobVertex.setParallelism(1);
        final JobGraph jobGraph = JobGraphTestUtils.batchJobGraph(jobVertex);

        try (final JobMaster jobMaster =
                new JobMasterBuilder(jobGraph, rpcService)
                        .withConfiguration(configuration)
                        .withHighAvailabilityServices(haServices)
                        .withHeartbeatServices(heartbeatServices)
                        .withBlocklistHandlerFactory(
                                new DefaultBlocklistHandler.Factory(Duration.ofMillis((100L))))
                        .createJobMaster()) {

            jobMaster.start();

            final JobMasterGateway jobMasterGateway =
                    jobMaster.getSelfGateway(JobMasterGateway.class);

            DefaultShuffleMetrics shuffleMetrics1 =
                    new DefaultShuffleMetrics(new ResultPartitionBytes(new long[] {1, 2, 3}));
            NettyShuffleDescriptor descriptor1 =
                    NettyShuffleDescriptorBuilder.newBuilder().buildLocal();
            Collection<PartitionWithMetrics> defaultPartitionWithMetrics1 =
                    Collections.singletonList(
                            new DefaultPartitionWithMetrics(descriptor1, shuffleMetrics1));

            DefaultShuffleMetrics shuffleMetrics2 =
                    new DefaultShuffleMetrics(new ResultPartitionBytes(new long[] {4, 5, 6}));
            NettyShuffleDescriptor descriptor2 =
                    NettyShuffleDescriptorBuilder.newBuilder().buildLocal();
            Collection<PartitionWithMetrics> defaultPartitionWithMetrics2 =
                    Collections.singletonList(
                            new DefaultPartitionWithMetrics(descriptor2, shuffleMetrics2));

            DefaultShuffleMetrics shuffleMetrics3 =
                    new DefaultShuffleMetrics(new ResultPartitionBytes(new long[] {7, 8, 9}));
            NettyShuffleDescriptor descriptor3 =
                    NettyShuffleDescriptorBuilder.newBuilder().buildLocal();
            Collection<PartitionWithMetrics> defaultPartitionWithMetrics3 =
                    Collections.singletonList(
                            new DefaultPartitionWithMetrics(descriptor3, shuffleMetrics3));

            DefaultShuffleMetrics shuffleMetrics4 =
                    new DefaultShuffleMetrics(new ResultPartitionBytes(new long[] {10, 11}));
            NettyShuffleDescriptor descriptor4 =
                    NettyShuffleDescriptorBuilder.newBuilder().buildLocal();
            Collection<PartitionWithMetrics> defaultPartitionWithMetrics4 =
                    Collections.singletonList(
                            new DefaultPartitionWithMetrics(descriptor4, shuffleMetrics4));

            // start fetch and retain partitions and then register tm1
            final TestingTaskExecutorGateway taskExecutorGateway1 =
                    new TestingTaskExecutorGatewayBuilder()
                            .setRequestPartitionWithMetricsFunction(
                                    ignored ->
                                            CompletableFuture.completedFuture(
                                                    defaultPartitionWithMetrics1))
                            .setAddress("tm1")
                            .createTestingTaskExecutorGateway();

            registerSlotsAtJobMaster(
                    1,
                    jobMasterGateway,
                    jobGraph.getJobID(),
                    taskExecutorGateway1,
                    new LocalUnresolvedTaskManagerLocation());

            jobMaster.startFetchAndRetainPartitionWithMetricsOnTaskManager();

            verifyPartitionMetrics(
                    jobMaster.getPartitionWithMetricsOnTaskManagers(),
                    defaultPartitionWithMetrics1);

            // register tm2
            TestingTaskExecutorGateway taskExecutorGateway2 =
                    new TestingTaskExecutorGatewayBuilder()
                            .setRequestPartitionWithMetricsFunction(
                                    ignored ->
                                            CompletableFuture.completedFuture(
                                                    defaultPartitionWithMetrics2))
                            .setAddress("tm2")
                            .createTestingTaskExecutorGateway();
            registerSlotsAtJobMaster(
                    1,
                    jobMasterGateway,
                    jobGraph.getJobID(),
                    taskExecutorGateway2,
                    new LocalUnresolvedTaskManagerLocation());

            Collection<PartitionWithMetrics> expectedMetrics =
                    new ArrayList<>(defaultPartitionWithMetrics1);
            expectedMetrics.addAll(defaultPartitionWithMetrics2);
            verifyPartitionMetrics(
                    jobMaster.getPartitionWithMetricsOnTaskManagers(), expectedMetrics);

            // register tm3 which received fetch request but not response on time.
            CompletableFuture<Tuple2<JobID, Set<ResultPartitionID>>> releaseRequest =
                    new CompletableFuture<>();
            CompletableFuture<Collection<PartitionWithMetrics>> fetchPartitionsFuture =
                    new CompletableFuture<>();
            TestingTaskExecutorGateway taskExecutorGateway3 =
                    new TestingTaskExecutorGatewayBuilder()
                            .setRequestPartitionWithMetricsFunction(
                                    ignored -> fetchPartitionsFuture)
                            .setReleasePartitionsConsumer(
                                    (id, partitions) ->
                                            releaseRequest.complete(Tuple2.of(id, partitions)))
                            .setAddress("tm3")
                            .createTestingTaskExecutorGateway();

            registerSlotsAtJobMaster(
                    1,
                    jobMasterGateway,
                    jobGraph.getJobID(),
                    taskExecutorGateway3,
                    new LocalUnresolvedTaskManagerLocation());

            // register tm4 which is not included in expected partitions
            TestingTaskExecutorGateway taskExecutorGateway4 =
                    new TestingTaskExecutorGatewayBuilder()
                            .setRequestPartitionWithMetricsFunction(
                                    ignored ->
                                            CompletableFuture.completedFuture(
                                                    defaultPartitionWithMetrics4))
                            .setAddress("tm4")
                            .createTestingTaskExecutorGateway();
            registerSlotsAtJobMaster(
                    1,
                    jobMasterGateway,
                    jobGraph.getJobID(),
                    taskExecutorGateway4,
                    new LocalUnresolvedTaskManagerLocation());

            Duration timeout = Duration.ofSeconds(10);
            Set<ResultPartitionID> expectedResultPartitions = new HashSet<>();
            expectedResultPartitions.add(descriptor1.getResultPartitionID());
            expectedResultPartitions.add(descriptor2.getResultPartitionID());
            expectedResultPartitions.add(descriptor4.getResultPartitionID());
            CompletableFuture<Collection<PartitionWithMetrics>> future =
                    jobMasterGateway.getPartitionWithMetrics(timeout, expectedResultPartitions);

            expectedMetrics = new ArrayList<>(defaultPartitionWithMetrics1);
            expectedMetrics.addAll(defaultPartitionWithMetrics2);
            expectedMetrics.addAll(defaultPartitionWithMetrics4);

            assertThat(future).succeedsWithin(timeout);
            verifyPartitionMetrics(
                    future.get().stream()
                            .collect(
                                    Collectors.toMap(
                                            metrics ->
                                                    metrics.getPartition().getResultPartitionID(),
                                            metrics -> metrics)),
                    expectedMetrics);

            // fetch partition after timeout
            fetchPartitionsFuture.complete(defaultPartitionWithMetrics3);
            assertThat(releaseRequest.get())
                    .isEqualTo(
                            Tuple2.of(
                                    jobGraph.getJobID(),
                                    Collections.singleton(descriptor3.getResultPartitionID())));

            // after partitions fetching finished and then register tm5
            CompletableFuture<Collection<PartitionWithMetrics>> requestFuture =
                    new CompletableFuture<>();
            TestingTaskExecutorGateway taskExecutorGateway5 =
                    new TestingTaskExecutorGatewayBuilder()
                            .setRequestPartitionWithMetricsFunction(
                                    ignored -> {
                                        requestFuture.complete(null);
                                        return requestFuture;
                                    })
                            .setAddress("tm5")
                            .createTestingTaskExecutorGateway();

            registerSlotsAtJobMaster(
                    1,
                    jobMasterGateway,
                    jobGraph.getJobID(),
                    taskExecutorGateway5,
                    new LocalUnresolvedTaskManagerLocation());

            assertThatThrownBy(() -> requestFuture.get(timeout.toMillis(), TimeUnit.MILLISECONDS))
                    .isInstanceOf(TimeoutException.class);
        }
    }