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

Example 6 with TestingTaskExecutorGateway

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

the class DeclarativeSlotManagerTest method testSlotRequestFailure.

/**
 * Tests that the SlotManager retries allocating a slot if the TaskExecutor#requestSlot call
 * fails.
 */
@Test
public void testSlotRequestFailure() throws Exception {
    final DefaultSlotTracker slotTracker = new DefaultSlotTracker();
    try (final DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder().setSlotTracker(slotTracker).buildAndStartWithDirectExec()) {
        ResourceRequirements requirements = createResourceRequirementsForSingleSlot();
        slotManager.processResourceRequirements(requirements);
        final BlockingQueue<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
        final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
        final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
        responseQueue.offer(firstManualSlotRequestResponse);
        final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
        responseQueue.offer(secondManualSlotRequestResponse);
        final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> {
            requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple6);
            try {
                return responseQueue.take();
            } catch (InterruptedException ignored) {
                return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
            }
        }).createTestingTaskExecutorGateway();
        final ResourceID taskExecutorResourceId = ResourceID.generate();
        final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
        final SlotReport slotReport = new SlotReport(createFreeSlotStatus(new SlotID(taskExecutorResourceId, 0)));
        slotManager.registerTaskManager(taskExecutionConnection, slotReport, ResourceProfile.ANY, ResourceProfile.ANY);
        final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> firstRequest = requestSlotQueue.take();
        // fail first request
        firstManualSlotRequestResponse.completeExceptionally(new SlotAllocationException("Test exception"));
        final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> secondRequest = requestSlotQueue.take();
        assertThat(secondRequest.f1, equalTo(firstRequest.f1));
        assertThat(secondRequest.f0, equalTo(firstRequest.f0));
        secondManualSlotRequestResponse.complete(Acknowledge.get());
        final DeclarativeTaskManagerSlot slot = slotTracker.getSlot(secondRequest.f0);
        assertThat(slot.getState(), equalTo(SlotState.ALLOCATED));
        assertThat(slot.getJobId(), equalTo(secondRequest.f1));
    }
}

Example 7 with TestingTaskExecutorGateway

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

the class ActiveResourceManagerTest method testCloseTaskManagerConnectionOnWorkerTerminated.

@Test
public void testCloseTaskManagerConnectionOnWorkerTerminated() throws Exception {
    new Context() {

        {
            final ResourceID tmResourceId = ResourceID.generate();
            final CompletableFuture<TaskExecutorProcessSpec> requestWorkerFromDriverFuture = new CompletableFuture<>();
            final CompletableFuture<Void> disconnectResourceManagerFuture = new CompletableFuture<>();
            final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setDisconnectResourceManagerConsumer((ignore) -> disconnectResourceManagerFuture.complete(null)).createTestingTaskExecutorGateway();
            driverBuilder.setRequestResourceFunction(taskExecutorProcessSpec -> {
                requestWorkerFromDriverFuture.complete(taskExecutorProcessSpec);
                return CompletableFuture.completedFuture(tmResourceId);
            });
            runTest(() -> {
                // request a new worker, terminate it after registered
                runInMainThread(() -> getResourceManager().startNewWorker(WORKER_RESOURCE_SPEC)).thenCompose((ignore) -> registerTaskExecutor(tmResourceId, taskExecutorGateway)).thenRun(() -> runInMainThread(() -> getResourceManager().onWorkerTerminated(tmResourceId, "terminate for testing")));
                // verify task manager connection is closed
                disconnectResourceManagerFuture.get(TIMEOUT_SEC, TimeUnit.SECONDS);
            });
        }
    };
}

Example 8 with TestingTaskExecutorGateway

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

the class DeclarativeSlotManagerTest method testSlotRequestRemovedIfTMReportsAllocation.

/**
 * Tests that pending request is removed if task executor reports a slot with the same job id.
 */
@Test
public void testSlotRequestRemovedIfTMReportsAllocation() throws Exception {
    final ResourceTracker resourceTracker = new DefaultResourceTracker();
    final DefaultSlotTracker slotTracker = new DefaultSlotTracker();
    try (final DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder().setResourceTracker(resourceTracker).setSlotTracker(slotTracker).buildAndStartWithDirectExec()) {
        final JobID jobID = new JobID();
        slotManager.processResourceRequirements(createResourceRequirementsForSingleSlot(jobID));
        final BlockingQueue<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
        final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
        final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
        responseQueue.offer(firstManualSlotRequestResponse);
        final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
        responseQueue.offer(secondManualSlotRequestResponse);
        final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> {
            requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple6);
            try {
                return responseQueue.take();
            } catch (InterruptedException ignored) {
                return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
            }
        }).createTestingTaskExecutorGateway();
        final ResourceID taskExecutorResourceId = ResourceID.generate();
        final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
        final SlotReport slotReport = new SlotReport(createFreeSlotStatus(new SlotID(taskExecutorResourceId, 0)));
        slotManager.registerTaskManager(taskExecutionConnection, slotReport, ResourceProfile.ANY, ResourceProfile.ANY);
        final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> firstRequest = requestSlotQueue.take();
        // fail first request
        firstManualSlotRequestResponse.completeExceptionally(new TimeoutException("Test exception to fail first allocation"));
        final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> secondRequest = requestSlotQueue.take();
        // fail second request
        secondManualSlotRequestResponse.completeExceptionally(new SlotOccupiedException("Test exception", new AllocationID(), jobID));
        assertThat(firstRequest.f1, equalTo(jobID));
        assertThat(secondRequest.f1, equalTo(jobID));
        assertThat(secondRequest.f0, equalTo(firstRequest.f0));
        final DeclarativeTaskManagerSlot slot = slotTracker.getSlot(secondRequest.f0);
        assertThat(slot.getState(), equalTo(SlotState.ALLOCATED));
        assertThat(slot.getJobId(), equalTo(firstRequest.f1));
        assertThat(slotManager.getNumberRegisteredSlots(), is(1));
        assertThat(getTotalResourceCount(resourceTracker.getAcquiredResources(jobID)), is(1));
    }
}

