/*
 * 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.scheduler.benchmark.scheduling;

import org.apache.flink.runtime.execution.ExecutionState;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.runtime.scheduler.benchmark.JobConfiguration;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.apache.flink.runtime.scheduler.strategy.PipelinedRegionSchedulingStrategy;

/**
 * The benchmark of scheduling downstream task in a BATCH job. The related method is {@link
 * PipelinedRegionSchedulingStrategy#onExecutionStateChange}.
 */
public class SchedulingDownstreamTasksInBatchJobBenchmark extends SchedulingBenchmarkBase {

    private ExecutionVertexID executionVertexID;
    private PipelinedRegionSchedulingStrategy schedulingStrategy;

    @Override
    public void setup(JobConfiguration jobConfiguration) throws Exception {
        super.setup(jobConfiguration);

        schedulingStrategy =
                new PipelinedRegionSchedulingStrategy(schedulerOperations, schedulingTopology);

        executionVertexID =
                executionGraph
                        .getJobVertex(jobVertices.get(0).getID())
                        .getTaskVertices()[0]
                        .getID();
        for (ExecutionVertex vertex :
                executionGraph.getJobVertex(jobVertices.get(0).getID()).getTaskVertices()) {
            vertex.finishAllBlockingPartitions();
        }
    }

    public void schedulingDownstreamTasks() {
        schedulingStrategy.onExecutionStateChange(executionVertexID, ExecutionState.FINISHED);
    }
}

@@ -1061,78 +1061,8 @@
 on;%0A
-import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;%0A
 impo
@@ -1401,58 +1401,8 @@
 se %7B
-%0A%0A    private ExecutionVertexID executionVertexID;
 %0A   
@@ -1465,16 +1465,46 @@
 ategy;%0A%0A
+    private int parallelism;%0A%0A
     @Ove
@@ -1763,32 +1763,162 @@
     
-executionVertexID =%0A    
+parallelism = jobConfiguration.getParallelism();%0A    %7D%0A%0A    public void schedulingDownstreamTasks() %7B%0A        for (int i = 0; i %3C parallelism - 1; i++) %7B%0A
     
@@ -1921,33 +1921,33 @@
             
-e
+E
 xecution
 Graph%0A      
@@ -1938,17 +1938,28 @@
 tion
-Graph%0A   
+Vertex taskVertex =%0A
     
@@ -1970,25 +1970,38 @@
             
- 
+executionGraph
 .getJobVerte
@@ -2033,33 +2033,8 @@
 D())
-%0A                        
 .get
@@ -2052,10 +2052,11 @@
 s()%5B
-0%5D
+i%5D;
 %0A   
@@ -2068,43 +2068,169 @@
     
-            .getID();%0A        for (
+taskVertex.finishAllBlockingPartitions();%0A%0A            schedulingStrategy.onExecutionStateChange(taskVertex.getID(), ExecutionState.FINISHED);%0A        %7D%0A        
 Exec
@@ -2245,16 +2245,20 @@
 tex 
-v
+lastV
 ertex 
-:
+=
 %0A   
@@ -2329,29 +2329,69 @@
 D())
-.getTaskVertices()) %7B
+%0A                        .getTaskVertices()%5Bparallelism - 1%5D;
 %0A   
@@ -2399,13 +2399,13 @@
     
-    v
+lastV
 erte
@@ -2441,71 +2441,8 @@
 ();%0A
-        %7D%0A    %7D%0A%0A    public void schedulingDownstreamTasks() %7B%0A
     
@@ -2491,25 +2491,26 @@
 nge(
-executionVertexID
+lastVertex.getID()
 , Ex

make schedulingDownStreamTasksInBatchJobBenchmark more in line with the real scene.