Examples with TestingTaskExecutorGateway org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway used on opensource projects

Example 11 with TestingTaskExecutorGateway

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway in project flink by apache.

the class DefaultSlotStatusSyncerTest method testAllocateSlot.

@Test
public void testAllocateSlot() throws Exception {
    final FineGrainedTaskManagerTracker taskManagerTracker = new FineGrainedTaskManagerTracker();
    final CompletableFuture<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestFuture = new CompletableFuture<>();
    final CompletableFuture<Acknowledge> responseFuture = new CompletableFuture<>();
    final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(tuple6 -> {
        requestFuture.complete(tuple6);
        return responseFuture;
    }).createTestingTaskExecutorGateway();
    final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(ResourceID.generate(), taskExecutorGateway);
    taskManagerTracker.addTaskManager(taskExecutorConnection, ResourceProfile.ANY, ResourceProfile.ANY);
    final ResourceTracker resourceTracker = new DefaultResourceTracker();
    final JobID jobId = new JobID();
    final SlotStatusSyncer slotStatusSyncer = new DefaultSlotStatusSyncer(TASK_MANAGER_REQUEST_TIMEOUT);
    slotStatusSyncer.initialize(taskManagerTracker, resourceTracker, ResourceManagerId.generate(), TestingUtils.defaultExecutor());
    final CompletableFuture<Void> allocatedFuture = slotStatusSyncer.allocateSlot(taskExecutorConnection.getInstanceID(), jobId, "address", ResourceProfile.ANY);
    final AllocationID allocationId = requestFuture.get().f2;
    assertThat(resourceTracker.getAcquiredResources(jobId), contains(ResourceRequirement.create(ResourceProfile.ANY, 1)));
    assertTrue(taskManagerTracker.getAllocatedOrPendingSlot(allocationId).isPresent());
    assertThat(taskManagerTracker.getAllocatedOrPendingSlot(allocationId).get().getJobId(), is(jobId));
    assertThat(taskManagerTracker.getAllocatedOrPendingSlot(allocationId).get().getState(), is(SlotState.PENDING));
    responseFuture.complete(Acknowledge.get());
    assertFalse(allocatedFuture.isCompletedExceptionally());
}

Example 12 with TestingTaskExecutorGateway

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway in project flink by apache.

the class JobMasterPartitionReleaseTest method testPartitionTableCleanupOnDisconnect.

@Test
public void testPartitionTableCleanupOnDisconnect() throws Exception {
    final CompletableFuture<JobID> disconnectTaskExecutorFuture = new CompletableFuture<>();
    final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setDisconnectJobManagerConsumer((jobID, throwable) -> disconnectTaskExecutorFuture.complete(jobID)).createTestingTaskExecutorGateway();
    try (final TestSetup testSetup = new TestSetup(rpcService, testingFatalErrorHandler, testingTaskExecutorGateway)) {
        final JobMasterGateway jobMasterGateway = testSetup.jobMaster.getSelfGateway(JobMasterGateway.class);
        jobMasterGateway.disconnectTaskManager(testSetup.getTaskExecutorResourceID(), new Exception("test"));
        disconnectTaskExecutorFuture.get();
        assertThat(testSetup.getStopTrackingPartitionsTargetResourceId().get(), equalTo(testSetup.getTaskExecutorResourceID()));
    }
}

Example 13 with TestingTaskExecutorGateway

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway in project flink by apache.

the class JobMasterPartitionReleaseTest method testPartitionReleaseOrPromotionOnJobTermination.

