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

Example 21 with TestingResourceManagerGateway

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

the class TaskExecutorTest method testOfferSlotToJobMasterAfterTimeout.

/**
 * Tests that offers slots to job master timeout and retry.
 */
@Test
public void testOfferSlotToJobMasterAfterTimeout() throws Exception {
    final TaskSlotTable<Task> taskSlotTable = TaskSlotUtils.createTaskSlotTable(2);
    final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();
    final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
    final AllocationID allocationId = new AllocationID();
    final CompletableFuture<ResourceID> initialSlotReportFuture = new CompletableFuture<>();
    final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
    testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
        initialSlotReportFuture.complete(null);
        return CompletableFuture.completedFuture(Acknowledge.get());
    });
    rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
    resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
    final CountDownLatch slotOfferings = new CountDownLatch(3);
    final CompletableFuture<AllocationID> offeredSlotFuture = new CompletableFuture<>();
    final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder().setOfferSlotsFunction((resourceID, slotOffers) -> {
        assertThat(slotOffers.size(), is(1));
        slotOfferings.countDown();
        if (slotOfferings.getCount() == 0) {
            offeredSlotFuture.complete(slotOffers.iterator().next().getAllocationId());
            return CompletableFuture.completedFuture(slotOffers);
        } else {
            return FutureUtils.completedExceptionally(new TimeoutException());
        }
    }).build();
    final String jobManagerAddress = jobMasterGateway.getAddress();
    rpc.registerGateway(jobManagerAddress, jobMasterGateway);
    jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());
    try {
        taskExecutor.start();
        final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);
        // wait for the connection to the ResourceManager
        initialSlotReportFuture.get();
        requestSlot(taskExecutorGateway, jobId, allocationId, new SlotID(taskExecutor.getResourceID(), 0), ResourceProfile.ZERO, jobManagerAddress, testingResourceManagerGateway.getFencingToken());
        slotOfferings.await();
        assertThat(offeredSlotFuture.get(), is(allocationId));
        assertTrue(taskSlotTable.isSlotFree(1));
    } finally {
        RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
    }
}

Example 22 with TestingResourceManagerGateway

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

the class JobMasterTest method testCloseUnestablishedResourceManagerConnection.

/**
 * Tests that we can close an unestablished ResourceManager connection.
 */
@Test
public void testCloseUnestablishedResourceManagerConnection() throws Exception {
    final JobMaster jobMaster = new JobMasterBuilder(jobGraph, rpcService).withConfiguration(configuration).withHighAvailabilityServices(haServices).createJobMaster();
    try {
        jobMaster.start();
        final TestingResourceManagerGateway firstResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
        final TestingResourceManagerGateway secondResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
        final OneShotLatch firstJobManagerRegistration = new OneShotLatch();
        final OneShotLatch secondJobManagerRegistration = new OneShotLatch();
        firstResourceManagerGateway.setRegisterJobManagerFunction((jobMasterId, resourceID, s, jobID) -> {
            firstJobManagerRegistration.trigger();
            return CompletableFuture.completedFuture(firstResourceManagerGateway.getJobMasterRegistrationSuccess());
        });
        secondResourceManagerGateway.setRegisterJobManagerFunction((jobMasterId, resourceID, s, jobID) -> {
            secondJobManagerRegistration.trigger();
            return CompletableFuture.completedFuture(secondResourceManagerGateway.getJobMasterRegistrationSuccess());
        });
        notifyResourceManagerLeaderListeners(firstResourceManagerGateway);
        // wait until we have seen the first registration attempt
        firstJobManagerRegistration.await();
        // this should stop the connection attempts towards the first RM
        notifyResourceManagerLeaderListeners(secondResourceManagerGateway);
        // check that we start registering at the second RM
        secondJobManagerRegistration.await();
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}

Example 23 with TestingResourceManagerGateway

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

the class TaskManagerDetailsHandlerTest method setup.

@Before
public void setup() throws HandlerRequestException {
    resourceManagerGateway = new TestingResourceManagerGateway();
    metricFetcher = new TestingMetricFetcher();
    testInstance = new TaskManagerDetailsHandler(() -> CompletableFuture.completedFuture(null), TestingUtils.TIMEOUT, Collections.emptyMap(), TaskManagerDetailsHeaders.getInstance(), () -> CompletableFuture.completedFuture(resourceManagerGateway), metricFetcher);
}

Example 24 with TestingResourceManagerGateway

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

the class TaskManagerLogListHandlerTest method setUp.

@Before
public void setUp() throws HandlerRequestException {
    resourceManagerGateway = new TestingResourceManagerGateway();
    taskManagerLogListHandler = new TaskManagerLogListHandler(() -> CompletableFuture.completedFuture(null), TestingUtils.TIMEOUT, Collections.emptyMap(), TaskManagerLogsHeaders.getInstance(), () -> CompletableFuture.completedFuture(resourceManagerGateway));
    handlerRequest = createRequest(EXPECTED_TASK_MANAGER_ID);
}

Example 25 with TestingResourceManagerGateway

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

the class JobMasterTest method testResourceManagerConnectionAfterStart.

/**
 * Tests that the a JM connects to the leading RM after regaining leadership.
 */
@Test
public void testResourceManagerConnectionAfterStart() throws Exception {
    final JobMaster jobMaster = new JobMasterBuilder(jobGraph, rpcService).withJobMasterId(jobMasterId).withConfiguration(configuration).withHighAvailabilityServices(haServices).withHeartbeatServices(heartbeatServices).createJobMaster();
    try {
        final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
        final BlockingQueue<JobMasterId> registrationQueue = new ArrayBlockingQueue<>(1);
        testingResourceManagerGateway.setRegisterJobManagerFunction((jobMasterId, resourceID, s, jobID) -> {
            registrationQueue.offer(jobMasterId);
            return CompletableFuture.completedFuture(testingResourceManagerGateway.getJobMasterRegistrationSuccess());
        });
        notifyResourceManagerLeaderListeners(testingResourceManagerGateway);
        jobMaster.start();
        final JobMasterId firstRegistrationAttempt = registrationQueue.take();
        assertThat(firstRegistrationAttempt, equalTo(jobMasterId));
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}