Example 9 with TestingTaskExecutorGateway

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

the class DefaultSlotStatusSyncerTest method testSlotStatusProcessing.

@Test
public void testSlotStatusProcessing() {
    final FineGrainedTaskManagerTracker taskManagerTracker = new FineGrainedTaskManagerTracker();
    final ResourceTracker resourceTracker = new DefaultResourceTracker();
    final SlotStatusSyncer slotStatusSyncer = new DefaultSlotStatusSyncer(TASK_MANAGER_REQUEST_TIMEOUT);
    slotStatusSyncer.initialize(taskManagerTracker, resourceTracker, ResourceManagerId.generate(), TestingUtils.defaultExecutor());
    final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(ignored -> new CompletableFuture<>()).createTestingTaskExecutorGateway();
    final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(ResourceID.generate(), taskExecutorGateway);
    final JobID jobId = new JobID();
    final AllocationID allocationId1 = new AllocationID();
    final AllocationID allocationId2 = new AllocationID();
    final SlotID slotId1 = new SlotID(taskExecutorConnection.getResourceID(), 0);
    final SlotID slotId2 = new SlotID(taskExecutorConnection.getResourceID(), 1);
    final SlotID slotId3 = new SlotID(taskExecutorConnection.getResourceID(), 2);
    final ResourceProfile totalResource = ResourceProfile.fromResources(5, 20);
    final ResourceProfile resource = ResourceProfile.fromResources(1, 4);
    final SlotReport slotReport1 = new SlotReport(Arrays.asList(new SlotStatus(slotId1, totalResource), new SlotStatus(slotId2, resource, jobId, allocationId1), new SlotStatus(slotId3, resource, jobId, allocationId2)));
    final SlotReport slotReport2 = new SlotReport(Arrays.asList(new SlotStatus(slotId3, resource), new SlotStatus(slotId2, resource, jobId, allocationId1)));
    taskManagerTracker.addTaskManager(taskExecutorConnection, totalResource, totalResource);
    slotStatusSyncer.reportSlotStatus(taskExecutorConnection.getInstanceID(), slotReport1);
    assertThat(resourceTracker.getAcquiredResources(jobId), contains(ResourceRequirement.create(resource, 2)));
    assertThat(taskManagerTracker.getRegisteredTaskManager(taskExecutorConnection.getInstanceID()).get().getAvailableResource(), equalTo(ResourceProfile.fromResources(3, 12)));
    assertTrue(taskManagerTracker.getAllocatedOrPendingSlot(allocationId1).isPresent());
    assertTrue(taskManagerTracker.getAllocatedOrPendingSlot(allocationId2).isPresent());
    slotStatusSyncer.allocateSlot(taskExecutorConnection.getInstanceID(), jobId, "address", resource);
    assertThat(resourceTracker.getAcquiredResources(jobId), contains(ResourceRequirement.create(resource, 3)));
    assertThat(taskManagerTracker.getRegisteredTaskManager(taskExecutorConnection.getInstanceID()).get().getAvailableResource(), equalTo(ResourceProfile.fromResources(2, 8)));
    final AllocationID allocationId3 = taskManagerTracker.getRegisteredTaskManager(taskExecutorConnection.getInstanceID()).get().getAllocatedSlots().keySet().stream().filter(allocationId -> !allocationId.equals(allocationId1) && !allocationId.equals(allocationId2)).findAny().get();
    // allocationId1 should still be allocated; allocationId2 should be freed; allocationId3
    // should continue to be in a pending state;
    slotStatusSyncer.reportSlotStatus(taskExecutorConnection.getInstanceID(), slotReport2);
    assertThat(resourceTracker.getAcquiredResources(jobId), contains(ResourceRequirement.create(resource, 2)));
    assertThat(taskManagerTracker.getRegisteredTaskManager(taskExecutorConnection.getInstanceID()).get().getAvailableResource(), equalTo(ResourceProfile.fromResources(3, 12)));
    assertTrue(taskManagerTracker.getAllocatedOrPendingSlot(allocationId1).isPresent());
    assertFalse(taskManagerTracker.getAllocatedOrPendingSlot(allocationId2).isPresent());
    assertTrue(taskManagerTracker.getAllocatedOrPendingSlot(allocationId3).isPresent());
    assertThat(taskManagerTracker.getAllocatedOrPendingSlot(allocationId1).get().getState(), is(SlotState.ALLOCATED));
    assertThat(taskManagerTracker.getAllocatedOrPendingSlot(allocationId3).get().getState(), is(SlotState.PENDING));
}

Example 10 with TestingTaskExecutorGateway

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

the class DefaultSlotStatusSyncerTest method testAllocateSlotFailsWithException.

@Test
public void testAllocateSlotFailsWithException() {
    final FineGrainedTaskManagerTracker taskManagerTracker = new FineGrainedTaskManagerTracker();
    final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(ignored -> FutureUtils.completedExceptionally(new TimeoutException("timeout"))).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);
    try {
        allocatedFuture.get();
    } catch (Exception e) {
        assertThat(e.getCause(), instanceOf(TimeoutException.class));
    }
    assertThat(resourceTracker.getAcquiredResources(jobId), is(empty()));
    assertThat(taskManagerTracker.getRegisteredTaskManager(taskExecutorConnection.getInstanceID()).get().getAllocatedSlots().keySet(), is(empty()));
}