Examples with TestCheckpointResponder org.apache.flink.runtime.taskmanager.TestCheckpointResponder used on opensource projects

Example 6 with TestCheckpointResponder

use of org.apache.flink.runtime.taskmanager.TestCheckpointResponder in project flink by apache.

the class SourceStreamTaskTest method testTriggeringStopWithSavepointWithDrain.

@Test
public void testTriggeringStopWithSavepointWithDrain() throws Exception {
    SourceFunction<String> testSource = new EmptySource();
    CompletableFuture<Boolean> checkpointCompleted = new CompletableFuture<>();
    CheckpointResponder checkpointResponder = new TestCheckpointResponder() {

        @Override
        public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
            super.acknowledgeCheckpoint(jobID, executionAttemptID, checkpointId, checkpointMetrics, subtaskState);
            checkpointCompleted.complete(null);
        }
    };
    try (StreamTaskMailboxTestHarness<String> harness = new StreamTaskMailboxTestHarnessBuilder<>(SourceStreamTask::new, STRING_TYPE_INFO).setTaskStateSnapshot(1, TaskStateSnapshot.FINISHED_ON_RESTORE).setCheckpointResponder(checkpointResponder).setupOutputForSingletonOperatorChain(new StreamSource<>(testSource)).build()) {
        CompletableFuture<Boolean> triggerResult = harness.streamTask.triggerCheckpointAsync(new CheckpointMetaData(2, 2), CheckpointOptions.alignedNoTimeout(SavepointType.terminate(SavepointFormatType.CANONICAL), CheckpointStorageLocationReference.getDefault()));
        checkpointCompleted.whenComplete((ignored, exception) -> harness.streamTask.notifyCheckpointCompleteAsync(2));
        // Run mailbox till the source thread finished and suspend the mailbox
        harness.streamTask.runMailboxLoop();
        harness.finishProcessing();
        assertTrue(triggerResult.isDone());
        assertTrue(triggerResult.get());
        assertTrue(checkpointCompleted.isDone());
    }
}

Example 7 with TestCheckpointResponder

use of org.apache.flink.runtime.taskmanager.TestCheckpointResponder in project flink by apache.

the class TaskSubmissionTestEnvironment method registerJobMasterConnection.

static void registerJobMasterConnection(JobTable jobTable, JobID jobId, RpcService testingRpcService, JobMasterGateway jobMasterGateway, TaskManagerActions taskManagerActions, Time timeout, MainThreadExecutable mainThreadExecutable) {
    mainThreadExecutable.runAsync(() -> {
        final JobTable.Job job = jobTable.getOrCreateJob(jobId, () -> TestingJobServices.newBuilder().build());
        job.connect(ResourceID.generate(), jobMasterGateway, taskManagerActions, new TestCheckpointResponder(), new TestGlobalAggregateManager(), new RpcResultPartitionConsumableNotifier(jobMasterGateway, testingRpcService.getScheduledExecutor(), timeout), TestingPartitionProducerStateChecker.newBuilder().setPartitionProducerStateFunction((jobID, intermediateDataSetID, resultPartitionID) -> CompletableFuture.completedFuture(ExecutionState.RUNNING)).build());
    });
}

Example 8 with TestCheckpointResponder

use of org.apache.flink.runtime.taskmanager.TestCheckpointResponder in project flink by apache.

the class SourceOperatorStreamTaskTest method testTriggeringStopWithSavepointWithDrain.

@Test
public void testTriggeringStopWithSavepointWithDrain() throws Exception {
    SourceOperatorFactory<Integer> sourceOperatorFactory = new SourceOperatorFactory<>(new MockSource(Boundedness.CONTINUOUS_UNBOUNDED, 2), WatermarkStrategy.noWatermarks());
    CompletableFuture<Boolean> checkpointCompleted = new CompletableFuture<>();
    CheckpointResponder checkpointResponder = new TestCheckpointResponder() {

        @Override
        public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
            super.acknowledgeCheckpoint(jobID, executionAttemptID, checkpointId, checkpointMetrics, subtaskState);
            checkpointCompleted.complete(null);
        }
    };
    try (StreamTaskMailboxTestHarness<Integer> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(SourceOperatorStreamTask::new, BasicTypeInfo.INT_TYPE_INFO).setupOutputForSingletonOperatorChain(sourceOperatorFactory).setCheckpointResponder(checkpointResponder).build()) {
        CompletableFuture<Boolean> triggerResult = testHarness.streamTask.triggerCheckpointAsync(new CheckpointMetaData(2, 2), CheckpointOptions.alignedNoTimeout(SavepointType.terminate(SavepointFormatType.CANONICAL), CheckpointStorageLocationReference.getDefault()));
        checkpointCompleted.whenComplete((ignored, exception) -> testHarness.streamTask.notifyCheckpointCompleteAsync(2));
        testHarness.waitForTaskCompletion();
        testHarness.finishProcessing();
        assertTrue(triggerResult.isDone());
        assertTrue(triggerResult.get());
        assertTrue(checkpointCompleted.isDone());
    }
}

Example 9 with TestCheckpointResponder

use of org.apache.flink.runtime.taskmanager.TestCheckpointResponder in project flink by apache.

the class StreamTaskFinalCheckpointsTest method testWaitingForUnalignedChannelStatesIfFinishedOnRestore.

