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

Example 1 with KeyedBackendSerializationProxy

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

the class HeapKeyedStateBackend method restorePartitionedState.

@SuppressWarnings({ "unchecked" })
private void restorePartitionedState(Collection<KeyGroupsStateHandle> state) throws Exception {
    final Map<Integer, String> kvStatesById = new HashMap<>();
    int numRegisteredKvStates = 0;
    stateTables.clear();
    for (KeyGroupsStateHandle keyGroupsHandle : state) {
        if (keyGroupsHandle == null) {
            continue;
        }
        FSDataInputStream fsDataInputStream = keyGroupsHandle.openInputStream();
        cancelStreamRegistry.registerClosable(fsDataInputStream);
        try {
            DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(fsDataInputStream);
            KeyedBackendSerializationProxy serializationProxy = new KeyedBackendSerializationProxy(userCodeClassLoader);
            serializationProxy.read(inView);
            List<KeyedBackendSerializationProxy.StateMetaInfo<?, ?>> metaInfoList = serializationProxy.getNamedStateSerializationProxies();
            for (KeyedBackendSerializationProxy.StateMetaInfo<?, ?> metaInfoSerializationProxy : metaInfoList) {
                StateTable<K, ?, ?> stateTable = stateTables.get(metaInfoSerializationProxy.getStateName());
                //important: only create a new table we did not already create it previously
                if (null == stateTable) {
                    RegisteredBackendStateMetaInfo<?, ?> registeredBackendStateMetaInfo = new RegisteredBackendStateMetaInfo<>(metaInfoSerializationProxy);
                    stateTable = newStateTable(registeredBackendStateMetaInfo);
                    stateTables.put(metaInfoSerializationProxy.getStateName(), stateTable);
                    kvStatesById.put(numRegisteredKvStates, metaInfoSerializationProxy.getStateName());
                    ++numRegisteredKvStates;
                }
            }
            for (Tuple2<Integer, Long> groupOffset : keyGroupsHandle.getGroupRangeOffsets()) {
                int keyGroupIndex = groupOffset.f0;
                long offset = groupOffset.f1;
                fsDataInputStream.seek(offset);
                int writtenKeyGroupIndex = inView.readInt();
                Preconditions.checkState(writtenKeyGroupIndex == keyGroupIndex, "Unexpected key-group in restore.");
                for (int i = 0; i < metaInfoList.size(); i++) {
                    int kvStateId = inView.readShort();
                    StateTable<K, ?, ?> stateTable = stateTables.get(kvStatesById.get(kvStateId));
                    StateTableByKeyGroupReader keyGroupReader = StateTableByKeyGroupReaders.readerForVersion(stateTable, serializationProxy.getRestoredVersion());
                    keyGroupReader.readMappingsInKeyGroup(inView, keyGroupIndex);
                }
            }
        } finally {
            cancelStreamRegistry.unregisterClosable(fsDataInputStream);
            IOUtils.closeQuietly(fsDataInputStream);
        }
    }
}

Example 2 with KeyedBackendSerializationProxy

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

the class HeapKeyedStateBackend method snapshot.

