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

Example 1 with TestingTaskExecutorGateway

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

the class JobMasterTester method createTaskExecutorGateway.

private TaskExecutorGateway createTaskExecutorGateway() {
    final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setSubmitTaskConsumer(this::onSubmitTaskConsumer).setTriggerCheckpointFunction(this::onTriggerCheckpoint).setConfirmCheckpointFunction(this::onConfirmCheckpoint).createTestingTaskExecutorGateway();
    rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);
    return taskExecutorGateway;
}

Example 2 with TestingTaskExecutorGateway

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

the class JobMasterTest method runHeartbeatTest.

private void runHeartbeatTest(TestingTaskExecutorGatewayBuilder testingTaskExecutorGatewayBuilder, HeartbeatServices heartbeatServices) throws Exception {
    final CompletableFuture<JobID> disconnectedJobManagerFuture = new CompletableFuture<>();
    final UnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
    final TestingTaskExecutorGateway taskExecutorGateway = testingTaskExecutorGatewayBuilder.setDisconnectJobManagerConsumer((jobId, throwable) -> disconnectedJobManagerFuture.complete(jobId)).createTestingTaskExecutorGateway();
    rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);
    final JobMaster jobMaster = new JobMasterBuilder(jobGraph, rpcService).withResourceId(jmResourceId).withConfiguration(configuration).withHighAvailabilityServices(haServices).withHeartbeatServices(heartbeatServices).createJobMaster();
    jobMaster.start();
    try {
        final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);
        // register task manager will trigger monitor heartbeat target, schedule heartbeat
        // request at interval time
        CompletableFuture<RegistrationResponse> registrationResponse = jobMasterGateway.registerTaskManager(jobGraph.getJobID(), TaskManagerRegistrationInformation.create(taskExecutorGateway.getAddress(), unresolvedTaskManagerLocation, TestingUtils.zeroUUID()), testingTimeout);
        // wait for the completion of the registration
        registrationResponse.get();
        final JobID disconnectedJobManager = disconnectedJobManagerFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);
        assertThat(disconnectedJobManager, equalTo(jobGraph.getJobID()));
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}

Example 3 with TestingTaskExecutorGateway

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

the class DefaultDeclarativeSlotPoolTest method withSlotPoolContainingOneTaskManagerWithTwoSlotsWithUniqueResourceProfiles.

private static void withSlotPoolContainingOneTaskManagerWithTwoSlotsWithUniqueResourceProfiles(QuadConsumer<DefaultDeclarativeSlotPool, SlotOffer, SlotOffer, TaskManagerLocation> test) {
    final DefaultDeclarativeSlotPool slotPool = DefaultDeclarativeSlotPoolBuilder.builder().build();
    final ResourceCounter resourceRequirements = ResourceCounter.withResource(RESOURCE_PROFILE_1, 1).add(RESOURCE_PROFILE_2, 1);
    final LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
    final FreeSlotConsumer freeSlotConsumer = new FreeSlotConsumer();
    final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setFreeSlotFunction(freeSlotConsumer).createTestingTaskExecutorGateway();
    final Iterator<SlotOffer> slotOffers = increaseRequirementsAndOfferSlotsToSlotPool(slotPool, resourceRequirements, taskManagerLocation, testingTaskExecutorGateway).iterator();
    final SlotOffer slot1 = slotOffers.next();
    final SlotOffer slot2 = slotOffers.next();
    test.accept(slotPool, slot1, slot2, taskManagerLocation);
}

Example 4 with TestingTaskExecutorGateway

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

the class DefaultDeclarativeSlotPoolTest method testReleaseSlotReturnsSlot.

@Test
public void testReleaseSlotReturnsSlot() throws InterruptedException {
    final NewSlotsService notifyNewSlots = new NewSlotsService();
    final DefaultDeclarativeSlotPool slotPool = createDefaultDeclarativeSlotPoolWithNewSlotsListener(notifyNewSlots);
    final ResourceCounter resourceRequirements = createResourceRequirements();
    final FreeSlotConsumer freeSlotConsumer = new FreeSlotConsumer();
    final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setFreeSlotFunction(freeSlotConsumer).createTestingTaskExecutorGateway();
    increaseRequirementsAndOfferSlotsToSlotPool(slotPool, resourceRequirements, new LocalTaskManagerLocation(), testingTaskExecutorGateway);
    final Collection<? extends PhysicalSlot> physicalSlots = notifyNewSlots.takeNewSlots();
    final PhysicalSlot physicalSlot = physicalSlots.iterator().next();
    slotPool.releaseSlot(physicalSlot.getAllocationId(), new FlinkException("Test failure"));
    final AllocationID freedSlot = Iterables.getOnlyElement(freeSlotConsumer.drainFreedSlots());
    assertThat(freedSlot, is(physicalSlot.getAllocationId()));
}

Example 5 with TestingTaskExecutorGateway

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

the class DeclarativeSlotManagerTest method testSlotReportWithConflictingJobIdDuringSlotAllocation.

/**
 * Tests that a pending slot allocation is cancelled if a slot report indicates that the slot is
 * already allocated by another job.
 */
@Test
public void testSlotReportWithConflictingJobIdDuringSlotAllocation() throws Exception {
    final ResourceRequirements resourceRequirements = createResourceRequirementsForSingleSlot();
    final ArrayBlockingQueue<SlotID> requestedSlotIds = new ArrayBlockingQueue<>(2);
    final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(FunctionUtils.uncheckedFunction(requestSlotParameters -> {
        requestedSlotIds.put(requestSlotParameters.f0);
        return new CompletableFuture<>();
    })).createTestingTaskExecutorGateway();
    final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection(taskExecutorGateway);
    final ResourceID resourceId = taskExecutorConnection.getResourceID();
    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 ScheduledExecutor mainThreadExecutor = new ManuallyTriggeredScheduledExecutor();
    try (final DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder().setScheduledExecutor(mainThreadExecutor).build()) {
        slotManager.start(ResourceManagerId.generate(), mainThreadExecutor, new TestingResourceActionsBuilder().build());
        slotManager.registerTaskManager(taskExecutorConnection, slotReport, ResourceProfile.ANY, ResourceProfile.ANY);
        slotManager.processResourceRequirements(resourceRequirements);
        final SlotID firstRequestedSlotId = requestedSlotIds.take();
        final SlotID freeSlotId = firstRequestedSlotId.equals(slotId1) ? slotId2 : slotId1;
        final SlotReport newSlotReport = new SlotReport(Arrays.asList(createAllocatedSlotStatus(firstRequestedSlotId), createFreeSlotStatus(freeSlotId)));
        slotManager.reportSlotStatus(taskExecutorConnection.getInstanceID(), newSlotReport);
        final SlotID secondRequestedSlotId = requestedSlotIds.take();
        assertEquals(freeSlotId, secondRequestedSlotId);
    }
}