Examples with TestingResourceManagerGateway org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway used on opensource projects

Example 11 with TestingResourceManagerGateway

use of org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway in project flink by apache.

the class TaskExecutorTest method testHeartbeatReporting.

/**
 * Tests that the correct partition/slot report is sent as part of the heartbeat response.
 */
@Test
public void testHeartbeatReporting() throws Exception {
    final String rmAddress = "rm";
    final UUID rmLeaderId = UUID.randomUUID();
    // register the mock resource manager gateway
    final TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway();
    final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
    final ResourceID rmResourceId = rmGateway.getOwnResourceId();
    final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), rmResourceId, new ClusterInformation("localhost", 1234)));
    rmGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
        taskExecutorRegistrationFuture.complete(taskExecutorRegistration.getResourceId());
        return registrationResponse;
    });
    final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
    rmGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
        initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
        return CompletableFuture.completedFuture(Acknowledge.get());
    });
    final CompletableFuture<TaskExecutorHeartbeatPayload> heartbeatPayloadCompletableFuture = new CompletableFuture<>();
    rmGateway.setTaskExecutorHeartbeatFunction((resourceID, heartbeatPayload) -> {
        heartbeatPayloadCompletableFuture.complete(heartbeatPayload);
        return FutureUtils.completedVoidFuture();
    });
    rpc.registerGateway(rmAddress, rmGateway);
    final SlotID slotId = buildSlotID(0);
    final ResourceProfile resourceProfile = ResourceProfile.fromResources(1.0, 1);
    final SlotReport slotReport1 = new SlotReport(new SlotStatus(slotId, resourceProfile));
    final SlotReport slotReport2 = new SlotReport(new SlotStatus(slotId, resourceProfile, new JobID(), new AllocationID()));
    final Queue<SlotReport> reports = new ArrayDeque<>(Arrays.asList(slotReport1, slotReport2));
    final TaskSlotTable<Task> taskSlotTable = TestingTaskSlotTable.<Task>newBuilder().createSlotReportSupplier(reports::poll).closeAsyncReturns(CompletableFuture.completedFuture(null)).build();
    final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();
    final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation).setTaskSlotTable(taskSlotTable).setTaskStateManager(localStateStoresManager).build();
    final TaskExecutorPartitionTracker partitionTracker = createPartitionTrackerWithFixedPartitionReport(taskManagerServices.getShuffleEnvironment());
    final TaskExecutor taskManager = createTaskExecutor(taskManagerServices, HEARTBEAT_SERVICES, partitionTracker);
    try {
        taskManager.start();
        // define a leader and see that a registration happens
        resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId);
        // register resource manager success will trigger monitoring heartbeat target between tm
        // and rm
        assertThat(taskExecutorRegistrationFuture.get(), equalTo(unresolvedTaskManagerLocation.getResourceID()));
        assertThat(initialSlotReportFuture.get(), equalTo(slotReport1));
        TaskExecutorGateway taskExecutorGateway = taskManager.getSelfGateway(TaskExecutorGateway.class);
        // trigger the heartbeat asynchronously
        taskExecutorGateway.heartbeatFromResourceManager(rmResourceId);
        // wait for heartbeat response
        SlotReport actualSlotReport = heartbeatPayloadCompletableFuture.get().getSlotReport();
        // the new slot report should be reported
        assertEquals(slotReport2, actualSlotReport);
        ClusterPartitionReport actualClusterPartitionReport = heartbeatPayloadCompletableFuture.get().getClusterPartitionReport();
        assertEquals(partitionTracker.createClusterPartitionReport(), actualClusterPartitionReport);
    } finally {
        RpcUtils.terminateRpcEndpoint(taskManager, timeout);
    }
}

Example 12 with TestingResourceManagerGateway

use of org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway in project flink by apache.

the class TaskExecutorTest method testIgnoringSlotRequestsIfNotRegistered.

/**
 * Tests that we ignore slot requests if the TaskExecutor is not registered at a
 * ResourceManager.
 */
