Examples with StreamStateHandle org.apache.flink.runtime.state.StreamStateHandle used on opensource projects

Example 26 with StreamStateHandle

use of org.apache.flink.runtime.state.StreamStateHandle in project flink by apache.

the class SavepointV1Test method createTaskStates.

static Collection<TaskState> createTaskStates(int numTaskStates, int numSubtasksPerTask) throws IOException {
    Random random = new Random(numTaskStates * 31 + numSubtasksPerTask);
    List<TaskState> taskStates = new ArrayList<>(numTaskStates);
    for (int stateIdx = 0; stateIdx < numTaskStates; ++stateIdx) {
        int chainLength = 1 + random.nextInt(8);
        TaskState taskState = new TaskState(new JobVertexID(), numSubtasksPerTask, 128, chainLength);
        int noNonPartitionableStateAtIndex = random.nextInt(chainLength);
        int noOperatorStateBackendAtIndex = random.nextInt(chainLength);
        int noOperatorStateStreamAtIndex = random.nextInt(chainLength);
        boolean hasKeyedBackend = random.nextInt(4) != 0;
        boolean hasKeyedStream = random.nextInt(4) != 0;
        for (int subtaskIdx = 0; subtaskIdx < numSubtasksPerTask; subtaskIdx++) {
            List<StreamStateHandle> nonPartitionableStates = new ArrayList<>(chainLength);
            List<OperatorStateHandle> operatorStatesBackend = new ArrayList<>(chainLength);
            List<OperatorStateHandle> operatorStatesStream = new ArrayList<>(chainLength);
            for (int chainIdx = 0; chainIdx < chainLength; ++chainIdx) {
                StreamStateHandle nonPartitionableState = new TestByteStreamStateHandleDeepCompare("a-" + chainIdx, ("Hi-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
                StreamStateHandle operatorStateBackend = new TestByteStreamStateHandleDeepCompare("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
                StreamStateHandle operatorStateStream = new TestByteStreamStateHandleDeepCompare("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
                Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
                offsetsMap.put("A", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 10, 20 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
                offsetsMap.put("B", new OperatorStateHandle.StateMetaInfo(new long[] { 30, 40, 50 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
                offsetsMap.put("C", new OperatorStateHandle.StateMetaInfo(new long[] { 60, 70, 80 }, OperatorStateHandle.Mode.BROADCAST));
                if (chainIdx != noNonPartitionableStateAtIndex) {
                    nonPartitionableStates.add(nonPartitionableState);
                }
                if (chainIdx != noOperatorStateBackendAtIndex) {
                    OperatorStateHandle operatorStateHandleBackend = new OperatorStateHandle(offsetsMap, operatorStateBackend);
                    operatorStatesBackend.add(operatorStateHandleBackend);
                }
                if (chainIdx != noOperatorStateStreamAtIndex) {
                    OperatorStateHandle operatorStateHandleStream = new OperatorStateHandle(offsetsMap, operatorStateStream);
                    operatorStatesStream.add(operatorStateHandleStream);
                }
            }
            KeyGroupsStateHandle keyedStateBackend = null;
            KeyGroupsStateHandle keyedStateStream = null;
            if (hasKeyedBackend) {
                keyedStateBackend = new KeyGroupsStateHandle(new KeyGroupRangeOffsets(1, 1, new long[] { 42 }), new TestByteStreamStateHandleDeepCompare("c", "Hello".getBytes(ConfigConstants.DEFAULT_CHARSET)));
            }
            if (hasKeyedStream) {
                keyedStateStream = new KeyGroupsStateHandle(new KeyGroupRangeOffsets(1, 1, new long[] { 23 }), new TestByteStreamStateHandleDeepCompare("d", "World".getBytes(ConfigConstants.DEFAULT_CHARSET)));
            }
            taskState.putState(subtaskIdx, new SubtaskState(new ChainedStateHandle<>(nonPartitionableStates), new ChainedStateHandle<>(operatorStatesBackend), new ChainedStateHandle<>(operatorStatesStream), keyedStateStream, keyedStateBackend));
        }
        taskStates.add(taskState);
    }
    return taskStates;
}

Example 27 with StreamStateHandle

use of org.apache.flink.runtime.state.StreamStateHandle in project flink by apache.

the class CompletedCheckpointStoreTest method createCheckpoint.

protected TestCompletedCheckpoint createCheckpoint(int id, int numberOfStates, CheckpointProperties props) throws IOException {
    JobVertexID jvid = new JobVertexID();
    Map<JobVertexID, TaskState> taskGroupStates = new HashMap<>();
    TaskState taskState = new TaskState(jvid, numberOfStates, numberOfStates, 1);
    taskGroupStates.put(jvid, taskState);
    for (int i = 0; i < numberOfStates; i++) {
        ChainedStateHandle<StreamStateHandle> stateHandle = CheckpointCoordinatorTest.generateChainedStateHandle(new CheckpointMessagesTest.MyHandle());
        taskState.putState(i, new SubtaskState(stateHandle, null, null, null, null));
    }
    return new TestCompletedCheckpoint(new JobID(), id, 0, taskGroupStates, props);
}

Example 28 with StreamStateHandle

use of org.apache.flink.runtime.state.StreamStateHandle in project flink by apache.

the class RoundRobinOperatorStateRepartitioner method groupByStateName.

/**
	 * Group by the different named states.
	 */
@SuppressWarnings("unchecked, rawtype")
private GroupByStateNameResults groupByStateName(List<OperatorStateHandle> previousParallelSubtaskStates) {
    //Reorganize: group by (State Name -> StreamStateHandle + StateMetaInfo)
    EnumMap<OperatorStateHandle.Mode, Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>>> nameToStateByMode = new EnumMap<>(OperatorStateHandle.Mode.class);
    for (OperatorStateHandle.Mode mode : OperatorStateHandle.Mode.values()) {
        Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> map = new HashMap<>();
        nameToStateByMode.put(mode, new HashMap<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>>());
    }
    for (OperatorStateHandle psh : previousParallelSubtaskStates) {
        for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> e : psh.getStateNameToPartitionOffsets().entrySet()) {
            OperatorStateHandle.StateMetaInfo metaInfo = e.getValue();
            Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> nameToState = nameToStateByMode.get(metaInfo.getDistributionMode());
            List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>> stateLocations = nameToState.get(e.getKey());
            if (stateLocations == null) {
                stateLocations = new ArrayList<>();
                nameToState.put(e.getKey(), stateLocations);
            }
            stateLocations.add(new Tuple2<>(psh.getDelegateStateHandle(), e.getValue()));
        }
    }
    return new GroupByStateNameResults(nameToStateByMode);
}

Example 29 with StreamStateHandle

use of org.apache.flink.runtime.state.StreamStateHandle 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 {
        final ChainedStateHandle<StreamStateHandle> serializedState = CheckpointCoordinatorTest.generateChainedStateHandle(new SerializableObject());
        KeyGroupRange keyGroupRange = KeyGroupRange.of(0, 0);
        List<SerializableObject> testStates = Collections.singletonList(new SerializableObject());
        final KeyGroupsStateHandle serializedKeyGroupStates = CheckpointCoordinatorTest.generateKeyGroupState(keyGroupRange, testStates);
        final JobID jid = new JobID();
        final JobVertexID statefulId = new JobVertexID();
        final JobVertexID statelessId = new JobVertexID();
        Execution statefulExec1 = mockExecution();
        Execution statefulExec2 = mockExecution();
        Execution statefulExec3 = mockExecution();
        Execution statelessExec1 = mockExecution();
        Execution statelessExec2 = mockExecution();
        ExecutionVertex stateful1 = mockExecutionVertex(statefulExec1, statefulId, 0, 3);
        ExecutionVertex stateful2 = mockExecutionVertex(statefulExec2, statefulId, 1, 3);
        ExecutionVertex stateful3 = mockExecutionVertex(statefulExec3, statefulId, 2, 3);
        ExecutionVertex stateless1 = mockExecutionVertex(statelessExec1, statelessId, 0, 2);
        ExecutionVertex stateless2 = mockExecutionVertex(statelessExec2, statelessId, 1, 2);
        ExecutionJobVertex stateful = mockExecutionJobVertex(statefulId, new ExecutionVertex[] { stateful1, stateful2, stateful3 });
        ExecutionJobVertex stateless = mockExecutionJobVertex(statelessId, new ExecutionVertex[] { stateless1, stateless2 });
        Map<JobVertexID, ExecutionJobVertex> map = new HashMap<JobVertexID, ExecutionJobVertex>();
        map.put(statefulId, stateful);
        map.put(statelessId, stateless);
        CheckpointCoordinator coord = new CheckpointCoordinator(jid, 200000L, 200000L, 0, Integer.MAX_VALUE, ExternalizedCheckpointSettings.none(), new ExecutionVertex[] { stateful1, stateful2, stateful3, stateless1, stateless2 }, new ExecutionVertex[] { stateful1, stateful2, stateful3, stateless1, stateless2 }, new ExecutionVertex[0], new StandaloneCheckpointIDCounter(), new StandaloneCompletedCheckpointStore(1), null, Executors.directExecutor());
        // create ourselves a checkpoint with state
        final long timestamp = 34623786L;
        coord.triggerCheckpoint(timestamp, false);
        PendingCheckpoint pending = coord.getPendingCheckpoints().values().iterator().next();
        final long checkpointId = pending.getCheckpointId();
        SubtaskState checkpointStateHandles = new SubtaskState(serializedState, null, null, serializedKeyGroupStates, null);
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec1.getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles));
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec2.getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles));
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec3.getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles));
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statelessExec1.getAttemptId(), checkpointId));
        coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statelessExec2.getAttemptId(), checkpointId));
        assertEquals(1, coord.getNumberOfRetainedSuccessfulCheckpoints());
        assertEquals(0, coord.getNumberOfPendingCheckpoints());
        // let the coordinator inject the state
        coord.restoreLatestCheckpointedState(map, true, false);
        // verify that each stateful vertex got the state
        final TaskStateHandles taskStateHandles = new TaskStateHandles(serializedState, Collections.<Collection<OperatorStateHandle>>singletonList(null), Collections.<Collection<OperatorStateHandle>>singletonList(null), Collections.singletonList(serializedKeyGroupStates), null);
        BaseMatcher<TaskStateHandles> matcher = new BaseMatcher<TaskStateHandles>() {

            @Override
            public boolean matches(Object o) {
                if (o instanceof TaskStateHandles) {
                    return o.equals(taskStateHandles);
                }
                return false;
            }

            @Override
            public void describeTo(Description description) {
                description.appendValue(taskStateHandles);
            }
        };
        verify(statefulExec1, times(1)).setInitialState(Mockito.argThat(matcher));
        verify(statefulExec2, times(1)).setInitialState(Mockito.argThat(matcher));
        verify(statefulExec3, times(1)).setInitialState(Mockito.argThat(matcher));
        verify(statelessExec1, times(0)).setInitialState(Mockito.<TaskStateHandles>any());
        verify(statelessExec2, times(0)).setInitialState(Mockito.<TaskStateHandles>any());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 30 with StreamStateHandle

