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

Example 31 with KeyedStateHandle

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

the class StreamTaskStateInitializerImpl method keyedStatedBackend.

protected <K> CheckpointableKeyedStateBackend<K> keyedStatedBackend(TypeSerializer<K> keySerializer, String operatorIdentifierText, PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates, CloseableRegistry backendCloseableRegistry, MetricGroup metricGroup, double managedMemoryFraction) throws Exception {
    if (keySerializer == null) {
        return null;
    }
    String logDescription = "keyed state backend for " + operatorIdentifierText;
    TaskInfo taskInfo = environment.getTaskInfo();
    final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(taskInfo.getMaxNumberOfParallelSubtasks(), taskInfo.getNumberOfParallelSubtasks(), taskInfo.getIndexOfThisSubtask());
    // Now restore processing is included in backend building/constructing process, so we need
    // to make sure
    // each stream constructed in restore could also be closed in case of task cancel, for
    // example the data
    // input stream opened for serDe during restore.
    CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
    backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
    BackendRestorerProcedure<CheckpointableKeyedStateBackend<K>, KeyedStateHandle> backendRestorer = new BackendRestorerProcedure<>((stateHandles) -> stateBackend.createKeyedStateBackend(environment, environment.getJobID(), operatorIdentifierText, keySerializer, taskInfo.getMaxNumberOfParallelSubtasks(), keyGroupRange, environment.getTaskKvStateRegistry(), ttlTimeProvider, metricGroup, stateHandles, cancelStreamRegistryForRestore, managedMemoryFraction), backendCloseableRegistry, logDescription);
    try {
        return backendRestorer.createAndRestore(prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
    } finally {
        if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
            IOUtils.closeQuietly(cancelStreamRegistryForRestore);
        }
    }
}

Example 32 with KeyedStateHandle

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

the class StreamTaskStateInitializerImpl method rawKeyedStateInputs.

protected CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs(Iterator<StateObjectCollection<KeyedStateHandle>> restoreStateAlternatives) {
    if (restoreStateAlternatives.hasNext()) {
        Collection<KeyedStateHandle> rawKeyedState = restoreStateAlternatives.next();
        // TODO currently this does not support local state recovery, so we expect there is only
        // one handle.
        Preconditions.checkState(!restoreStateAlternatives.hasNext(), "Local recovery is currently not implemented for raw keyed state, but found state alternative.");
        if (rawKeyedState != null) {
            Collection<KeyGroupsStateHandle> keyGroupsStateHandles = transform(rawKeyedState);
            final CloseableRegistry closeableRegistry = new CloseableRegistry();
            return new CloseableIterable<KeyGroupStatePartitionStreamProvider>() {

                @Override
                public void close() throws IOException {
                    closeableRegistry.close();
                }

                @Override
                public Iterator<KeyGroupStatePartitionStreamProvider> iterator() {
                    return new KeyGroupStreamIterator(keyGroupsStateHandles.iterator(), closeableRegistry);
                }
            };
        }
    }
    return CloseableIterable.empty();
}

Example 33 with KeyedStateHandle

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

the class OperatorSnapshotUtil method readStateHandle.

public static OperatorSubtaskState readStateHandle(String path) throws IOException, ClassNotFoundException {
    FileInputStream in = new FileInputStream(path);
    try (DataInputStream dis = new DataInputStream(in)) {
        // required for backwards compatibility.
        final int v = dis.readInt();
        // still required for compatibility to consume the bytes.
        MetadataV3Serializer.deserializeStreamStateHandle(dis);
        List<OperatorStateHandle> rawOperatorState = null;
        int numRawOperatorStates = dis.readInt();
        if (numRawOperatorStates >= 0) {
            rawOperatorState = new ArrayList<>();
            for (int i = 0; i < numRawOperatorStates; i++) {
                OperatorStateHandle operatorState = MetadataV3Serializer.deserializeOperatorStateHandleUtil(dis);
                rawOperatorState.add(operatorState);
            }
        }
        List<OperatorStateHandle> managedOperatorState = null;
        int numManagedOperatorStates = dis.readInt();
        if (numManagedOperatorStates >= 0) {
            managedOperatorState = new ArrayList<>();
            for (int i = 0; i < numManagedOperatorStates; i++) {
                OperatorStateHandle operatorState = MetadataV3Serializer.deserializeOperatorStateHandleUtil(dis);
                managedOperatorState.add(operatorState);
            }
        }
        List<KeyedStateHandle> rawKeyedState = null;
        int numRawKeyedStates = dis.readInt();
        if (numRawKeyedStates >= 0) {
            rawKeyedState = new ArrayList<>();
            for (int i = 0; i < numRawKeyedStates; i++) {
                KeyedStateHandle keyedState = MetadataV3Serializer.deserializeKeyedStateHandleUtil(dis);
                rawKeyedState.add(keyedState);
            }
        }
        List<KeyedStateHandle> managedKeyedState = null;
        int numManagedKeyedStates = dis.readInt();
        if (numManagedKeyedStates >= 0) {
            managedKeyedState = new ArrayList<>();
            for (int i = 0; i < numManagedKeyedStates; i++) {
                KeyedStateHandle keyedState = MetadataV3Serializer.deserializeKeyedStateHandleUtil(dis);
                managedKeyedState.add(keyedState);
            }
        }
        final StateObjectCollection<InputChannelStateHandle> inputChannelStateHandles = v == MetadataV3Serializer.VERSION ? MetadataV3Serializer.deserializeInputChannelStateHandle(dis) : StateObjectCollection.empty();
        final StateObjectCollection<ResultSubpartitionStateHandle> resultSubpartitionStateHandles = v == MetadataV3Serializer.VERSION ? MetadataV3Serializer.INSTANCE.deserializeResultSubpartitionStateHandle(dis) : StateObjectCollection.empty();
        return OperatorSubtaskState.builder().setManagedOperatorState(new StateObjectCollection<>(managedOperatorState)).setRawOperatorState(new StateObjectCollection<>(rawOperatorState)).setManagedKeyedState(new StateObjectCollection<>(managedKeyedState)).setRawKeyedState(new StateObjectCollection<>(rawKeyedState)).setInputChannelState(inputChannelStateHandles).setResultSubpartitionState(resultSubpartitionStateHandles).build();
    }
}