@Test
public void testIgnoringSlotRequestsIfNotRegistered() throws Exception {
    final TaskExecutor taskExecutor = createTaskExecutor(1);
    taskExecutor.start();
    try {
        final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
        final CompletableFuture<RegistrationResponse> registrationFuture = new CompletableFuture<>();
        final CompletableFuture<ResourceID> taskExecutorResourceIdFuture = new CompletableFuture<>();
        testingResourceManagerGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
            taskExecutorResourceIdFuture.complete(taskExecutorRegistration.getResourceId());
            return registrationFuture;
        });
        rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
        resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
        final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);
        final ResourceID resourceId = taskExecutorResourceIdFuture.get();
        final CompletableFuture<Acknowledge> slotRequestResponse = taskExecutorGateway.requestSlot(new SlotID(resourceId, 0), jobId, new AllocationID(), ResourceProfile.ZERO, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
        try {
            slotRequestResponse.get();
            fail("We should not be able to request slots before the TaskExecutor is registered at the ResourceManager.");
        } catch (ExecutionException ee) {
            assertThat(ExceptionUtils.stripExecutionException(ee), instanceOf(TaskManagerException.class));
        }
    } finally {
        RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
    }
}

Example 13 with TestingResourceManagerGateway

use of org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway in project flink by apache.

the class TaskExecutorTest method testInitialSlotReportFailure.

/**
 * Tests that the {@link TaskExecutor} tries to reconnect if the initial slot report fails.
 */
@Test
public void testInitialSlotReportFailure() throws Exception {
    final TaskSlotTable<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(1);
    final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
    final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation).build();
    final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
    taskExecutor.start();
    try {
        final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
        final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
        testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
            try {
                return responseQueue.take();
            } catch (InterruptedException e) {
                return FutureUtils.completedExceptionally(e);
            }
        });
        final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), testingResourceManagerGateway.getOwnResourceId(), new ClusterInformation("foobar", 1234)));
        final CountDownLatch numberRegistrations = new CountDownLatch(2);
        testingResourceManagerGateway.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
            numberRegistrations.countDown();
            return registrationResponse;
        });
        responseQueue.offer(FutureUtils.completedExceptionally(new FlinkException("Test exception")));
        responseQueue.offer(CompletableFuture.completedFuture(Acknowledge.get()));
        rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
        resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
        // wait for the second registration attempt
        numberRegistrations.await();
    } finally {
        RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
    }
}

Example 14 with TestingResourceManagerGateway

use of org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway in project flink by apache.

the class TaskExecutorTest method testTriggerRegistrationOnLeaderChange.

