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

Example 11 with KeyedStateHandle

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

the class TtlStateTestBase method testRestoreTtlAndRegisterNonTtlStateCompatFailure.

@Test(expected = StateMigrationException.class)
public void testRestoreTtlAndRegisterNonTtlStateCompatFailure() throws Exception {
    assumeThat(this, not(instanceOf(MockTtlStateTest.class)));
    initTest();
    timeProvider.time = 0;
    ctx().update(ctx().updateEmpty);
    KeyedStateHandle snapshot = sbetc.takeSnapshot();
    sbetc.createAndRestoreKeyedStateBackend(snapshot);
    sbetc.setCurrentKey("defaultKey");
    sbetc.createState(ctx().createStateDescriptor(), "");
}

Example 12 with KeyedStateHandle

use of org.apache.flink.runtime.state.KeyedStateHandle 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 13 with KeyedStateHandle

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

the class RocksDBIncrementalRestoreOperation method restore.

/**
 * Root method that branches for different implementations of {@link KeyedStateHandle}.
 */
@Override
public RocksDBRestoreResult restore() throws Exception {
    if (restoreStateHandles == null || restoreStateHandles.isEmpty()) {
        return null;
    }
    final KeyedStateHandle theFirstStateHandle = restoreStateHandles.iterator().next();
    boolean isRescaling = (restoreStateHandles.size() > 1 || !Objects.equals(theFirstStateHandle.getKeyGroupRange(), keyGroupRange));
    if (isRescaling) {
        restoreWithRescaling(restoreStateHandles);
    } else {
        restoreWithoutRescaling(theFirstStateHandle);
    }
    return new RocksDBRestoreResult(this.rocksHandle.getDb(), this.rocksHandle.getDefaultColumnFamilyHandle(), this.rocksHandle.getNativeMetricMonitor(), lastCompletedCheckpointId, backendUID, restoredSstFiles);
}

Example 14 with KeyedStateHandle

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

the class RocksDBKeyedStateBackendBuilder method getRocksDBRestoreOperation.

private RocksDBRestoreOperation getRocksDBRestoreOperation(int keyGroupPrefixBytes, CloseableRegistry cancelStreamRegistry, LinkedHashMap<String, RocksDBKeyedStateBackend.RocksDbKvStateInfo> kvStateInformation, LinkedHashMap<String, HeapPriorityQueueSnapshotRestoreWrapper<?>> registeredPQStates, RocksDbTtlCompactFiltersManager ttlCompactFiltersManager) {
    DBOptions dbOptions = optionsContainer.getDbOptions();
    if (restoreStateHandles.isEmpty()) {
        return new RocksDBNoneRestoreOperation<>(kvStateInformation, instanceRocksDBPath, dbOptions, columnFamilyOptionsFactory, nativeMetricOptions, metricGroup, ttlCompactFiltersManager, optionsContainer.getWriteBufferManagerCapacity());
    }
    KeyedStateHandle firstStateHandle = restoreStateHandles.iterator().next();
    if (firstStateHandle instanceof IncrementalKeyedStateHandle) {
        return new RocksDBIncrementalRestoreOperation<>(operatorIdentifier, keyGroupRange, keyGroupPrefixBytes, numberOfTransferingThreads, cancelStreamRegistry, userCodeClassLoader, kvStateInformation, keySerializerProvider, instanceBasePath, instanceRocksDBPath, dbOptions, columnFamilyOptionsFactory, nativeMetricOptions, metricGroup, restoreStateHandles, ttlCompactFiltersManager, writeBatchSize, optionsContainer.getWriteBufferManagerCapacity());
    } else if (priorityQueueStateType == EmbeddedRocksDBStateBackend.PriorityQueueStateType.HEAP) {
        return new RocksDBHeapTimersFullRestoreOperation<>(keyGroupRange, numberOfKeyGroups, userCodeClassLoader, kvStateInformation, registeredPQStates, createHeapQueueFactory(), keySerializerProvider, instanceRocksDBPath, dbOptions, columnFamilyOptionsFactory, nativeMetricOptions, metricGroup, restoreStateHandles, ttlCompactFiltersManager, writeBatchSize, optionsContainer.getWriteBufferManagerCapacity());
    } else {
        return new RocksDBFullRestoreOperation<>(keyGroupRange, userCodeClassLoader, kvStateInformation, keySerializerProvider, instanceRocksDBPath, dbOptions, columnFamilyOptionsFactory, nativeMetricOptions, metricGroup, restoreStateHandles, ttlCompactFiltersManager, writeBatchSize, optionsContainer.getWriteBufferManagerCapacity());
    }
}

