zhijiangW commented on a change in pull request #8630: [FLINK-12667][runtime] 
Add JobID to TaskExecutorGateway#releasePartitions
URL: https://github.com/apache/flink/pull/8630#discussion_r291095641
 
 

 ##########
 File path: 
flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionTest.java
 ##########
 @@ -293,6 +299,81 @@ public void 
testSlotAllocationCancellationWhenExecutionCancelled() throws Except
                assertThat(canceledSlotRequests, equalTo(slotRequests));
        }
 
+       /**
+        * Tests that the partitions are released in case of a execution 
cancellation after the execution is already finished.
+        */
+       @Test
+       public void testPartitionReleaseOnCancelAfterFinished() throws 
Exception {
+               testPartitionReleaseAfterFinished(Execution::cancel);
+       }
+
+       /**
+        * Tests that the partitions are released in case of a execution 
suspension after the execution is already finished.
+        */
+       @Test
+       public void testPartitionReleaseOnSuspendAfterFinished() throws 
Exception {
+               testPartitionReleaseAfterFinished(Execution::suspend);
+       }
+
+       private void testPartitionReleaseAfterFinished(Consumer<Execution> 
postFinishedExecutionAction) throws Exception {
+               final Tuple2<JobID, Collection<ResultPartitionID>> 
releasedPartitions = Tuple2.of(null, null);
+               final SimpleAckingTaskManagerGateway taskManagerGateway = new 
SimpleAckingTaskManagerGateway();
+               
taskManagerGateway.setReleasePartitionsConsumer(releasedPartitions::setFields);
+
+               final SimpleSlot slot = new SimpleSlot(
+                       new SingleSlotTestingSlotOwner(),
+                       new LocalTaskManagerLocation(),
+                       0,
+                       taskManagerGateway);
+
+               final JobVertex producerVertex = createNoOpJobVertex();
+               final JobVertex consumerVertex = createNoOpJobVertex();
+               consumerVertex.connectNewDataSetAsInput(producerVertex, 
DistributionPattern.ALL_TO_ALL, ResultPartitionType.BLOCKING);
+
+               final ProgrammedSlotProvider slotProvider = new 
ProgrammedSlotProvider(1);
+               slotProvider.addSlot(producerVertex.getID(), 0, 
CompletableFuture.completedFuture(slot));
+
+               ExecutionGraph executionGraph = 
ExecutionGraphTestUtils.createSimpleTestGraph(
+                       new JobID(),
+                       slotProvider,
+                       new NoRestartStrategy(),
+                       producerVertex,
+                       consumerVertex);
+
+               
executionGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());
+
+               ExecutionJobVertex executionJobVertex = 
executionGraph.getJobVertex(producerVertex.getID());
+               ExecutionVertex executionVertex = 
executionJobVertex.getTaskVertices()[0];
+
+               final Execution execution = 
executionVertex.getCurrentExecutionAttempt();
+
+               execution.allocateAndAssignSlotForExecution(
+                       slotProvider,
+                       false,
+                       LocationPreferenceConstraint.ALL,
+                       Collections.emptySet(),
+                       TestingUtils.infiniteTime());
+
+               execution.deploy();
 
 Review comment:
   Maybe we could further provide an utility for creating `Execution` like 
below:
   
   `
   private CompletableFuture<Execution> createExecution(
                TaskManagerGateway taskManagerGateway,
                JobVertex... vertices) throws Exception {
   
                SimpleSlot slot = new SimpleSlot(
                        new SingleSlotTestingSlotOwner(),
                        new LocalTaskManagerLocation(),
                        0,
                        taskManagerGateway);
   
                ProgrammedSlotProvider slotProvider = new 
ProgrammedSlotProvider(1);
                slotProvider.addSlot(vertices[0].getID(), 0, 
CompletableFuture.completedFuture(slot));
   
                ExecutionGraph executionGraph = 
ExecutionGraphTestUtils.createSimpleTestGraph(
                        new JobID(),
                        slotProvider,
                        new NoRestartStrategy(),
                        vertices);
   
                
executionGraph.start(TestingComponentMainThreadExecutorServiceAdapter.forMainThread());
   
                ExecutionJobVertex executionJobVertex = 
executionGraph.getJobVertex(vertices[0].getID());
                ExecutionVertex executionVertex = 
executionJobVertex.getTaskVertices()[0];
                Execution execution = 
executionVertex.getCurrentExecutionAttempt();
   
                CompletableFuture<Execution> allocationFuture = 
execution.allocateAndAssignSlotForExecution(
                        slotProvider,
                        false,
                        LocationPreferenceConstraint.ALL,
                        Collections.emptySet(),
                        TestingUtils.infiniteTime());
   
                return allocationFuture;
        }
   `
   Then we could get `Execution` from future, and further get `ExecutionVertex` 
and `ExecutionGraph` from `Execution`. Then this helper could be reused for 
many existing tests.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to