@VisibleForTesting
    BlockingResultInfo getBlockingResultInfo(IntermediateDataSetID resultId) {
        return blockingResultInfos.get(resultId);
    }
@Test
    void testUpdateBlockingResultInfoWhileScheduling() throws Exception {
        JobGraph jobGraph = createJobGraph();
        Iterator<JobVertex> jobVertexIterator = jobGraph.getVertices().iterator();
        JobVertex source1 = jobVertexIterator.next();
        JobVertex source2 = jobVertexIterator.next();
        JobVertex sink = jobVertexIterator.next();

        final TestingJobMasterPartitionTracker partitionTracker =
                new TestingJobMasterPartitionTracker();
        partitionTracker.setIsPartitionTrackedFunction(ignore -> true);
        int maxParallelism = 6;

        AdaptiveBatchScheduler scheduler =
                new DefaultSchedulerBuilder(
                                jobGraph, mainThreadExecutor, EXECUTOR_RESOURCE.getExecutor())
                        .setDelayExecutor(taskRestartExecutor)
                        .setPartitionTracker(partitionTracker)
                        .setRestartBackoffTimeStrategy(
                                new FixedDelayRestartBackoffTimeStrategyFactory(10, 0).create())
                        .setVertexParallelismAndInputInfosDecider(
                                createCustomParallelismDecider(maxParallelism))
                        .setDefaultMaxParallelism(maxParallelism)
                        .buildAdaptiveBatchJobScheduler();

        final DefaultExecutionGraph graph = (DefaultExecutionGraph) scheduler.getExecutionGraph();
        final ExecutionJobVertex source1ExecutionJobVertex = graph.getJobVertex(source1.getID());
        final ExecutionJobVertex sinkExecutionJobVertex = graph.getJobVertex(sink.getID());

        PointwiseBlockingResultInfo blockingResultInfo;

        scheduler.startScheduling();
        // trigger source1 finished.
        transitionExecutionsState(scheduler, ExecutionState.FINISHED, source1);
        blockingResultInfo =
                (PointwiseBlockingResultInfo) getBlockingResultInfo(scheduler, source1);
        assertThat(blockingResultInfo.getNumOfRecordedPartitions()).isEqualTo(SOURCE_PARALLELISM_1);

        // trigger source2 finished.
        transitionExecutionsState(scheduler, ExecutionState.FINISHED, source2);
        blockingResultInfo =
                (PointwiseBlockingResultInfo) getBlockingResultInfo(scheduler, source2);
        assertThat(blockingResultInfo.getNumOfRecordedPartitions()).isEqualTo(SOURCE_PARALLELISM_2);

        // trigger sink fail with partition not found
        triggerFailedByPartitionNotFound(
                scheduler,
                source1ExecutionJobVertex.getTaskVertices()[0],
                sinkExecutionJobVertex.getTaskVertices()[0]);

        taskRestartExecutor.triggerScheduledTasks();

        // check the partition info is reset
        assertThat(
                        ((PointwiseBlockingResultInfo) getBlockingResultInfo(scheduler, source1))
                                .getNumOfRecordedPartitions())
                .isEqualTo(SOURCE_PARALLELISM_1 - 1);
    }