@Override
@SuppressWarnings("unchecked")
public RunnableFuture<KeyGroupsStateHandle> snapshot(final long checkpointId, final long timestamp, final CheckpointStreamFactory streamFactory, CheckpointOptions checkpointOptions) throws Exception {
    if (!hasRegisteredState()) {
        return DoneFuture.nullValue();
    }
    long syncStartTime = System.currentTimeMillis();
    Preconditions.checkState(stateTables.size() <= Short.MAX_VALUE, "Too many KV-States: " + stateTables.size() + ". Currently at most " + Short.MAX_VALUE + " states are supported");
    List<KeyedBackendSerializationProxy.StateMetaInfo<?, ?>> metaInfoProxyList = new ArrayList<>(stateTables.size());
    final Map<String, Integer> kVStateToId = new HashMap<>(stateTables.size());
    final Map<StateTable<K, ?, ?>, StateTableSnapshot> cowStateStableSnapshots = new HashedMap(stateTables.size());
    for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
        RegisteredBackendStateMetaInfo<?, ?> metaInfo = kvState.getValue().getMetaInfo();
        KeyedBackendSerializationProxy.StateMetaInfo<?, ?> metaInfoProxy = new KeyedBackendSerializationProxy.StateMetaInfo(metaInfo.getStateType(), metaInfo.getName(), metaInfo.getNamespaceSerializer(), metaInfo.getStateSerializer());
        metaInfoProxyList.add(metaInfoProxy);
        kVStateToId.put(kvState.getKey(), kVStateToId.size());
        StateTable<K, ?, ?> stateTable = kvState.getValue();
        if (null != stateTable) {
            cowStateStableSnapshots.put(stateTable, stateTable.createSnapshot());
        }
    }
    final KeyedBackendSerializationProxy serializationProxy = new KeyedBackendSerializationProxy(keySerializer, metaInfoProxyList);
    //--------------------------------------------------- this becomes the end of sync part
    // implementation of the async IO operation, based on FutureTask
    final AbstractAsyncIOCallable<KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream> ioCallable = new AbstractAsyncIOCallable<KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream>() {

        AtomicBoolean open = new AtomicBoolean(false);

        @Override
        public CheckpointStreamFactory.CheckpointStateOutputStream openIOHandle() throws Exception {
            if (open.compareAndSet(false, true)) {
                CheckpointStreamFactory.CheckpointStateOutputStream stream = streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
                try {
                    cancelStreamRegistry.registerClosable(stream);
                    return stream;
                } catch (Exception ex) {
                    open.set(false);
                    throw ex;
                }
            } else {
                throw new IOException("Operation already opened.");
            }
        }

        @Override
        public KeyGroupsStateHandle performOperation() throws Exception {
            long asyncStartTime = System.currentTimeMillis();
            CheckpointStreamFactory.CheckpointStateOutputStream stream = getIoHandle();
            DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(stream);
            serializationProxy.write(outView);
            long[] keyGroupRangeOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
            for (int keyGroupPos = 0; keyGroupPos < keyGroupRange.getNumberOfKeyGroups(); ++keyGroupPos) {
                int keyGroupId = keyGroupRange.getKeyGroupId(keyGroupPos);
                keyGroupRangeOffsets[keyGroupPos] = stream.getPos();
                outView.writeInt(keyGroupId);
                for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
                    outView.writeShort(kVStateToId.get(kvState.getKey()));
                    cowStateStableSnapshots.get(kvState.getValue()).writeMappingsInKeyGroup(outView, keyGroupId);
                }
            }
            if (open.compareAndSet(true, false)) {
                StreamStateHandle streamStateHandle = stream.closeAndGetHandle();
                KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(keyGroupRange, keyGroupRangeOffsets);
                final KeyGroupsStateHandle keyGroupsStateHandle = new KeyGroupsStateHandle(offsets, streamStateHandle);
                if (asynchronousSnapshots) {
                    LOG.info("Heap backend snapshot ({}, asynchronous part) in thread {} took {} ms.", streamFactory, Thread.currentThread(), (System.currentTimeMillis() - asyncStartTime));
                }
                return keyGroupsStateHandle;
            } else {
                throw new IOException("Checkpoint stream already closed.");
            }
        }

        @Override
        public void done(boolean canceled) {
            if (open.compareAndSet(true, false)) {
                CheckpointStreamFactory.CheckpointStateOutputStream stream = getIoHandle();
                if (null != stream) {
                    cancelStreamRegistry.unregisterClosable(stream);
                    IOUtils.closeQuietly(stream);
                }
            }
            for (StateTableSnapshot snapshot : cowStateStableSnapshots.values()) {
                snapshot.release();
            }
        }
    };
    AsyncStoppableTaskWithCallback<KeyGroupsStateHandle> task = AsyncStoppableTaskWithCallback.from(ioCallable);
    if (!asynchronousSnapshots) {
        task.run();
    }
    LOG.info("Heap backend snapshot (" + streamFactory + ", synchronous part) in thread " + Thread.currentThread() + " took " + (System.currentTimeMillis() - syncStartTime) + " ms.");
    return task;
}