use of org.apache.flink.runtime.state.StreamStateHandle in project flink by apache.

the class CheckpointCoordinatorTest method testRestoreLatestCheckpointFailureWhenMaxParallelismChanges.

/**
	 * Tests that the checkpoint restoration fails if the max parallelism of the job vertices has
	 * changed.
	 *
	 * @throws Exception
	 */
@Test(expected = IllegalStateException.class)
public void testRestoreLatestCheckpointFailureWhenMaxParallelismChanges() throws Exception {
    final JobID jid = new JobID();
    final long timestamp = System.currentTimeMillis();
    final JobVertexID jobVertexID1 = new JobVertexID();
    final JobVertexID jobVertexID2 = new JobVertexID();
    int parallelism1 = 3;
    int parallelism2 = 2;
    int maxParallelism1 = 42;
    int maxParallelism2 = 13;
    final ExecutionJobVertex jobVertex1 = mockExecutionJobVertex(jobVertexID1, parallelism1, maxParallelism1);
    final ExecutionJobVertex jobVertex2 = mockExecutionJobVertex(jobVertexID2, parallelism2, maxParallelism2);
    List<ExecutionVertex> allExecutionVertices = new ArrayList<>(parallelism1 + parallelism2);
    allExecutionVertices.addAll(Arrays.asList(jobVertex1.getTaskVertices()));
    allExecutionVertices.addAll(Arrays.asList(jobVertex2.getTaskVertices()));
    ExecutionVertex[] arrayExecutionVertices = allExecutionVertices.toArray(new ExecutionVertex[allExecutionVertices.size()]);
    // set up the coordinator and validate the initial state
    CheckpointCoordinator coord = new CheckpointCoordinator(jid, 600000, 600000, 0, Integer.MAX_VALUE, ExternalizedCheckpointSettings.none(), arrayExecutionVertices, arrayExecutionVertices, arrayExecutionVertices, new StandaloneCheckpointIDCounter(), new StandaloneCompletedCheckpointStore(1), null, Executors.directExecutor());
    // trigger the checkpoint
    coord.triggerCheckpoint(timestamp, false);
    assertTrue(coord.getPendingCheckpoints().keySet().size() == 1);
    long checkpointId = Iterables.getOnlyElement(coord.getPendingCheckpoints().keySet());
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, 0L);
    List<KeyGroupRange> keyGroupPartitions1 = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism1, parallelism1);
    List<KeyGroupRange> keyGroupPartitions2 = StateAssignmentOperation.createKeyGroupPartitions(maxParallelism2, parallelism2);
    for (int index = 0; index < jobVertex1.getParallelism(); index++) {
        ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID1, index);
        KeyGroupsStateHandle keyGroupState = generateKeyGroupState(jobVertexID1, keyGroupPartitions1.get(index), false);
        SubtaskState checkpointStateHandles = new SubtaskState(valueSizeTuple, null, null, keyGroupState, null);
        AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(jid, jobVertex1.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles);
        coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
    }
    for (int index = 0; index < jobVertex2.getParallelism(); index++) {
        ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID2, index);
        KeyGroupsStateHandle keyGroupState = generateKeyGroupState(jobVertexID2, keyGroupPartitions2.get(index), false);
        SubtaskState checkpointStateHandles = new SubtaskState(valueSizeTuple, null, null, keyGroupState, null);
        AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(jid, jobVertex2.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(), checkpointId, new CheckpointMetrics(), checkpointStateHandles);
        coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
    }
    List<CompletedCheckpoint> completedCheckpoints = coord.getSuccessfulCheckpoints();
    assertEquals(1, completedCheckpoints.size());
    Map<JobVertexID, ExecutionJobVertex> tasks = new HashMap<>();
    int newMaxParallelism1 = 20;
    int newMaxParallelism2 = 42;
    final ExecutionJobVertex newJobVertex1 = mockExecutionJobVertex(jobVertexID1, parallelism1, newMaxParallelism1);
    final ExecutionJobVertex newJobVertex2 = mockExecutionJobVertex(jobVertexID2, parallelism2, newMaxParallelism2);
    tasks.put(jobVertexID1, newJobVertex1);
    tasks.put(jobVertexID2, newJobVertex2);
    coord.restoreLatestCheckpointedState(tasks, true, false);
    fail("The restoration should have failed because the max parallelism changed.");
}