Example 15 with KeyedStateHandle

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

the class OperatorSnapshotFinalizerTest method testRunAndExtract.

/**
 * Test that the runnable futures are executed and the result is correctly extracted.
 */
@Test
public void testRunAndExtract() throws Exception {
    Random random = new Random(0x42);
    KeyedStateHandle keyedTemplate = StateHandleDummyUtil.createNewKeyedStateHandle(new KeyGroupRange(0, 0));
    OperatorStateHandle operatorTemplate = StateHandleDummyUtil.createNewOperatorStateHandle(2, random);
    InputChannelStateHandle inputChannelTemplate = StateHandleDummyUtil.createNewInputChannelStateHandle(2, random);
    ResultSubpartitionStateHandle resultSubpartitionTemplate = StateHandleDummyUtil.createNewResultSubpartitionStateHandle(2, random);
    SnapshotResult<KeyedStateHandle> manKeyed = withLocalState(deepDummyCopy(keyedTemplate), deepDummyCopy(keyedTemplate));
    SnapshotResult<KeyedStateHandle> rawKeyed = withLocalState(deepDummyCopy(keyedTemplate), deepDummyCopy(keyedTemplate));
    SnapshotResult<OperatorStateHandle> manOper = withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
    SnapshotResult<OperatorStateHandle> rawOper = withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
    SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannel = withLocalState(singleton(deepDummyCopy(inputChannelTemplate)), singleton(deepDummyCopy(inputChannelTemplate)));
    SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartition = withLocalState(singleton(deepDummyCopy(resultSubpartitionTemplate)), singleton(deepDummyCopy(resultSubpartitionTemplate)));
    OperatorSnapshotFutures snapshotFutures = new OperatorSnapshotFutures(new PseudoNotDoneFuture<>(manKeyed), new PseudoNotDoneFuture<>(rawKeyed), new PseudoNotDoneFuture<>(manOper), new PseudoNotDoneFuture<>(rawOper), new PseudoNotDoneFuture<>(inputChannel), new PseudoNotDoneFuture<>(resultSubpartition));
    for (Future<?> f : snapshotFutures.getAllFutures()) {
        assertFalse(f.isDone());
    }
    OperatorSnapshotFinalizer finalizer = new OperatorSnapshotFinalizer(snapshotFutures);
    for (Future<?> f : snapshotFutures.getAllFutures()) {
        assertTrue(f.isDone());
    }
    Map<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> map = new HashMap<>();
    map.put(manKeyed, headExtractor(OperatorSubtaskState::getManagedKeyedState));
    map.put(rawKeyed, headExtractor(OperatorSubtaskState::getRawKeyedState));
    map.put(manOper, headExtractor(OperatorSubtaskState::getManagedOperatorState));
    map.put(rawOper, headExtractor(OperatorSubtaskState::getRawOperatorState));
    map.put(inputChannel, OperatorSubtaskState::getInputChannelState);
    map.put(resultSubpartition, OperatorSubtaskState::getResultSubpartitionState);
    for (Map.Entry<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> e : map.entrySet()) {
        assertEquals(e.getKey().getJobManagerOwnedSnapshot(), e.getValue().apply(finalizer.getJobManagerOwnedState()));
    }
    for (Map.Entry<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> e : map.entrySet()) {
        assertEquals(e.getKey().getTaskLocalSnapshot(), e.getValue().apply(finalizer.getTaskLocalState()));
    }
}