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

Example 11 with TaskExecutorGateway

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

the class DeclarativeSlotManagerTest method testTaskManagerRegistration.

/**
 * Tests that we can register task manager and their slots at the slot manager.
 */
@Test
public void testTaskManagerRegistration() throws Exception {
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
    final ResourceID resourceId = ResourceID.generate();
    final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceId, taskExecutorGateway);
    final SlotID slotId1 = new SlotID(resourceId, 0);
    final SlotID slotId2 = new SlotID(resourceId, 1);
    final SlotReport slotReport = new SlotReport(Arrays.asList(createFreeSlotStatus(slotId1), createFreeSlotStatus(slotId2)));
    final DefaultSlotTracker slotTracker = new DefaultSlotTracker();
    try (DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder().setSlotTracker(slotTracker).buildAndStartWithDirectExec()) {
        slotManager.registerTaskManager(taskManagerConnection, slotReport, ResourceProfile.ANY, ResourceProfile.ANY);
        assertThat("The number registered slots does not equal the expected number.", slotManager.getNumberRegisteredSlots(), is(2));
        assertNotNull(slotTracker.getSlot(slotId1));
        assertNotNull(slotTracker.getSlot(slotId2));
    }
}

Example 12 with TaskExecutorGateway

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

the class ResourceManagerTest method testHeartbeatTimeoutWithTaskExecutor.

@Test
public void testHeartbeatTimeoutWithTaskExecutor() throws Exception {
    final ResourceID taskExecutorId = ResourceID.generate();
    final CompletableFuture<ResourceID> heartbeatRequestFuture = new CompletableFuture<>();
    final CompletableFuture<Exception> disconnectFuture = new CompletableFuture<>();
    final CompletableFuture<ResourceID> stopWorkerFuture = new CompletableFuture<>();
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setDisconnectResourceManagerConsumer(disconnectFuture::complete).setHeartbeatResourceManagerFunction(resourceId -> {
        heartbeatRequestFuture.complete(resourceId);
        return FutureUtils.completedVoidFuture();
    }).createTestingTaskExecutorGateway();
    rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);
    runHeartbeatTimeoutTest(builder -> builder.withStopWorkerFunction((worker) -> {
        stopWorkerFuture.complete(worker);
        return true;
    }), resourceManagerGateway -> {
        registerTaskExecutor(resourceManagerGateway, taskExecutorId, taskExecutorGateway.getAddress());
    }, resourceManagerResourceId -> {
        // might have been completed or not depending whether the timeout was triggered
        // first
        final ResourceID optionalHeartbeatRequestOrigin = heartbeatRequestFuture.getNow(null);
        assertThat(optionalHeartbeatRequestOrigin, anyOf(is(resourceManagerResourceId), is(nullValue())));
        assertThat(disconnectFuture.get(), instanceOf(TimeoutException.class));
        assertThat(stopWorkerFuture.get(), is(taskExecutorId));
    });
}

Example 13 with TaskExecutorGateway

use of org.apache.flink.runtime.taskexecutor.TaskExecutorGateway 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 14 with TaskExecutorGateway

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

the class ResourceManagerPartitionLifecycleTest method runTest.

private void runTest(TaskExecutorSetup taskExecutorBuilderSetup, TestAction testAction) throws Exception {
    final ResourceManagerGateway resourceManagerGateway = createAndStartResourceManager();
    TestingTaskExecutorGatewayBuilder testingTaskExecutorGateway1Builder = new TestingTaskExecutorGatewayBuilder();
    taskExecutorBuilderSetup.accept(testingTaskExecutorGateway1Builder);
    final TaskExecutorGateway taskExecutorGateway1 = testingTaskExecutorGateway1Builder.setAddress(UUID.randomUUID().toString()).createTestingTaskExecutorGateway();
    rpcService.registerGateway(taskExecutorGateway1.getAddress(), taskExecutorGateway1);
    final TaskExecutorGateway taskExecutorGateway2 = new TestingTaskExecutorGatewayBuilder().setAddress(UUID.randomUUID().toString()).createTestingTaskExecutorGateway();
    rpcService.registerGateway(taskExecutorGateway2.getAddress(), taskExecutorGateway2);
    final ResourceID taskManagerId1 = ResourceID.generate();
    final ResourceID taskManagerId2 = ResourceID.generate();
    registerTaskExecutor(resourceManagerGateway, taskManagerId1, taskExecutorGateway1.getAddress());
    registerTaskExecutor(resourceManagerGateway, taskManagerId2, taskExecutorGateway2.getAddress());
    testAction.accept(resourceManagerGateway, taskManagerId1, taskManagerId2);
}

Example 15 with TaskExecutorGateway

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

the class DeclarativeSlotManagerTest method testDuplicateResourceRequirementDeclarationAfterSuccessfulAllocation.

/**
 * Tests that duplicate resource requirement declaration do not result in additional slots being
 * allocated after a pending slot request has been fulfilled but not yet freed.
 */
@Test
public void testDuplicateResourceRequirementDeclarationAfterSuccessfulAllocation() throws Exception {
    final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
    final AtomicInteger allocateResourceCalls = new AtomicInteger(0);
    final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder().setAllocateResourceConsumer(ignored -> allocateResourceCalls.incrementAndGet()).build();
    ResourceRequirements requirements = createResourceRequirementsForSingleSlot();
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
    final ResourceID resourceID = ResourceID.generate();
    final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);
    final SlotID slotId = new SlotID(resourceID, 0);
    final SlotReport slotReport = new SlotReport(createFreeSlotStatus(slotId));
    final DefaultSlotTracker slotTracker = new DefaultSlotTracker();
    try (DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder().setSlotTracker(slotTracker).buildAndStartWithDirectExec(resourceManagerId, resourceManagerActions)) {
        slotManager.registerTaskManager(taskManagerConnection, slotReport, ResourceProfile.ANY, ResourceProfile.ANY);
        slotManager.processResourceRequirements(requirements);
        DeclarativeTaskManagerSlot slot = slotTracker.getSlot(slotId);
        assertThat(slot.getState(), is(SlotState.ALLOCATED));
        slotManager.processResourceRequirements(requirements);
    }
    // check that we have only called the resource allocation only for the first slot request,
    // since the second request is a duplicate
    assertThat(allocateResourceCalls.get(), is(0));
}