Examples with TaskStateSnapshot org.apache.flink.runtime.checkpoint.TaskStateSnapshot used on opensource projects

Example 31 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class AbstractStreamOperatorTestHarness method initializeState.

/**
 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#initializeState()}.
 * Calls {@link
 * org.apache.flink.streaming.api.operators.SetupableStreamOperator#setup(StreamTask,
 * StreamConfig, Output)} if it was not called before.
 *
 * @param jmOperatorStateHandles the primary state (owned by JM)
 * @param tmOperatorStateHandles the (optional) local state (owned by TM) or null.
 * @throws Exception
 */
public void initializeState(OperatorSubtaskState jmOperatorStateHandles, OperatorSubtaskState tmOperatorStateHandles) throws Exception {
    checkState(!initializeCalled, "TestHarness has already been initialized. Have you " + "opened this harness before initializing it?");
    if (!setupCalled) {
        setup();
    }
    if (jmOperatorStateHandles != null) {
        TaskStateSnapshot jmTaskStateSnapshot = new TaskStateSnapshot();
        jmTaskStateSnapshot.putSubtaskStateByOperatorID(operator.getOperatorID(), jmOperatorStateHandles);
        taskStateManager.setReportedCheckpointId(0);
        taskStateManager.setJobManagerTaskStateSnapshotsByCheckpointId(Collections.singletonMap(0L, jmTaskStateSnapshot));
        if (tmOperatorStateHandles != null) {
            TaskStateSnapshot tmTaskStateSnapshot = new TaskStateSnapshot();
            tmTaskStateSnapshot.putSubtaskStateByOperatorID(operator.getOperatorID(), tmOperatorStateHandles);
            taskStateManager.setTaskManagerTaskStateSnapshotsByCheckpointId(Collections.singletonMap(0L, tmTaskStateSnapshot));
        }
    }
    operator.initializeState(mockTask.createStreamTaskStateInitializer());
    initializeCalled = true;
}

Example 32 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot 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();
            }
        }
    }
}

Example 33 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class TaskStateManagerImplTest method testStateReportingAndRetrieving.

/**
 * Test reporting and retrieving prioritized local and remote state.
 */