private void testPartitionReleaseOrPromotionOnJobTermination(Function<TestSetup, CompletableFuture<Collection<ResultPartitionID>>> callSelector, ExecutionState finalExecutionState) throws Exception {
    final CompletableFuture<TaskDeploymentDescriptor> taskDeploymentDescriptorFuture = new CompletableFuture<>();
    final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setSubmitTaskConsumer((tdd, ignored) -> {
        taskDeploymentDescriptorFuture.complete(tdd);
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).createTestingTaskExecutorGateway();
    try (final TestSetup testSetup = new TestSetup(rpcService, testingFatalErrorHandler, testingTaskExecutorGateway)) {
        ResultPartitionID partitionID0 = new ResultPartitionID();
        ResultPartitionID partitionID1 = new ResultPartitionID();
        testSetup.getPartitionTracker().setGetAllTrackedPartitionsSupplier(() -> {
            ResultPartitionDeploymentDescriptor partitionDesc0 = AbstractPartitionTrackerTest.createResultPartitionDeploymentDescriptor(partitionID0, true);
            ResultPartitionDeploymentDescriptor partitionDesc1 = AbstractPartitionTrackerTest.createResultPartitionDeploymentDescriptor(partitionID1, false);
            return Arrays.asList(partitionDesc0, partitionDesc1);
        });
        final JobMasterGateway jobMasterGateway = testSetup.getJobMasterGateway();
        // update the execution state of the only execution to target state
        // this should trigger the job to finish
        final TaskDeploymentDescriptor taskDeploymentDescriptor = taskDeploymentDescriptorFuture.get();
        jobMasterGateway.updateTaskExecutionState(new TaskExecutionState(taskDeploymentDescriptor.getExecutionAttemptId(), finalExecutionState));
        assertThat(callSelector.apply(testSetup).get(), containsInAnyOrder(partitionID0, partitionID1));
    }
}

Example 14 with TestingTaskExecutorGateway

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway in project flink by apache.

the class JobMasterTest method testRequestPartitionState.

/**
 * Tests the {@link JobMaster#requestPartitionState(IntermediateDataSetID, ResultPartitionID)}
 * call for a finished result partition.
 */
@Test
public void testRequestPartitionState() throws Exception {
    final JobGraph producerConsumerJobGraph = producerConsumerJobGraph();
    final JobMaster jobMaster = new JobMasterBuilder(producerConsumerJobGraph, rpcService).withConfiguration(configuration).withHighAvailabilityServices(haServices).withHeartbeatServices(heartbeatServices).createJobMaster();
    jobMaster.start();
    try {
        final CompletableFuture<TaskDeploymentDescriptor> tddFuture = new CompletableFuture<>();
        final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setSubmitTaskConsumer((taskDeploymentDescriptor, jobMasterId) -> {
            tddFuture.complete(taskDeploymentDescriptor);
            return CompletableFuture.completedFuture(Acknowledge.get());
        }).createTestingTaskExecutorGateway();
        final LocalUnresolvedTaskManagerLocation taskManagerLocation = new LocalUnresolvedTaskManagerLocation();
        final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);
        final Collection<SlotOffer> slotOffers = registerSlotsAtJobMaster(1, jobMasterGateway, producerConsumerJobGraph.getJobID(), testingTaskExecutorGateway, taskManagerLocation);
        assertThat(slotOffers, hasSize(1));
        // obtain tdd for the result partition ids
        final TaskDeploymentDescriptor tdd = tddFuture.get();
        assertThat(tdd.getProducedPartitions(), hasSize(1));
        final ResultPartitionDeploymentDescriptor partition = tdd.getProducedPartitions().iterator().next();
        final ExecutionAttemptID executionAttemptId = tdd.getExecutionAttemptId();
        final ExecutionAttemptID copiedExecutionAttemptId = new ExecutionAttemptID(executionAttemptId);
        // finish the producer task
        jobMasterGateway.updateTaskExecutionState(new TaskExecutionState(executionAttemptId, ExecutionState.FINISHED)).get();
        // request the state of the result partition of the producer
        final ResultPartitionID partitionId = new ResultPartitionID(partition.getPartitionId(), copiedExecutionAttemptId);
        CompletableFuture<ExecutionState> partitionStateFuture = jobMasterGateway.requestPartitionState(partition.getResultId(), partitionId);
        assertThat(partitionStateFuture.get(), equalTo(ExecutionState.FINISHED));
        // ask for unknown result partition
        partitionStateFuture = jobMasterGateway.requestPartitionState(partition.getResultId(), new ResultPartitionID());
        try {
            partitionStateFuture.get();
            fail("Expected failure.");
        } catch (ExecutionException e) {
            assertThat(ExceptionUtils.findThrowable(e, IllegalArgumentException.class).isPresent(), is(true));
        }
        // ask for wrong intermediate data set id
        partitionStateFuture = jobMasterGateway.requestPartitionState(new IntermediateDataSetID(), partitionId);
        try {
            partitionStateFuture.get();
            fail("Expected failure.");
        } catch (ExecutionException e) {
            assertThat(ExceptionUtils.findThrowable(e, IllegalArgumentException.class).isPresent(), is(true));
        }
        // ask for "old" execution
        partitionStateFuture = jobMasterGateway.requestPartitionState(partition.getResultId(), new ResultPartitionID(partition.getPartitionId(), new ExecutionAttemptID()));
        try {
            partitionStateFuture.get();
            fail("Expected failure.");
        } catch (ExecutionException e) {
            assertThat(ExceptionUtils.findThrowable(e, PartitionProducerDisposedException.class).isPresent(), is(true));
        }
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}

Example 15 with TestingTaskExecutorGateway

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway in project flink by apache.

the class JobMasterTest method testTaskExecutorNotReleasedOnFailedAllocationIfPartitionIsAllocated.

@Test
public void testTaskExecutorNotReleasedOnFailedAllocationIfPartitionIsAllocated() throws Exception {
    final JobManagerSharedServices jobManagerSharedServices = new TestingJobManagerSharedServicesBuilder().build();
    final JobGraph jobGraph = JobGraphTestUtils.singleNoOpJobGraph();
    final LocalUnresolvedTaskManagerLocation taskManagerUnresolvedLocation = new LocalUnresolvedTaskManagerLocation();
    final AtomicBoolean isTrackingPartitions = new AtomicBoolean(true);
    final TestingJobMasterPartitionTracker partitionTracker = new TestingJobMasterPartitionTracker();
    partitionTracker.setIsTrackingPartitionsForFunction(ignored -> isTrackingPartitions.get());
    final JobMaster jobMaster = new JobMasterBuilder(jobGraph, rpcService).withConfiguration(configuration).withHighAvailabilityServices(haServices).withJobManagerSharedServices(jobManagerSharedServices).withHeartbeatServices(heartbeatServices).withPartitionTrackerFactory(ignored -> partitionTracker).createJobMaster();
    final CompletableFuture<JobID> disconnectTaskExecutorFuture = new CompletableFuture<>();
    final CompletableFuture<AllocationID> freedSlotFuture = new CompletableFuture<>();
    final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setFreeSlotFunction((allocationID, throwable) -> {
        freedSlotFuture.complete(allocationID);
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).setDisconnectJobManagerConsumer((jobID, throwable) -> disconnectTaskExecutorFuture.complete(jobID)).createTestingTaskExecutorGateway();
    try {
        jobMaster.start();
        final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);
        final Collection<SlotOffer> slotOffers = registerSlotsAtJobMaster(1, jobMasterGateway, jobGraph.getJobID(), testingTaskExecutorGateway, taskManagerUnresolvedLocation);
        // check that we accepted the offered slot
        assertThat(slotOffers, hasSize(1));
        final AllocationID allocationId = slotOffers.iterator().next().getAllocationId();
        jobMasterGateway.failSlot(taskManagerUnresolvedLocation.getResourceID(), allocationId, new FlinkException("Fail allocation test exception"));
        // we should free the slot, but not disconnect from the TaskExecutor as we still have an
        // allocated partition
        assertThat(freedSlotFuture.get(), equalTo(allocationId));
        // trigger some request to guarantee ensure the slotAllocationFailure processing if
        // complete
        jobMasterGateway.requestJobStatus(Time.seconds(5)).get();
        assertThat(disconnectTaskExecutorFuture.isDone(), is(false));
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}