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

Example 1 with IncrementalRemoteKeyedStateHandle

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

the class CheckpointCoordinatorTest method ackCheckpoint.

private void ackCheckpoint(long checkpointId, CheckpointCoordinator coordinator, JobVertexID ackVertexID, ExecutionGraph graph, TestingStreamStateHandle metaState, TestingStreamStateHandle privateState, TestingStreamStateHandle sharedState) throws CheckpointException {
    Map<StateHandleID, StreamStateHandle> sharedStateMap = new HashMap<>(singletonMap(new StateHandleID("shared-state-key"), sharedState));
    Map<StateHandleID, StreamStateHandle> privateStateMap = new HashMap<>(singletonMap(new StateHandleID("private-state-key"), privateState));
    ExecutionJobVertex jobVertex = graph.getJobVertex(ackVertexID);
    OperatorID operatorID = jobVertex.getOperatorIDs().get(0).getGeneratedOperatorID();
    coordinator.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(graph.getJobID(), jobVertex.getTaskVertices()[0].getCurrentExecutionAttempt().getAttemptId(), checkpointId, new CheckpointMetrics(), new TaskStateSnapshot(singletonMap(operatorID, OperatorSubtaskState.builder().setManagedKeyedState(new IncrementalRemoteKeyedStateHandle(UUID.randomUUID(), KeyGroupRange.of(0, 9), checkpointId, sharedStateMap, privateStateMap, metaState)).build()))), "test");
}

Example 2 with IncrementalRemoteKeyedStateHandle

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

the class MetadataV2V3SerializerBase method serializeKeyedStateHandle.

// ------------------------------------------------------------------------
// keyed state
// ------------------------------------------------------------------------
@VisibleForTesting
static void serializeKeyedStateHandle(KeyedStateHandle stateHandle, DataOutputStream dos) throws IOException {
    if (stateHandle == null) {
        dos.writeByte(NULL_HANDLE);
    } else if (stateHandle instanceof KeyGroupsStateHandle) {
        KeyGroupsStateHandle keyGroupsStateHandle = (KeyGroupsStateHandle) stateHandle;
        if (stateHandle instanceof KeyGroupsSavepointStateHandle) {
            dos.writeByte(SAVEPOINT_KEY_GROUPS_HANDLE);
        } else {
            dos.writeByte(KEY_GROUPS_HANDLE_V2);
        }
        dos.writeInt(keyGroupsStateHandle.getKeyGroupRange().getStartKeyGroup());
        dos.writeInt(keyGroupsStateHandle.getKeyGroupRange().getNumberOfKeyGroups());
        for (int keyGroup : keyGroupsStateHandle.getKeyGroupRange()) {
            dos.writeLong(keyGroupsStateHandle.getOffsetForKeyGroup(keyGroup));
        }
        serializeStreamStateHandle(keyGroupsStateHandle.getDelegateStateHandle(), dos);
        // savepoint state handle would not need to persist state handle id out.
        if (!(stateHandle instanceof KeyGroupsSavepointStateHandle)) {
            writeStateHandleId(stateHandle, dos);
        }
    } else if (stateHandle instanceof IncrementalRemoteKeyedStateHandle) {
        IncrementalRemoteKeyedStateHandle incrementalKeyedStateHandle = (IncrementalRemoteKeyedStateHandle) stateHandle;
        dos.writeByte(INCREMENTAL_KEY_GROUPS_HANDLE_V2);
        dos.writeLong(incrementalKeyedStateHandle.getCheckpointId());
        dos.writeUTF(String.valueOf(incrementalKeyedStateHandle.getBackendIdentifier()));
        dos.writeInt(incrementalKeyedStateHandle.getKeyGroupRange().getStartKeyGroup());
        dos.writeInt(incrementalKeyedStateHandle.getKeyGroupRange().getNumberOfKeyGroups());
        dos.writeLong(incrementalKeyedStateHandle.getCheckpointedSize());
        serializeStreamStateHandle(incrementalKeyedStateHandle.getMetaStateHandle(), dos);
        serializeStreamStateHandleMap(incrementalKeyedStateHandle.getSharedState(), dos);
        serializeStreamStateHandleMap(incrementalKeyedStateHandle.getPrivateState(), dos);
        writeStateHandleId(incrementalKeyedStateHandle, dos);
    } else if (stateHandle instanceof ChangelogStateBackendHandle) {
        ChangelogStateBackendHandle handle = (ChangelogStateBackendHandle) stateHandle;
        dos.writeByte(CHANGELOG_HANDLE);
        dos.writeInt(handle.getKeyGroupRange().getStartKeyGroup());
        dos.writeInt(handle.getKeyGroupRange().getNumberOfKeyGroups());
        dos.writeLong(handle.getCheckpointedSize());
        dos.writeInt(handle.getMaterializedStateHandles().size());
        for (KeyedStateHandle keyedStateHandle : handle.getMaterializedStateHandles()) {
            serializeKeyedStateHandle(keyedStateHandle, dos);
        }
        dos.writeInt(handle.getNonMaterializedStateHandles().size());
        for (KeyedStateHandle k : handle.getNonMaterializedStateHandles()) {
            serializeKeyedStateHandle(k, dos);
        }
        dos.writeLong(handle.getMaterializationID());
        writeStateHandleId(handle, dos);
    } else if (stateHandle instanceof InMemoryChangelogStateHandle) {
        InMemoryChangelogStateHandle handle = (InMemoryChangelogStateHandle) stateHandle;
        dos.writeByte(CHANGELOG_BYTE_INCREMENT_HANDLE);
        dos.writeInt(handle.getKeyGroupRange().getStartKeyGroup());
        dos.writeInt(handle.getKeyGroupRange().getNumberOfKeyGroups());
        dos.writeLong(handle.getFrom());
        dos.writeLong(handle.getTo());
        dos.writeInt(handle.getChanges().size());
        for (StateChange change : handle.getChanges()) {
            dos.writeInt(change.getKeyGroup());
            dos.writeInt(change.getChange().length);
            dos.write(change.getChange());
        }
        writeStateHandleId(handle, dos);
    } else if (stateHandle instanceof ChangelogStateHandleStreamImpl) {
        ChangelogStateHandleStreamImpl handle = (ChangelogStateHandleStreamImpl) stateHandle;
        dos.writeByte(CHANGELOG_FILE_INCREMENT_HANDLE);
        dos.writeInt(handle.getKeyGroupRange().getStartKeyGroup());
        dos.writeInt(handle.getKeyGroupRange().getNumberOfKeyGroups());
        dos.writeInt(handle.getHandlesAndOffsets().size());
        for (Tuple2<StreamStateHandle, Long> streamHandleAndOffset : handle.getHandlesAndOffsets()) {
            dos.writeLong(streamHandleAndOffset.f1);
            serializeStreamStateHandle(streamHandleAndOffset.f0, dos);
        }
        dos.writeLong(handle.getStateSize());
        dos.writeLong(handle.getCheckpointedSize());
        writeStateHandleId(handle, dos);
    } else {
        throw new IllegalStateException("Unknown KeyedStateHandle type: " + stateHandle.getClass());
    }
}