@Test
public void testStateReportingAndRetrieving() {
    JobID jobID = new JobID();
    ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
    TestCheckpointResponder testCheckpointResponder = new TestCheckpointResponder();
    TestTaskLocalStateStore testTaskLocalStateStore = new TestTaskLocalStateStore();
    InMemoryStateChangelogStorage changelogStorage = new InMemoryStateChangelogStorage();
    TaskStateManager taskStateManager = taskStateManager(jobID, executionAttemptID, testCheckpointResponder, null, testTaskLocalStateStore, changelogStorage);
    // ---------------------------------------- test reporting
    // -----------------------------------------
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(74L, 11L);
    CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
    TaskStateSnapshot jmTaskStateSnapshot = new TaskStateSnapshot();
    OperatorID operatorID_1 = new OperatorID(1L, 1L);
    OperatorID operatorID_2 = new OperatorID(2L, 2L);
    OperatorID operatorID_3 = new OperatorID(3L, 3L);
    Assert.assertFalse(taskStateManager.prioritizedOperatorState(operatorID_1).isRestored());
    Assert.assertFalse(taskStateManager.prioritizedOperatorState(operatorID_2).isRestored());
    Assert.assertFalse(taskStateManager.prioritizedOperatorState(operatorID_3).isRestored());
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
    // Remote state of operator 1 has only managed keyed state.
    OperatorSubtaskState jmOperatorSubtaskState_1 = OperatorSubtaskState.builder().setManagedKeyedState(StateHandleDummyUtil.createNewKeyedStateHandle(keyGroupRange)).build();
    // Remote state of operator 1 has only raw keyed state.
    OperatorSubtaskState jmOperatorSubtaskState_2 = OperatorSubtaskState.builder().setRawKeyedState(StateHandleDummyUtil.createNewKeyedStateHandle(keyGroupRange)).build();
    jmTaskStateSnapshot.putSubtaskStateByOperatorID(operatorID_1, jmOperatorSubtaskState_1);
    jmTaskStateSnapshot.putSubtaskStateByOperatorID(operatorID_2, jmOperatorSubtaskState_2);
    TaskStateSnapshot tmTaskStateSnapshot = new TaskStateSnapshot();
    // Only operator 1 has a local alternative for the managed keyed state.
    OperatorSubtaskState tmOperatorSubtaskState_1 = OperatorSubtaskState.builder().setManagedKeyedState(StateHandleDummyUtil.createNewKeyedStateHandle(keyGroupRange)).build();
    tmTaskStateSnapshot.putSubtaskStateByOperatorID(operatorID_1, tmOperatorSubtaskState_1);
    taskStateManager.reportTaskStateSnapshots(checkpointMetaData, checkpointMetrics, jmTaskStateSnapshot, tmTaskStateSnapshot);
    TestCheckpointResponder.AcknowledgeReport acknowledgeReport = testCheckpointResponder.getAcknowledgeReports().get(0);
    // checks that the checkpoint responder and the local state store received state as
    // expected.
    Assert.assertEquals(checkpointMetaData.getCheckpointId(), acknowledgeReport.getCheckpointId());
    Assert.assertEquals(checkpointMetrics, acknowledgeReport.getCheckpointMetrics());
    Assert.assertEquals(executionAttemptID, acknowledgeReport.getExecutionAttemptID());
    Assert.assertEquals(jobID, acknowledgeReport.getJobID());
    Assert.assertEquals(jmTaskStateSnapshot, acknowledgeReport.getSubtaskState());
    Assert.assertEquals(tmTaskStateSnapshot, testTaskLocalStateStore.retrieveLocalState(checkpointMetaData.getCheckpointId()));
    // -------------------------------------- test prio retrieving
    // ---------------------------------------
    JobManagerTaskRestore taskRestore = new JobManagerTaskRestore(checkpointMetaData.getCheckpointId(), acknowledgeReport.getSubtaskState());
    taskStateManager = taskStateManager(jobID, executionAttemptID, testCheckpointResponder, taskRestore, testTaskLocalStateStore, changelogStorage);
    // this has remote AND local managed keyed state.
    PrioritizedOperatorSubtaskState prioritized_1 = taskStateManager.prioritizedOperatorState(operatorID_1);
    // this has only remote raw keyed state.
    PrioritizedOperatorSubtaskState prioritized_2 = taskStateManager.prioritizedOperatorState(operatorID_2);
    // not restored.
    PrioritizedOperatorSubtaskState prioritized_3 = taskStateManager.prioritizedOperatorState(operatorID_3);
    Assert.assertTrue(prioritized_1.isRestored());
    Assert.assertTrue(prioritized_2.isRestored());
    Assert.assertTrue(prioritized_3.isRestored());
    Assert.assertTrue(taskStateManager.prioritizedOperatorState(new OperatorID()).isRestored());
    // checks for operator 1.
    Iterator<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState_1 = prioritized_1.getPrioritizedManagedKeyedState().iterator();
    Assert.assertTrue(prioritizedManagedKeyedState_1.hasNext());
    StateObjectCollection<KeyedStateHandle> current = prioritizedManagedKeyedState_1.next();
    KeyedStateHandle keyedStateHandleExp = tmOperatorSubtaskState_1.getManagedKeyedState().iterator().next();
    KeyedStateHandle keyedStateHandleAct = current.iterator().next();
    Assert.assertTrue(keyedStateHandleExp == keyedStateHandleAct);
    Assert.assertTrue(prioritizedManagedKeyedState_1.hasNext());
    current = prioritizedManagedKeyedState_1.next();
    keyedStateHandleExp = jmOperatorSubtaskState_1.getManagedKeyedState().iterator().next();
    keyedStateHandleAct = current.iterator().next();
    Assert.assertTrue(keyedStateHandleExp == keyedStateHandleAct);
    Assert.assertFalse(prioritizedManagedKeyedState_1.hasNext());
    // checks for operator 2.
    Iterator<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState_2 = prioritized_2.getPrioritizedRawKeyedState().iterator();
    Assert.assertTrue(prioritizedRawKeyedState_2.hasNext());
    current = prioritizedRawKeyedState_2.next();
    keyedStateHandleExp = jmOperatorSubtaskState_2.getRawKeyedState().iterator().next();
    keyedStateHandleAct = current.iterator().next();
    Assert.assertTrue(keyedStateHandleExp == keyedStateHandleAct);
    Assert.assertFalse(prioritizedRawKeyedState_2.hasNext());
}

Example 34 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class TaskStateManagerImplTest method testStateRetrievingWithFinishedOperator.

@Test
public void testStateRetrievingWithFinishedOperator() {
    TaskStateSnapshot taskStateSnapshot = TaskStateSnapshot.FINISHED_ON_RESTORE;
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(2, taskStateSnapshot);
    TaskStateManagerImpl stateManager = new TaskStateManagerImpl(new JobID(), new ExecutionAttemptID(), new TestTaskLocalStateStore(), null, jobManagerTaskRestore, new TestCheckpointResponder());
    Assert.assertTrue(stateManager.isTaskDeployedAsFinished());
}

Example 35 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class TaskLocalStateStoreImplTest method retrievePersistedLocalStateFromDisc.

@Test
public void retrievePersistedLocalStateFromDisc() {
    final TaskStateSnapshot taskStateSnapshot = createTaskStateSnapshot();
    final long checkpointId = 0L;
    taskLocalStateStore.storeLocalState(checkpointId, taskStateSnapshot);
    final TaskLocalStateStoreImpl newTaskLocalStateStore = createTaskLocalStateStoreImpl(allocationBaseDirs, jobID, allocationID, jobVertexID, 0);
    final TaskStateSnapshot retrievedTaskStateSnapshot = newTaskLocalStateStore.retrieveLocalState(checkpointId);
    assertThat(retrievedTaskStateSnapshot).isEqualTo(taskStateSnapshot);
}