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

Example 26 with KeyedStateHandle

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

the class HeapRestoreOperation method restore.

@Override
public Void restore() throws Exception {
    registeredKVStates.clear();
    registeredPQStates.clear();
    boolean keySerializerRestored = false;
    for (KeyedStateHandle keyedStateHandle : restoreStateHandles) {
        if (keyedStateHandle == null) {
            continue;
        }
        if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
            throw unexpectedStateHandleException(KeyGroupsStateHandle.class, keyedStateHandle.getClass());
        }
        LOG.info("Starting to restore from state handle: {}.", keyedStateHandle);
        KeyGroupsStateHandle keyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
        FSDataInputStream fsDataInputStream = keyGroupsStateHandle.openInputStream();
        cancelStreamRegistry.registerCloseable(fsDataInputStream);
        try {
            DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(fsDataInputStream);
            KeyedBackendSerializationProxy<K> serializationProxy = new KeyedBackendSerializationProxy<>(userCodeClassLoader);
            serializationProxy.read(inView);
            if (!keySerializerRestored) {
                // fetch current serializer now because if it is incompatible, we can't access
                // it anymore to improve the error message
                TypeSerializer<K> currentSerializer = keySerializerProvider.currentSchemaSerializer();
                // check for key serializer compatibility; this also reconfigures the
                // key serializer to be compatible, if it is required and is possible
                TypeSerializerSchemaCompatibility<K> keySerializerSchemaCompat = keySerializerProvider.setPreviousSerializerSnapshotForRestoredState(serializationProxy.getKeySerializerSnapshot());
                if (keySerializerSchemaCompat.isCompatibleAfterMigration() || keySerializerSchemaCompat.isIncompatible()) {
                    throw new StateMigrationException("The new key serializer (" + currentSerializer + ") must be compatible with the previous key serializer (" + keySerializerProvider.previousSchemaSerializer() + ").");
                }
                keySerializerRestored = true;
            }
            List<StateMetaInfoSnapshot> restoredMetaInfos = serializationProxy.getStateMetaInfoSnapshots();
            final Map<Integer, StateMetaInfoSnapshot> kvStatesById = this.heapMetaInfoRestoreOperation.createOrCheckStateForMetaInfo(restoredMetaInfos, registeredKVStates, registeredPQStates);
            readStateHandleStateData(fsDataInputStream, inView, keyGroupsStateHandle.getGroupRangeOffsets(), kvStatesById, restoredMetaInfos.size(), serializationProxy.getReadVersion(), serializationProxy.isUsingKeyGroupCompression());
            LOG.info("Finished restoring from state handle: {}.", keyedStateHandle);
        } finally {
            if (cancelStreamRegistry.unregisterCloseable(fsDataInputStream)) {
                IOUtils.closeQuietly(fsDataInputStream);
            }
        }
    }
    return null;
}

Example 27 with KeyedStateHandle

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

the class HeapKeyedStateBackendBuilder method restoreState.

private void restoreState(Map<String, StateTable<K, ?, ?>> registeredKVStates, Map<String, HeapPriorityQueueSnapshotRestoreWrapper<?>> registeredPQStates, InternalKeyContext<K> keyContext, StateTableFactory<K> stateTableFactory) throws BackendBuildingException {
    final RestoreOperation<Void> restoreOperation;
    final KeyedStateHandle firstHandle;
    if (restoreStateHandles.isEmpty()) {
        firstHandle = null;
    } else {
        firstHandle = restoreStateHandles.iterator().next();
    }
    if (firstHandle instanceof SavepointKeyedStateHandle) {
        restoreOperation = new HeapSavepointRestoreOperation<>(restoreStateHandles, keySerializerProvider, userCodeClassLoader, registeredKVStates, registeredPQStates, priorityQueueSetFactory, keyGroupRange, numberOfKeyGroups, stateTableFactory, keyContext);
    } else {
        restoreOperation = new HeapRestoreOperation<>(restoreStateHandles, keySerializerProvider, userCodeClassLoader, registeredKVStates, registeredPQStates, cancelStreamRegistry, priorityQueueSetFactory, keyGroupRange, numberOfKeyGroups, stateTableFactory, keyContext);
    }
    try {
        restoreOperation.restore();
        logger.info("Finished to build heap keyed state-backend.");
    } catch (Exception e) {
        throw new BackendBuildingException("Failed when trying to restore heap backend", e);
    }
}

Example 28 with KeyedStateHandle

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

the class CheckpointTestUtils method randomlySetSubtaskState.