Example 34 with KeyedStateHandle

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

the class OperatorSnapshotUtil method writeStateHandle.

public static void writeStateHandle(OperatorSubtaskState state, String path) throws IOException {
    FileOutputStream out = new FileOutputStream(path);
    try (DataOutputStream dos = new DataOutputStream(out)) {
        // required for backwards compatibility.
        dos.writeInt(MetadataV3Serializer.VERSION);
        // still required for compatibility
        MetadataV3Serializer.serializeStreamStateHandle(null, dos);
        Collection<OperatorStateHandle> rawOperatorState = state.getRawOperatorState();
        if (rawOperatorState != null) {
            dos.writeInt(rawOperatorState.size());
            for (OperatorStateHandle operatorStateHandle : rawOperatorState) {
                MetadataV3Serializer.serializeOperatorStateHandleUtil(operatorStateHandle, dos);
            }
        } else {
            // this means no states, not even an empty list
            dos.writeInt(-1);
        }
        Collection<OperatorStateHandle> managedOperatorState = state.getManagedOperatorState();
        if (managedOperatorState != null) {
            dos.writeInt(managedOperatorState.size());
            for (OperatorStateHandle operatorStateHandle : managedOperatorState) {
                MetadataV3Serializer.serializeOperatorStateHandleUtil(operatorStateHandle, dos);
            }
        } else {
            // this means no states, not even an empty list
            dos.writeInt(-1);
        }
        Collection<KeyedStateHandle> rawKeyedState = state.getRawKeyedState();
        if (rawKeyedState != null) {
            dos.writeInt(rawKeyedState.size());
            for (KeyedStateHandle keyedStateHandle : rawKeyedState) {
                MetadataV3Serializer.serializeKeyedStateHandleUtil(keyedStateHandle, dos);
            }
        } else {
            // this means no operator states, not even an empty list
            dos.writeInt(-1);
        }
        Collection<KeyedStateHandle> managedKeyedState = state.getManagedKeyedState();
        if (managedKeyedState != null) {
            dos.writeInt(managedKeyedState.size());
            for (KeyedStateHandle keyedStateHandle : managedKeyedState) {
                MetadataV3Serializer.serializeKeyedStateHandleUtil(keyedStateHandle, dos);
            }
        } else {
            // this means no operator states, not even an empty list
            dos.writeInt(-1);
        }
        Collection<InputChannelStateHandle> inputChannelStateHandles = state.getInputChannelState();
        dos.writeInt(inputChannelStateHandles.size());
        for (InputChannelStateHandle inputChannelStateHandle : inputChannelStateHandles) {
            MetadataV3Serializer.INSTANCE.serializeInputChannelStateHandle(inputChannelStateHandle, dos);
        }
        Collection<ResultSubpartitionStateHandle> resultSubpartitionStateHandles = state.getResultSubpartitionState();
        dos.writeInt(inputChannelStateHandles.size());
        for (ResultSubpartitionStateHandle resultSubpartitionStateHandle : resultSubpartitionStateHandles) {
            MetadataV3Serializer.INSTANCE.serializeResultSubpartitionStateHandle(resultSubpartitionStateHandle, dos);
        }
        dos.flush();
    }
}

Example 35 with KeyedStateHandle

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

the class FullSnapshotRestoreOperation method restore.

@Override
public ThrowingIterator<SavepointRestoreResult> restore() throws IOException, StateMigrationException {
    return new ThrowingIterator<SavepointRestoreResult>() {

        private final Iterator<KeyedStateHandle> keyedStateHandlesIterator = restoreStateHandles.iterator();

        @Override
        public boolean hasNext() {
            return keyedStateHandlesIterator.hasNext();
        }

        @Override
        public SavepointRestoreResult next() throws IOException, StateMigrationException {
            KeyedStateHandle keyedStateHandle = keyedStateHandlesIterator.next();
            if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
                throw unexpectedStateHandleException(KeyGroupsStateHandle.class, keyedStateHandle.getClass());
            }
            KeyGroupsStateHandle groupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
            return restoreKeyGroupsInStateHandle(groupsStateHandle);
        }

        @Override
        public void close() {
        }
    };
}