/**
 * This test verifies for tasks that finished on restore, when taking unaligned checkpoint the
 * asynchronous part would wait for the channel states futures get completed, which means the
 * barriers are aligned.
 */
@Test
public void testWaitingForUnalignedChannelStatesIfFinishedOnRestore() throws Exception {
    OperatorID operatorId = new OperatorID();
    try (StreamTaskMailboxTestHarness<String> harness = new StreamTaskMailboxTestHarnessBuilder<>(OneInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).modifyStreamConfig(streamConfig -> streamConfig.setUnalignedCheckpointsEnabled(true)).addInput(BasicTypeInfo.STRING_TYPE_INFO, 3).setCollectNetworkEvents().setTaskStateSnapshot(1, TaskStateSnapshot.FINISHED_ON_RESTORE).setupOperatorChain(new TestFinishedOnRestoreStreamOperator()).chain(operatorId, new TestFinishedOnRestoreStreamOperator(operatorId), StringSerializer.INSTANCE).finish().build()) {
        // Finish the restore, including state initialization and open.
        harness.processAll();
        TestCheckpointResponder checkpointResponder = harness.getCheckpointResponder();
        checkpointResponder.setAcknowledgeLatch(new OneShotLatch());
        checkpointResponder.setDeclinedLatch(new OneShotLatch());
        CheckpointBarrier unalignedBarrier = new CheckpointBarrier(2, 2, CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, getDefault()));
        // On first unaligned barrier, the task would take snapshot and start the asynchronous
        // part. We slightly extend the process to make the asynchronous part start executing
        // before the other barriers arrived.
        harness.processEvent(unalignedBarrier, 0, 0);
        Thread.sleep(CONCURRENT_EVENT_WAIT_PERIOD_MS);
        // Finish the unaligned checkpoint.
        harness.processEvent(unalignedBarrier, 0, 1);
        harness.processEvent(unalignedBarrier, 0, 2);
        // Wait till the asynchronous part finished either normally or exceptionally.
        CommonTestUtils.waitUntilCondition(() -> checkpointResponder.getAcknowledgeLatch().isTriggered() || checkpointResponder.getDeclinedLatch().isTriggered(), Deadline.fromNow(Duration.ofSeconds(10)));
        assertEquals(Collections.singletonList(2L), checkpointResponder.getAcknowledgeReports().stream().map(TestCheckpointResponder.AbstractReport::getCheckpointId).collect(Collectors.toList()));
        assertEquals(Collections.emptyList(), checkpointResponder.getDeclineReports().stream().map(TestCheckpointResponder.AbstractReport::getCheckpointId).collect(Collectors.toList()));
    }
}

Example 10 with TestCheckpointResponder

use of org.apache.flink.runtime.taskmanager.TestCheckpointResponder in project flink by apache.

the class SourceStreamTaskTest method testTriggeringCheckpointAfterSourceThreadFinished.

@Test
public void testTriggeringCheckpointAfterSourceThreadFinished() throws Exception {
    ResultPartition[] partitionWriters = new ResultPartition[2];
    try (NettyShuffleEnvironment env = new NettyShuffleEnvironmentBuilder().setNumNetworkBuffers(partitionWriters.length * 2).build()) {
        for (int i = 0; i < partitionWriters.length; ++i) {
            partitionWriters[i] = PartitionTestUtils.createPartition(env, ResultPartitionType.PIPELINED_BOUNDED, 1);
            partitionWriters[i].setup();
        }
        final CompletableFuture<Long> checkpointCompleted = new CompletableFuture<>();
        try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(SourceStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).modifyStreamConfig(config -> config.setCheckpointingEnabled(true)).setCheckpointResponder(new TestCheckpointResponder() {

            @Override
            public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
                super.acknowledgeCheckpoint(jobID, executionAttemptID, checkpointId, checkpointMetrics, subtaskState);
                checkpointCompleted.complete(checkpointId);
            }
        }).addAdditionalOutput(partitionWriters).setupOperatorChain(new StreamSource<>(new MockSource(0, 0, 1))).finishForSingletonOperatorChain(StringSerializer.INSTANCE).build()) {
            testHarness.processAll();
            CompletableFuture<Void> taskFinished = testHarness.getStreamTask().getCompletionFuture();
            do {
                testHarness.processAll();
            } while (!taskFinished.isDone());
            Future<Boolean> checkpointFuture = triggerCheckpoint(testHarness, 2);
            // Notifies the result partition that all records are processed after the
            // last checkpoint is triggered.
            checkState(checkpointFuture instanceof CompletableFuture, "The trigger future should " + " be also CompletableFuture.");
            ((CompletableFuture<?>) checkpointFuture).thenAccept((ignored) -> {
                for (ResultPartition resultPartition : partitionWriters) {
                    resultPartition.onSubpartitionAllDataProcessed(0);
                }
            });
            checkpointCompleted.whenComplete((id, error) -> testHarness.getStreamTask().notifyCheckpointCompleteAsync(2));
            testHarness.finishProcessing();
            assertTrue(checkpointFuture.isDone());
            // EndOfUserRecordEvent.
            for (ResultPartition resultPartition : partitionWriters) {
                assertEquals(3, resultPartition.getNumberOfQueuedBuffers());
            }
        }
    } finally {
        for (ResultPartitionWriter writer : partitionWriters) {
            if (writer != null) {
                writer.close();
            }
        }
    }
}