Examples with SerializableObject org.apache.flink.util.SerializableObject used on opensource projects

Example 1 with SerializableObject

use of org.apache.flink.util.SerializableObject in project flink by apache.

the class CheckpointStateRestoreTest method testSetState.

/**
 * Tests that on restore the task state is reset for each stateful task.
 */
@Test
public void testSetState() {
    try {
        KeyGroupRange keyGroupRange = KeyGroupRange.of(0, 0);
        List<SerializableObject> testStates = Collections.singletonList(new SerializableObject());
        final KeyedStateHandle serializedKeyGroupStates = CheckpointCoordinatorTestingUtils.generateKeyGroupState(keyGroupRange, testStates);
        final JobVertexID statefulId = new JobVertexID();
        final JobVertexID statelessId = new JobVertexID();
        ExecutionGraph graph = new CheckpointCoordinatorTestingUtils.CheckpointExecutionGraphBuilder().addJobVertex(statefulId, 3, 256).addJobVertex(statelessId, 2, 256).build();
        ExecutionJobVertex stateful = graph.getJobVertex(statefulId);
        ExecutionJobVertex stateless = graph.getJobVertex(statelessId);
        ExecutionVertex stateful1 = stateful.getTaskVertices()[0];
        ExecutionVertex stateful2 = stateful.getTaskVertices()[1];
        ExecutionVertex stateful3 = stateful.getTaskVertices()[2];
        ExecutionVertex stateless1 = stateless.getTaskVertices()[0];
        ExecutionVertex stateless2 = stateless.getTaskVertices()[1];
        Execution statefulExec1 = stateful1.getCurrentExecutionAttempt();
        Execution statefulExec2 = stateful2.getCurrentExecutionAttempt();
        Execution statefulExec3 = stateful3.getCurrentExecutionAttempt();
        Execution statelessExec1 = stateless1.getCurrentExecutionAttempt();
        Execution statelessExec2 = stateless2.getCurrentExecutionAttempt();
        ManuallyTriggeredScheduledExecutor manuallyTriggeredScheduledExecutor = new ManuallyTriggeredScheduledExecutor();
        CheckpointCoordinator coord = new CheckpointCoordinatorBuilder().setExecutionGraph(graph).setTimer(manuallyTriggeredScheduledExecutor).build();
        // create ourselves a checkpoint with state
        coord.triggerCheckpoint(false);
        manuallyTriggeredScheduledExecutor.triggerAll();
        PendingCheckpoint pending = coord.getPendingCheckpoints().values().iterator().next();
        final long checkpointId = pending.getCheckpointId();
        final TaskStateSnapshot subtaskStates = new TaskStateSnapshot();
        subtaskStates.putSubtaskStateByOperatorID(OperatorID.fromJobVertexID(statefulId), OperatorSubtaskState.builder().setManagedKeyedState(serializedKeyGroupStates).build());
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statefulExec1.getAttemptId(), checkpointId, new CheckpointMetrics(), subtaskStates), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statefulExec2.getAttemptId(), checkpointId, new CheckpointMetrics(), subtaskStates), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statefulExec3.getAttemptId(), checkpointId, new CheckpointMetrics(), subtaskStates), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statelessExec1.getAttemptId(), checkpointId), TASK_MANAGER_LOCATION_INFO);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), statelessExec2.getAttemptId(), checkpointId), TASK_MANAGER_LOCATION_INFO);
        assertEquals(1, coord.getNumberOfRetainedSuccessfulCheckpoints());
        assertEquals(0, coord.getNumberOfPendingCheckpoints());
        // let the coordinator inject the state
        assertTrue(coord.restoreLatestCheckpointedStateToAll(new HashSet<>(Arrays.asList(stateful, stateless)), false));
        // verify that each stateful vertex got the state
        assertEquals(subtaskStates, statefulExec1.getTaskRestore().getTaskStateSnapshot());
        assertEquals(subtaskStates, statefulExec2.getTaskRestore().getTaskStateSnapshot());
        assertEquals(subtaskStates, statefulExec3.getTaskRestore().getTaskStateSnapshot());
        assertNull(statelessExec1.getTaskRestore());
        assertNull(statelessExec2.getTaskRestore());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}