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

Example 1 with TestingTaskExecutorGatewayBuilder

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

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

the class JobMasterTest method testJobMasterAcceptsSlotsWhenJobIsRestarting.

@Test
public void testJobMasterAcceptsSlotsWhenJobIsRestarting() throws Exception {
    configuration.set(RestartStrategyOptions.RESTART_STRATEGY, "fixed-delay");
    configuration.set(RestartStrategyOptions.RESTART_STRATEGY_FIXED_DELAY_DELAY, Duration.ofDays(1));
    final int numberSlots = 1;
    final JobMaster jobMaster = new JobMasterBuilder(jobGraph, rpcService).withConfiguration(configuration).createJobMaster();
    try {
        jobMaster.start();
        final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);
        final LocalUnresolvedTaskManagerLocation unresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation();
        registerSlotsAtJobMaster(numberSlots, jobMasterGateway, jobGraph.getJobID(), new TestingTaskExecutorGatewayBuilder().setAddress("firstTaskManager").createTestingTaskExecutorGateway(), unresolvedTaskManagerLocation);
        CommonTestUtils.waitUntilCondition(() -> jobMasterGateway.requestJobStatus(testingTimeout).get() == JobStatus.RUNNING, Deadline.fromNow(TimeUtils.toDuration(testingTimeout)));
        jobMasterGateway.disconnectTaskManager(unresolvedTaskManagerLocation.getResourceID(), new FlinkException("Test exception."));
        CommonTestUtils.waitUntilCondition(() -> jobMasterGateway.requestJobStatus(testingTimeout).get() == JobStatus.RESTARTING, Deadline.fromNow(TimeUtils.toDuration(testingTimeout)));
        assertThat(registerSlotsAtJobMaster(numberSlots, jobMasterGateway, jobGraph.getJobID(), new TestingTaskExecutorGatewayBuilder().setAddress("secondTaskManager").createTestingTaskExecutorGateway(), new LocalUnresolvedTaskManagerLocation()), hasSize(numberSlots));
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}

Example 3 with TestingTaskExecutorGatewayBuilder

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

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

the class JobMasterTest method registerSlotsRequiredForJobExecution.

private static void registerSlotsRequiredForJobExecution(JobMasterGateway jobMasterGateway, JobID jobId, int numSlots) throws ExecutionException, InterruptedException {
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setCancelTaskFunction(executionAttemptId -> {
        jobMasterGateway.updateTaskExecutionState(new TaskExecutionState(executionAttemptId, ExecutionState.CANCELED));
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).createTestingTaskExecutorGateway();
    JobMasterTestUtils.registerTaskExecutorAndOfferSlots(rpcService, jobMasterGateway, jobId, numSlots, taskExecutorGateway, testingTimeout);
}

Example 5 with TestingTaskExecutorGatewayBuilder

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

the class DeclarativeSlotPoolServiceTest method testSlotOfferingOfKnownTaskManager.

@Test
public void testSlotOfferingOfKnownTaskManager() throws Exception {
    final AtomicReference<Collection<? extends SlotOffer>> receivedSlotOffers = new AtomicReference<>();
    try (DeclarativeSlotPoolService declarativeSlotPoolService = createDeclarativeSlotPoolService(new TestingDeclarativeSlotPoolFactory(new TestingDeclarativeSlotPoolBuilder().setOfferSlotsFunction((slotOffers, taskManagerLocation, taskManagerGateway, aLong) -> {
        receivedSlotOffers.set(slotOffers);
        return new ArrayList<>(slotOffers);
    })))) {
        final LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
        declarativeSlotPoolService.registerTaskManager(taskManagerLocation.getResourceID());
        final Collection<SlotOffer> slotOffers = Collections.singletonList(new SlotOffer(new AllocationID(), 0, ResourceProfile.UNKNOWN));
        declarativeSlotPoolService.offerSlots(taskManagerLocation, new RpcTaskManagerGateway(new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway(), jobMasterId), slotOffers);
        assertThat(receivedSlotOffers.get(), is(slotOffers));
    }
}