Example 3 with IncrementalRemoteKeyedStateHandle

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

the class RocksDBIncrementalRestoreOperation method restoreWithoutRescaling.

/**
 * Recovery from a single remote incremental state without rescaling.
 */
@SuppressWarnings("unchecked")
private void restoreWithoutRescaling(KeyedStateHandle keyedStateHandle) throws Exception {
    logger.info("Starting to restore from state handle: {} without rescaling.", keyedStateHandle);
    if (keyedStateHandle instanceof IncrementalRemoteKeyedStateHandle) {
        IncrementalRemoteKeyedStateHandle incrementalRemoteKeyedStateHandle = (IncrementalRemoteKeyedStateHandle) keyedStateHandle;
        restorePreviousIncrementalFilesStatus(incrementalRemoteKeyedStateHandle);
        restoreFromRemoteState(incrementalRemoteKeyedStateHandle);
    } else if (keyedStateHandle instanceof IncrementalLocalKeyedStateHandle) {
        IncrementalLocalKeyedStateHandle incrementalLocalKeyedStateHandle = (IncrementalLocalKeyedStateHandle) keyedStateHandle;
        restorePreviousIncrementalFilesStatus(incrementalLocalKeyedStateHandle);
        restoreFromLocalState(incrementalLocalKeyedStateHandle);
    } else {
        throw unexpectedStateHandleException(new Class[] { IncrementalRemoteKeyedStateHandle.class, IncrementalLocalKeyedStateHandle.class }, keyedStateHandle.getClass());
    }
    logger.info("Finished restoring from state handle: {} without rescaling.", keyedStateHandle);
}

Example 4 with IncrementalRemoteKeyedStateHandle

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

the class RocksDBIncrementalRestoreOperation method restoreWithRescaling.

/**
 * Recovery from multi incremental states with rescaling. For rescaling, this method creates a
 * temporary RocksDB instance for a key-groups shard. All contents from the temporary instance
 * are copied into the real restore instance and then the temporary instance is discarded.
 */