private static void randomlySetSubtaskState(OperatorState taskState, int[] subtasksToSet, Random random, String basePath) {
    boolean hasOperatorStateBackend = random.nextBoolean();
    boolean hasOperatorStateStream = random.nextBoolean();
    boolean hasKeyedBackend = random.nextInt(4) != 0;
    boolean hasKeyedStream = random.nextInt(4) != 0;
    boolean isIncremental = random.nextInt(3) == 0;
    for (int subtaskIdx : subtasksToSet) {
        StreamStateHandle operatorStateBackend = new ByteStreamStateHandle("b", ("Beautiful").getBytes(ConfigConstants.DEFAULT_CHARSET));
        StreamStateHandle operatorStateStream = new ByteStreamStateHandle("b", ("Beautiful").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.UNION));
        final OperatorSubtaskState.Builder state = OperatorSubtaskState.builder();
        if (hasOperatorStateBackend) {
            state.setManagedOperatorState(new OperatorStreamStateHandle(offsetsMap, operatorStateBackend));
        }
        if (hasOperatorStateStream) {
            state.setRawOperatorState(new OperatorStreamStateHandle(offsetsMap, operatorStateStream));
        }
        if (hasKeyedBackend) {
            final KeyedStateHandle stateHandle;
            if (isSavepoint(basePath)) {
                stateHandle = createDummyKeyGroupSavepointStateHandle(random, basePath);
            } else if (isIncremental) {
                stateHandle = createDummyIncrementalKeyedStateHandle(random);
            } else {
                stateHandle = createDummyKeyGroupStateHandle(random, null);
            }
            state.setRawKeyedState(stateHandle);
        }
        if (hasKeyedStream) {
            final KeyedStateHandle stateHandle;
            if (isSavepoint(basePath)) {
                stateHandle = createDummyKeyGroupSavepointStateHandle(random, basePath);
            } else {
                stateHandle = createDummyKeyGroupStateHandle(random, null);
            }
            state.setManagedKeyedState(stateHandle);
        }
        state.setInputChannelState((random.nextBoolean() && !isSavepoint(basePath)) ? singleton(createNewInputChannelStateHandle(random.nextInt(5), random)) : empty());
        state.setResultSubpartitionState((random.nextBoolean() && !isSavepoint(basePath)) ? singleton(createNewResultSubpartitionStateHandle(random.nextInt(5), random)) : empty());
        taskState.putState(subtaskIdx, state.build());
    }
}

Example 29 with KeyedStateHandle

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

Example 30 with KeyedStateHandle

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

the class PrioritizedOperatorSubtaskStateTest method generateForConfiguration.

/**
 * Generator for all 3^4 = 81 possible configurations of a OperatorSubtaskState: - 4 different
 * sub-states: managed/raw + operator/keyed. - 3 different options per sub-state: empty
 * (simulate no state), single handle (simulate recovery), 2 handles (simulate e.g. rescaling)
 */
private OperatorSubtaskState generateForConfiguration(int conf) {
    // 3^4
    Preconditions.checkState(conf >= 0 && conf <= 80);
    final int numModes = 3;
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 4);
    KeyGroupRange keyGroupRange1 = new KeyGroupRange(0, 2);
    KeyGroupRange keyGroupRange2 = new KeyGroupRange(3, 4);
    int div = 1;
    int mode = (conf / div) % numModes;
    StateObjectCollection<OperatorStateHandle> s1 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewOperatorStateHandle(2, RANDOM))) : new StateObjectCollection<>(Arrays.asList(createNewOperatorStateHandle(2, RANDOM), createNewOperatorStateHandle(2, RANDOM)));
    div *= numModes;
    mode = (conf / div) % numModes;
    StateObjectCollection<OperatorStateHandle> s2 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewOperatorStateHandle(2, RANDOM))) : new StateObjectCollection<>(Arrays.asList(createNewOperatorStateHandle(2, RANDOM), createNewOperatorStateHandle(2, RANDOM)));
    div *= numModes;
    mode = (conf / div) % numModes;
    StateObjectCollection<KeyedStateHandle> s3 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewKeyedStateHandle(keyGroupRange))) : new StateObjectCollection<>(Arrays.asList(createNewKeyedStateHandle(keyGroupRange1), createNewKeyedStateHandle(keyGroupRange2)));
    div *= numModes;
    mode = (conf / div) % numModes;
    StateObjectCollection<KeyedStateHandle> s4 = mode == 0 ? StateObjectCollection.empty() : mode == 1 ? new StateObjectCollection<>(Collections.singletonList(createNewKeyedStateHandle(keyGroupRange))) : new StateObjectCollection<>(Arrays.asList(createNewKeyedStateHandle(keyGroupRange1), createNewKeyedStateHandle(keyGroupRange2)));
    return OperatorSubtaskState.builder().setManagedOperatorState(s1).setRawOperatorState(s2).setManagedKeyedState(s3).setRawKeyedState(s4).build();
}