@Test
public void testTriggerRegistrationOnLeaderChange() throws Exception {
    final UUID leaderId1 = UUID.randomUUID();
    final UUID leaderId2 = UUID.randomUUID();
    // register the mock resource manager gateways
    final CompletableFuture<TaskExecutorRegistration> rmGateway1TaskExecutorRegistration = new CompletableFuture<>();
    TestingResourceManagerGateway rmGateway1 = new TestingResourceManagerGateway();
    rmGateway1.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
        rmGateway1TaskExecutorRegistration.complete(taskExecutorRegistration);
        return createRegistrationResponse(rmGateway1);
    });
    final CompletableFuture<TaskExecutorRegistration> rmGateway2TaskExecutorRegistration = new CompletableFuture<>();
    TestingResourceManagerGateway rmGateway2 = new TestingResourceManagerGateway();
    rmGateway2.setRegisterTaskExecutorFunction(taskExecutorRegistration -> {
        rmGateway2TaskExecutorRegistration.complete(taskExecutorRegistration);
        return createRegistrationResponse(rmGateway2);
    });
    rpc.registerGateway(rmGateway1.getAddress(), rmGateway1);
    rpc.registerGateway(rmGateway2.getAddress(), rmGateway2);
    final TaskSlotTable<Task> taskSlotTable = TestingTaskSlotTable.<Task>newBuilder().createSlotReportSupplier(SlotReport::new).closeAsyncReturns(CompletableFuture.completedFuture(null)).build();
    final TaskExecutorLocalStateStoresManager localStateStoresManager = createTaskExecutorLocalStateStoresManager();
    final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setUnresolvedTaskManagerLocation(unresolvedTaskManagerLocation).setTaskSlotTable(taskSlotTable).setTaskStateManager(localStateStoresManager).build();
    TaskExecutor taskManager = createTaskExecutor(taskManagerServices);
    try {
        taskManager.start();
        String taskManagerAddress = taskManager.getAddress();
        // no connection initially, since there is no leader
        assertNull(taskManager.getResourceManagerConnection());
        // define a leader and see that a registration happens
        resourceManagerLeaderRetriever.notifyListener(rmGateway1.getAddress(), leaderId1);
        final TaskExecutorRegistration taskExecutorRegistration1 = rmGateway1TaskExecutorRegistration.join();
        assertThat(taskExecutorRegistration1.getTaskExecutorAddress(), is(taskManagerAddress));
        assertThat(taskExecutorRegistration1.getResourceId(), is(unresolvedTaskManagerLocation.getResourceID()));
        assertNotNull(taskManager.getResourceManagerConnection());
        // cancel the leader
        resourceManagerLeaderRetriever.notifyListener(null, null);
        // set a new leader, see that a registration happens
        resourceManagerLeaderRetriever.notifyListener(rmGateway2.getAddress(), leaderId2);
        final TaskExecutorRegistration taskExecutorRegistration2 = rmGateway2TaskExecutorRegistration.join();
        assertThat(taskExecutorRegistration2.getTaskExecutorAddress(), is(taskManagerAddress));
        assertThat(taskExecutorRegistration2.getResourceId(), is(unresolvedTaskManagerLocation.getResourceID()));
        assertNotNull(taskManager.getResourceManagerConnection());
    } finally {
        RpcUtils.terminateRpcEndpoint(taskManager, timeout);
    }
}

Example 15 with TestingResourceManagerGateway

use of org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway in project flink by apache.

the class TaskExecutorTest method testDisconnectFromJobMasterWhenNewLeader.

/**
 * Tests that the TaskExecutor disconnects from the JobMaster if a new leader is detected.
 */
@Test
public void testDisconnectFromJobMasterWhenNewLeader() throws Exception {
    final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(TaskSlotUtils.createTaskSlotTable(1)).build();
    final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
    final CompletableFuture<Integer> offeredSlotsFuture = new CompletableFuture<>();
    final CompletableFuture<ResourceID> disconnectFuture = new CompletableFuture<>();
    final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder().setOfferSlotsFunction((resourceID, slotOffers) -> {
        offeredSlotsFuture.complete(slotOffers.size());
        return CompletableFuture.completedFuture(slotOffers);
    }).setDisconnectTaskManagerFunction(resourceID -> {
        disconnectFuture.complete(resourceID);
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).build();
    final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
    final CompletableFuture<Void> initialSlotReportFuture = new CompletableFuture<>();
    resourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
        initialSlotReportFuture.complete(null);
        return CompletableFuture.completedFuture(Acknowledge.get());
    });
    rpc.registerGateway(resourceManagerGateway.getAddress(), resourceManagerGateway);
    rpc.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway);
    try {
        taskExecutor.start();
        TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);
        resourceManagerLeaderRetriever.notifyListener(resourceManagerGateway.getAddress(), resourceManagerGateway.getFencingToken().toUUID());
        initialSlotReportFuture.get();
        ResourceID resourceID = taskManagerServices.getUnresolvedTaskManagerLocation().getResourceID();
        requestSlot(taskExecutorGateway, jobId, new AllocationID(), new SlotID(resourceID, 0), ResourceProfile.ZERO, "foobar", resourceManagerGateway.getFencingToken());
        jobManagerLeaderRetriever.notifyListener(jobMasterGateway.getAddress(), UUID.randomUUID());
        assertThat(offeredSlotsFuture.get(), is(1));
        // notify loss of leadership
        jobManagerLeaderRetriever.notifyListener(null, null);
        assertThat(disconnectFuture.get(), is(resourceID));
    } finally {
        RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
    }
}