Examples with TestExecutionSlotAllocatorFactory org.apache.flink.runtime.scheduler.TestExecutionSlotAllocatorFactory used on opensource projects

Example 1 with TestExecutionSlotAllocatorFactory

use of org.apache.flink.runtime.scheduler.TestExecutionSlotAllocatorFactory in project flink by apache.

the class UpdatePartitionConsumersTest method testUpdatePartitionConsumers.

/**
 * Test BLOCKING partition information are properly updated to consumers when its producer
 * finishes.
 */
@Test
public void testUpdatePartitionConsumers() throws Exception {
    final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();
    final SchedulerBase scheduler = SchedulerTestingUtils.newSchedulerBuilder(jobGraph, ComponentMainThreadExecutorServiceAdapter.forMainThread()).setExecutionSlotAllocatorFactory(new TestExecutionSlotAllocatorFactory(taskManagerGateway)).build();
    final ExecutionVertex ev1 = scheduler.getExecutionVertex(new ExecutionVertexID(v1.getID(), 0));
    final ExecutionVertex ev2 = scheduler.getExecutionVertex(new ExecutionVertexID(v2.getID(), 0));
    final ExecutionVertex ev3 = scheduler.getExecutionVertex(new ExecutionVertexID(v3.getID(), 0));
    final ExecutionVertex ev4 = scheduler.getExecutionVertex(new ExecutionVertexID(v4.getID(), 0));
    final CompletableFuture<TaskDeploymentDescriptor> ev4TddFuture = new CompletableFuture<>();
    taskManagerGateway.setSubmitConsumer(tdd -> {
        if (tdd.getExecutionAttemptId().equals(ev4.getCurrentExecutionAttempt().getAttemptId())) {
            ev4TddFuture.complete(tdd);
        }
    });
    scheduler.startScheduling();
    assertThat(ev1.getExecutionState(), is(ExecutionState.DEPLOYING));
    assertThat(ev2.getExecutionState(), is(ExecutionState.DEPLOYING));
    assertThat(ev3.getExecutionState(), is(ExecutionState.DEPLOYING));
    assertThat(ev4.getExecutionState(), is(ExecutionState.DEPLOYING));
    updateState(scheduler, ev1, ExecutionState.INITIALIZING);
    updateState(scheduler, ev1, ExecutionState.RUNNING);
    updateState(scheduler, ev2, ExecutionState.INITIALIZING);
    updateState(scheduler, ev2, ExecutionState.RUNNING);
    updateState(scheduler, ev3, ExecutionState.INITIALIZING);
    updateState(scheduler, ev3, ExecutionState.RUNNING);
    updateState(scheduler, ev4, ExecutionState.INITIALIZING);
    updateState(scheduler, ev4, ExecutionState.RUNNING);
    final InputGateDeploymentDescriptor ev4Igdd2 = ev4TddFuture.get(TIMEOUT, TimeUnit.MILLISECONDS).getInputGates().get(1);
    assertThat(ev4Igdd2.getShuffleDescriptors()[0], instanceOf(UnknownShuffleDescriptor.class));
    final CompletableFuture<Void> updatePartitionFuture = new CompletableFuture<>();
    taskManagerGateway.setUpdatePartitionsConsumer((attemptId, partitionInfos, time) -> {
        assertThat(attemptId, equalTo(ev4.getCurrentExecutionAttempt().getAttemptId()));
        final List<PartitionInfo> partitionInfoList = IterableUtils.toStream(partitionInfos).collect(Collectors.toList());
        assertThat(partitionInfoList, hasSize(1));
        final PartitionInfo partitionInfo = partitionInfoList.get(0);
        assertThat(partitionInfo.getIntermediateDataSetID(), equalTo(v3.getProducedDataSets().get(0).getId()));
        assertThat(partitionInfo.getShuffleDescriptor(), instanceOf(NettyShuffleDescriptor.class));
        updatePartitionFuture.complete(null);
    });
    updateState(scheduler, ev1, ExecutionState.FINISHED);
    updateState(scheduler, ev3, ExecutionState.FINISHED);
    updatePartitionFuture.get(TIMEOUT, TimeUnit.MILLISECONDS);
}