private void restoreWithRescaling(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
    // Prepare for restore with rescaling
    KeyedStateHandle initialHandle = RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(restoreStateHandles, keyGroupRange);
    // Init base DB instance
    if (initialHandle != null) {
        restoreStateHandles.remove(initialHandle);
        initDBWithRescaling(initialHandle);
    } else {
        this.rocksHandle.openDB();
    }
    // Transfer remaining key-groups from temporary instance into base DB
    byte[] startKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
    CompositeKeySerializationUtils.serializeKeyGroup(keyGroupRange.getStartKeyGroup(), startKeyGroupPrefixBytes);
    byte[] stopKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
    CompositeKeySerializationUtils.serializeKeyGroup(keyGroupRange.getEndKeyGroup() + 1, stopKeyGroupPrefixBytes);
    for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
        if (!(rawStateHandle instanceof IncrementalRemoteKeyedStateHandle)) {
            throw unexpectedStateHandleException(IncrementalRemoteKeyedStateHandle.class, rawStateHandle.getClass());
        }
        logger.info("Starting to restore from state handle: {} with rescaling.", rawStateHandle);
        Path temporaryRestoreInstancePath = instanceBasePath.getAbsoluteFile().toPath().resolve(UUID.randomUUID().toString());
        try (RestoredDBInstance tmpRestoreDBInfo = restoreDBInstanceFromStateHandle((IncrementalRemoteKeyedStateHandle) rawStateHandle, temporaryRestoreInstancePath);
            RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.rocksHandle.getDb(), writeBatchSize)) {
            List<ColumnFamilyDescriptor> tmpColumnFamilyDescriptors = tmpRestoreDBInfo.columnFamilyDescriptors;
            List<ColumnFamilyHandle> tmpColumnFamilyHandles = tmpRestoreDBInfo.columnFamilyHandles;
            // family handle
            for (int i = 0; i < tmpColumnFamilyDescriptors.size(); ++i) {
                ColumnFamilyHandle tmpColumnFamilyHandle = tmpColumnFamilyHandles.get(i);
                ColumnFamilyHandle targetColumnFamilyHandle = this.rocksHandle.getOrRegisterStateColumnFamilyHandle(null, tmpRestoreDBInfo.stateMetaInfoSnapshots.get(i)).columnFamilyHandle;
                try (RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(tmpRestoreDBInfo.db, tmpColumnFamilyHandle, tmpRestoreDBInfo.readOptions)) {
                    iterator.seek(startKeyGroupPrefixBytes);
                    while (iterator.isValid()) {
                        if (RocksDBIncrementalCheckpointUtils.beforeThePrefixBytes(iterator.key(), stopKeyGroupPrefixBytes)) {
                            writeBatchWrapper.put(targetColumnFamilyHandle, iterator.key(), iterator.value());
                        } else {
                            // we can just break here.
                            break;
                        }
                        iterator.next();
                    }
                }
            // releases native iterator resources
            }
            logger.info("Finished restoring from state handle: {} with rescaling.", rawStateHandle);
        } finally {
            cleanUpPathQuietly(temporaryRestoreInstancePath);
        }
    }
}

Example 5 with IncrementalRemoteKeyedStateHandle

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

the class RocksDBStateDownloaderTest method testMultiThreadRestoreCorrectly.

/**
 * Tests that download files with multi-thread correctly.
 */
@Test
public void testMultiThreadRestoreCorrectly() throws Exception {
    Random random = new Random();
    int contentNum = 6;
    byte[][] contents = new byte[contentNum][];
    for (int i = 0; i < contentNum; ++i) {
        contents[i] = new byte[random.nextInt(100000) + 1];
        random.nextBytes(contents[i]);
    }
    List<StreamStateHandle> handles = new ArrayList<>(contentNum);
    for (int i = 0; i < contentNum; ++i) {
        handles.add(new ByteStreamStateHandle(String.format("state%d", i), contents[i]));
    }
    Map<StateHandleID, StreamStateHandle> sharedStates = new HashMap<>(contentNum);
    Map<StateHandleID, StreamStateHandle> privateStates = new HashMap<>(contentNum);
    for (int i = 0; i < contentNum; ++i) {
        sharedStates.put(new StateHandleID(String.format("sharedState%d", i)), handles.get(i));
        privateStates.put(new StateHandleID(String.format("privateState%d", i)), handles.get(i));
    }
    IncrementalRemoteKeyedStateHandle incrementalKeyedStateHandle = new IncrementalRemoteKeyedStateHandle(UUID.randomUUID(), KeyGroupRange.of(0, 1), 1, sharedStates, privateStates, handles.get(0));
    Path dstPath = temporaryFolder.newFolder().toPath();
    try (RocksDBStateDownloader rocksDBStateDownloader = new RocksDBStateDownloader(5)) {
        rocksDBStateDownloader.transferAllStateDataToDirectory(incrementalKeyedStateHandle, dstPath, new CloseableRegistry());
    }
    for (int i = 0; i < contentNum; ++i) {
        assertStateContentEqual(contents[i], dstPath.resolve(String.format("sharedState%d", i)));
    }
}