org.apache.flink.runtime.state.KeyedStateHandle Java Examples

@VisibleForTesting
public static KeyedStateHandle deserializeKeyedStateHandle(DataInputStream dis) throws IOException {
	final int type = dis.readByte();
	if (NULL_HANDLE == type) {
		return null;
	} else if (KEY_GROUPS_HANDLE == type) {
		int startKeyGroup = dis.readInt();
		int numKeyGroups = dis.readInt();
		KeyGroupRange keyGroupRange = KeyGroupRange.of(startKeyGroup, startKeyGroup + numKeyGroups - 1);
		long[] offsets = new long[numKeyGroups];
		for (int i = 0; i < numKeyGroups; ++i) {
			offsets[i] = dis.readLong();
		}
		KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(
			keyGroupRange, offsets);
		StreamStateHandle stateHandle = deserializeStreamStateHandle(dis);
		return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
	} else {
		throw new IllegalStateException("Reading invalid KeyedStateHandle, type: " + type);
	}
}
 

@Nonnull
@Override
@SuppressWarnings("unchecked")
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws IOException {

	long startTime = System.currentTimeMillis();

	final RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		snapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	snapshotStrategy.logSyncCompleted(streamFactory, startTime);
	return snapshotRunner;
}
 

private void initDBWithRescaling(KeyedStateHandle initialHandle) throws Exception {

		assert (initialHandle instanceof IncrementalRemoteKeyedStateHandle);

		// 1. Restore base DB from selected initial handle
		restoreFromRemoteState((IncrementalRemoteKeyedStateHandle) initialHandle);

		// 2. Clip the base DB instance
		try {
			RocksDBIncrementalCheckpointUtils.clipDBWithKeyGroupRange(
				db,
				columnFamilyHandles,
				keyGroupRange,
				initialHandle.getKeyGroupRange(),
				keyGroupPrefixBytes);
		} catch (RocksDBException e) {
			String errMsg = "Failed to clip DB after initialization.";
			LOG.error(errMsg, e);
			throw new BackendBuildingException(errMsg, e);
		}
	}
 

/**
 * Choose the best state handle according to the {@link #STATE_HANDLE_EVALUATOR}
 * to init the initial db.
 *
 * @param restoreStateHandles The candidate state handles.
 * @param targetKeyGroupRange The target key group range.
 * @return The best candidate or null if no candidate was a good fit.
 */
@Nullable
public static KeyedStateHandle chooseTheBestStateHandleForInitial(
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	@Nonnull KeyGroupRange targetKeyGroupRange) {

	KeyedStateHandle bestStateHandle = null;
	double bestScore = 0;
	for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
		double handleScore = STATE_HANDLE_EVALUATOR.apply(rawStateHandle, targetKeyGroupRange);
		if (handleScore > bestScore) {
			bestStateHandle = rawStateHandle;
			bestScore = handleScore;
		}
	}

	return bestStateHandle;
}
 

@VisibleForTesting
public 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;

		dos.writeByte(KEY_GROUPS_HANDLE);
		dos.writeInt(keyGroupsStateHandle.getKeyGroupRange().getStartKeyGroup());
		dos.writeInt(keyGroupsStateHandle.getKeyGroupRange().getNumberOfKeyGroups());
		for (int keyGroup : keyGroupsStateHandle.getKeyGroupRange()) {
			dos.writeLong(keyGroupsStateHandle.getOffsetForKeyGroup(keyGroup));
		}
		serializeStreamStateHandle(keyGroupsStateHandle.getDelegateStateHandle(), dos);
	} else {
		throw new IllegalStateException("Unknown KeyedStateHandle type: " + stateHandle.getClass());
	}
}
 

public MockKeyedStateBackendBuilder(
	TaskKvStateRegistry kvStateRegistry,
	TypeSerializer<K> keySerializer,
	ClassLoader userCodeClassLoader,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	StreamCompressionDecorator keyGroupCompressionDecorator,
	CloseableRegistry cancelStreamRegistry) {
	super(
		kvStateRegistry,
		keySerializer,
		userCodeClassLoader,
		numberOfKeyGroups,
		keyGroupRange,
		executionConfig,
		ttlTimeProvider,
		stateHandles,
		keyGroupCompressionDecorator,
		cancelStreamRegistry);
}
 

@Override
protected SnapshotResult<KeyedStateHandle> callInternal() throws Exception {
	final KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupRange);
	final CheckpointStreamWithResultProvider checkpointStreamWithResultProvider =
		checkpointStreamSupplier.get();

	snapshotCloseableRegistry.registerCloseable(checkpointStreamWithResultProvider);
	writeSnapshotToOutputStream(checkpointStreamWithResultProvider, keyGroupRangeOffsets);

	if (snapshotCloseableRegistry.unregisterCloseable(checkpointStreamWithResultProvider)) {
		return CheckpointStreamWithResultProvider.toKeyedStateHandleSnapshotResult(
			checkpointStreamWithResultProvider.closeAndFinalizeCheckpointStreamResult(),
			keyGroupRangeOffsets);
	} else {
		throw new IOException("Stream is already unregistered/closed.");
	}
}
 

public <K> HeapKeyedStateBackend<K> createKeyedBackend(
	TypeSerializer<K> keySerializer,
	Collection<KeyedStateHandle> stateHandles) throws Exception {
	final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 15);
	final int numKeyGroups = keyGroupRange.getNumberOfKeyGroups();
	ExecutionConfig executionConfig = new ExecutionConfig();

	return new HeapKeyedStateBackendBuilder<>(
		mock(TaskKvStateRegistry.class),
		keySerializer,
		HeapStateBackendTestBase.class.getClassLoader(),
		numKeyGroups,
		keyGroupRange,
		executionConfig,
		TtlTimeProvider.DEFAULT,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		TestLocalRecoveryConfig.disabled(),
		new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128),
		async,
		new CloseableRegistry()).build();
}
 

/**
 * Determine the subset of {@link KeyGroupsStateHandle KeyGroupsStateHandles} with correct
 * key group index for the given subtask {@link KeyGroupRange}.
 *
 * <p>This is publicly visible to be used in tests.
 */
public static List<KeyedStateHandle> getKeyedStateHandles(
	Collection<? extends KeyedStateHandle> keyedStateHandles,
	KeyGroupRange subtaskKeyGroupRange) {

	List<KeyedStateHandle> subtaskKeyedStateHandles = new ArrayList<>(keyedStateHandles.size());

	for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
		KeyedStateHandle intersectedKeyedStateHandle = keyedStateHandle.getIntersection(subtaskKeyGroupRange);

		if (intersectedKeyedStateHandle != null) {
			subtaskKeyedStateHandles.add(intersectedKeyedStateHandle);
		}
	}

	return subtaskKeyedStateHandles;
}
 

HeapRestoreOperation(
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	StateSerializerProvider<K> keySerializerProvider,
	ClassLoader userCodeClassLoader,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	CloseableRegistry cancelStreamRegistry,
	HeapPriorityQueueSetFactory priorityQueueSetFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	int numberOfKeyGroups,
	HeapSnapshotStrategy<K> snapshotStrategy,
	HeapKeyedStateBackend<K> backend) {
	this.restoreStateHandles = restoreStateHandles;
	this.keySerializerProvider = keySerializerProvider;
	this.userCodeClassLoader = userCodeClassLoader;
	this.registeredKVStates = registeredKVStates;
	this.registeredPQStates = registeredPQStates;
	this.cancelStreamRegistry = cancelStreamRegistry;
	this.priorityQueueSetFactory = priorityQueueSetFactory;
	this.keyGroupRange = keyGroupRange;
	this.numberOfKeyGroups = numberOfKeyGroups;
	this.snapshotStrategy = snapshotStrategy;
	this.backend = backend;
}
 

/**
 * Triggers an asynchronous snapshot of the keyed state backend from RocksDB. This snapshot can be canceled and
 * is also stopped when the backend is closed through {@link #dispose()}. For each backend, this method must always
 * be called by the same thread.
 *
 * @param checkpointId  The Id of the checkpoint.
 * @param timestamp     The timestamp of the checkpoint.
 * @param streamFactory The factory that we can use for writing our state to streams.
 * @param checkpointOptions Options for how to perform this checkpoint.
 * @return Future to the state handle of the snapshot data.
 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
 */
@Nonnull
@Override
public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
	final long checkpointId,
	final long timestamp,
	@Nonnull final CheckpointStreamFactory streamFactory,
	@Nonnull CheckpointOptions checkpointOptions) throws Exception {

	long startTime = System.currentTimeMillis();

	// flush everything into db before taking a snapshot
	writeBatchWrapper.flush();

	RocksDBSnapshotStrategyBase<K> chosenSnapshotStrategy =
			checkpointOptions.getCheckpointType().isSavepoint() ? savepointSnapshotStrategy : checkpointSnapshotStrategy;

	RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunner =
		chosenSnapshotStrategy.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);

	chosenSnapshotStrategy.logSyncCompleted(streamFactory, startTime);

	return snapshotRunner;
}
 

/**
 * Determine the subset of {@link KeyGroupsStateHandle KeyGroupsStateHandles} with correct
 * key group index for the given subtask {@link KeyGroupRange}.
 *
 * <p>This is publicly visible to be used in tests.
 */
public static List<KeyedStateHandle> getKeyedStateHandles(
	Collection<? extends KeyedStateHandle> keyedStateHandles,
	KeyGroupRange subtaskKeyGroupRange) {

	List<KeyedStateHandle> subtaskKeyedStateHandles = new ArrayList<>(keyedStateHandles.size());

	for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
		KeyedStateHandle intersectedKeyedStateHandle = keyedStateHandle.getIntersection(subtaskKeyGroupRange);

		if (intersectedKeyedStateHandle != null) {
			subtaskKeyedStateHandles.add(intersectedKeyedStateHandle);
		}
	}

	return subtaskKeyedStateHandles;
}
 

HeapRestoreOperation(
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	StateSerializerProvider<K> keySerializerProvider,
	ClassLoader userCodeClassLoader,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	CloseableRegistry cancelStreamRegistry,
	HeapPriorityQueueSetFactory priorityQueueSetFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	int numberOfKeyGroups,
	HeapSnapshotStrategy<K> snapshotStrategy,
	InternalKeyContext<K> keyContext) {
	this.restoreStateHandles = restoreStateHandles;
	this.keySerializerProvider = keySerializerProvider;
	this.userCodeClassLoader = userCodeClassLoader;
	this.registeredKVStates = registeredKVStates;
	this.registeredPQStates = registeredPQStates;
	this.cancelStreamRegistry = cancelStreamRegistry;
	this.priorityQueueSetFactory = priorityQueueSetFactory;
	this.keyGroupRange = keyGroupRange;
	this.numberOfKeyGroups = numberOfKeyGroups;
	this.snapshotStrategy = snapshotStrategy;
	this.keyContext = keyContext;
}
 

public <K> HeapKeyedStateBackend<K> createKeyedBackend(
	TypeSerializer<K> keySerializer,
	Collection<KeyedStateHandle> stateHandles) throws Exception {
	final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 15);
	final int numKeyGroups = keyGroupRange.getNumberOfKeyGroups();
	ExecutionConfig executionConfig = new ExecutionConfig();

	return new HeapKeyedStateBackendBuilder<>(
		mock(TaskKvStateRegistry.class),
		keySerializer,
		HeapStateBackendTestBase.class.getClassLoader(),
		numKeyGroups,
		keyGroupRange,
		executionConfig,
		TtlTimeProvider.DEFAULT,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		TestLocalRecoveryConfig.disabled(),
		new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128),
		async,
		new CloseableRegistry()).build();
}
 

/**
 * Recovery from a single remote incremental state without rescaling.
 */
private void restoreWithoutRescaling(KeyedStateHandle keyedStateHandle) throws Exception {
	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 new BackendBuildingException("Unexpected state handle type, " +
			"expected " + IncrementalRemoteKeyedStateHandle.class + " or " + IncrementalLocalKeyedStateHandle.class +
			", but found " + keyedStateHandle.getClass());
	}
}
 

public OperatorSubtaskState(
	@Nonnull StateObjectCollection<OperatorStateHandle> managedOperatorState,
	@Nonnull StateObjectCollection<OperatorStateHandle> rawOperatorState,
	@Nonnull StateObjectCollection<KeyedStateHandle> managedKeyedState,
	@Nonnull StateObjectCollection<KeyedStateHandle> rawKeyedState) {

	this.managedOperatorState = Preconditions.checkNotNull(managedOperatorState);
	this.rawOperatorState = Preconditions.checkNotNull(rawOperatorState);
	this.managedKeyedState = Preconditions.checkNotNull(managedKeyedState);
	this.rawKeyedState = Preconditions.checkNotNull(rawKeyedState);

	long calculateStateSize = managedOperatorState.getStateSize();
	calculateStateSize += rawOperatorState.getStateSize();
	calculateStateSize += managedKeyedState.getStateSize();
	calculateStateSize += rawKeyedState.getStateSize();
	stateSize = calculateStateSize;
}
 

void dispose() throws Exception {
	disposeKeyedStateBackend();
	for (KeyedStateHandle snapshot : snapshots) {
		snapshot.discardState();
	}
	snapshots.clear();
	sharedStateRegistry.close();
}
 

public static void compareKeyedState(
		Collection<KeyGroupsStateHandle> expectPartitionedKeyGroupState,
		Collection<? extends KeyedStateHandle> actualPartitionedKeyGroupState) throws Exception {

	KeyGroupsStateHandle expectedHeadOpKeyGroupStateHandle = expectPartitionedKeyGroupState.iterator().next();
	int expectedTotalKeyGroups = expectedHeadOpKeyGroupStateHandle.getKeyGroupRange().getNumberOfKeyGroups();
	int actualTotalKeyGroups = 0;
	for(KeyedStateHandle keyedStateHandle: actualPartitionedKeyGroupState) {
		assertTrue(keyedStateHandle instanceof KeyGroupsStateHandle);

		actualTotalKeyGroups += keyedStateHandle.getKeyGroupRange().getNumberOfKeyGroups();
	}

	assertEquals(expectedTotalKeyGroups, actualTotalKeyGroups);

	try (FSDataInputStream inputStream = expectedHeadOpKeyGroupStateHandle.openInputStream()) {
		for (int groupId : expectedHeadOpKeyGroupStateHandle.getKeyGroupRange()) {
			long offset = expectedHeadOpKeyGroupStateHandle.getOffsetForKeyGroup(groupId);
			inputStream.seek(offset);
			int expectedKeyGroupState =
					InstantiationUtil.deserializeObject(inputStream, Thread.currentThread().getContextClassLoader());
			for (KeyedStateHandle oneActualKeyedStateHandle : actualPartitionedKeyGroupState) {

				assertTrue(oneActualKeyedStateHandle instanceof KeyGroupsStateHandle);

				KeyGroupsStateHandle oneActualKeyGroupStateHandle = (KeyGroupsStateHandle) oneActualKeyedStateHandle;
				if (oneActualKeyGroupStateHandle.getKeyGroupRange().contains(groupId)) {
					long actualOffset = oneActualKeyGroupStateHandle.getOffsetForKeyGroup(groupId);
					try (FSDataInputStream actualInputStream = oneActualKeyGroupStateHandle.openInputStream()) {
						actualInputStream.seek(actualOffset);
						int actualGroupState = InstantiationUtil.
								deserializeObject(actualInputStream, Thread.currentThread().getContextClassLoader());
						assertEquals(expectedKeyGroupState, actualGroupState);
					}
				}
			}
		}
	}
}
 

public RocksDBNoneRestoreOperation(
	KeyGroupRange keyGroupRange,
	int keyGroupPrefixBytes,
	int numberOfTransferringThreads,
	CloseableRegistry cancelStreamRegistry,
	ClassLoader userCodeClassLoader,
	Map<String, RocksDbKvStateInfo> kvStateInformation,
	StateSerializerProvider<K> keySerializerProvider,
	File instanceBasePath,
	File instanceRocksDBPath,
	DBOptions dbOptions,
	Function<String, ColumnFamilyOptions> columnFamilyOptionsFactory,
	RocksDBNativeMetricOptions nativeMetricOptions,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	@Nonnull RocksDbTtlCompactFiltersManager ttlCompactFiltersManager
) {
	super(keyGroupRange,
		keyGroupPrefixBytes,
		numberOfTransferringThreads,
		cancelStreamRegistry,
		userCodeClassLoader,
		kvStateInformation,
		keySerializerProvider,
		instanceBasePath,
		instanceRocksDBPath,
		dbOptions,
		columnFamilyOptionsFactory,
		nativeMetricOptions,
		metricGroup,
		restoreStateHandles,
		ttlCompactFiltersManager);
}
 

public OperatorSnapshotFinalizer(
	@Nonnull OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {

	SnapshotResult<KeyedStateHandle> keyedManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());

	SnapshotResult<KeyedStateHandle> keyedRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());

	SnapshotResult<OperatorStateHandle> operatorManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());

	SnapshotResult<OperatorStateHandle> operatorRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());

	jobManagerOwnedState = new OperatorSubtaskState(
		operatorManaged.getJobManagerOwnedSnapshot(),
		operatorRaw.getJobManagerOwnedSnapshot(),
		keyedManaged.getJobManagerOwnedSnapshot(),
		keyedRaw.getJobManagerOwnedSnapshot()
	);

	taskLocalState = new OperatorSubtaskState(
		operatorManaged.getTaskLocalSnapshot(),
		operatorRaw.getTaskLocalSnapshot(),
		keyedManaged.getTaskLocalSnapshot(),
		keyedRaw.getTaskLocalSnapshot()
	);
}
 

public RocksDBFullRestoreOperation(
	KeyGroupRange keyGroupRange,
	int keyGroupPrefixBytes,
	int numberOfTransferringThreads,
	CloseableRegistry cancelStreamRegistry,
	ClassLoader userCodeClassLoader,
	Map<String, RocksDbKvStateInfo> kvStateInformation,
	StateSerializerProvider<K> keySerializerProvider,
	File instanceBasePath,
	File instanceRocksDBPath,
	DBOptions dbOptions,
	Function<String, ColumnFamilyOptions> columnFamilyOptionsFactory,
	RocksDBNativeMetricOptions nativeMetricOptions,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	@Nonnull RocksDbTtlCompactFiltersManager ttlCompactFiltersManager) {
	super(
		keyGroupRange,
		keyGroupPrefixBytes,
		numberOfTransferringThreads,
		cancelStreamRegistry,
		userCodeClassLoader,
		kvStateInformation,
		keySerializerProvider,
		instanceBasePath,
		instanceRocksDBPath,
		dbOptions,
		columnFamilyOptionsFactory,
		nativeMetricOptions,
		metricGroup,
		restoreStateHandles,
		ttlCompactFiltersManager);
}
 

private Tuple2<List<KeyedStateHandle>, List<KeyedStateHandle>> reAssignSubKeyedStates(
		OperatorState operatorState,
		List<KeyGroupRange> keyGroupPartitions,
		int subTaskIndex,
		int newParallelism,
		int oldParallelism) {

	List<KeyedStateHandle> subManagedKeyedState;
	List<KeyedStateHandle> subRawKeyedState;

	if (newParallelism == oldParallelism) {
		if (operatorState.getState(subTaskIndex) != null) {
			subManagedKeyedState = operatorState.getState(subTaskIndex).getManagedKeyedState().asList();
			subRawKeyedState = operatorState.getState(subTaskIndex).getRawKeyedState().asList();
		} else {
			subManagedKeyedState = Collections.emptyList();
			subRawKeyedState = Collections.emptyList();
		}
	} else {
		subManagedKeyedState = getManagedKeyedStateHandles(operatorState, keyGroupPartitions.get(subTaskIndex));
		subRawKeyedState = getRawKeyedStateHandles(operatorState, keyGroupPartitions.get(subTaskIndex));
	}

	if (subManagedKeyedState.isEmpty() && subRawKeyedState.isEmpty()) {
		return new Tuple2<>(Collections.emptyList(), Collections.emptyList());
	} else {
		return new Tuple2<>(subManagedKeyedState, subRawKeyedState);
	}
}
 

@Test
public void testDismissingSnapshot() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		snapshot.cancel(true);
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

@Test
public void testSnapshotChangeRestore() throws Exception {
	initTest();

	timeProvider.time = 0;
	sbetc.setCurrentKey("k1");
	ctx().update(ctx().updateEmpty);

	timeProvider.time = 50;
	sbetc.setCurrentKey("k1");
	ctx().update(ctx().updateUnexpired);

	timeProvider.time = 100;
	sbetc.setCurrentKey("k2");
	ctx().update(ctx().updateEmpty);

	KeyedStateHandle snapshot = sbetc.takeSnapshot();

	timeProvider.time = 170;
	sbetc.setCurrentKey("k1");
	ctx().update(ctx().updateExpired);
	sbetc.setCurrentKey("k2");
	ctx().update(ctx().updateUnexpired);

	restoreSnapshot(snapshot, StateBackendTestContext.NUMBER_OF_KEY_GROUPS);

	timeProvider.time = 180;
	sbetc.setCurrentKey("k1");
	assertEquals(EXPIRED_UNAVAIL, ctx().emptyValue, ctx().get());
	sbetc.setCurrentKey("k2");
	assertEquals(UNEXPIRED_AVAIL, ctx().getUpdateEmpty, ctx().get());
}
 

private static void registerSharedState(
	SharedStateRegistry sharedStateRegistry,
	Iterable<KeyedStateHandle> stateHandles) {
	for (KeyedStateHandle stateHandle : stateHandles) {
		if (stateHandle != null) {
			stateHandle.registerSharedStates(sharedStateRegistry);
		}
	}
}
 

/**
 * For convenience because the size of the collections is typically 0 or 1. Null values are translated into empty
 * Collections.
 */
public OperatorSubtaskState(
	@Nullable OperatorStateHandle managedOperatorState,
	@Nullable OperatorStateHandle rawOperatorState,
	@Nullable KeyedStateHandle managedKeyedState,
	@Nullable KeyedStateHandle rawKeyedState) {

	this(
		singletonOrEmptyOnNull(managedOperatorState),
		singletonOrEmptyOnNull(rawOperatorState),
		singletonOrEmptyOnNull(managedKeyedState),
		singletonOrEmptyOnNull(rawKeyedState));
}
 

@Test
public void testDismissingSnapshot() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		snapshot.cancel(true);
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

public RocksDBNoneRestoreOperation(
	KeyGroupRange keyGroupRange,
	int keyGroupPrefixBytes,
	int numberOfTransferringThreads,
	CloseableRegistry cancelStreamRegistry,
	ClassLoader userCodeClassLoader,
	Map<String, RocksDbKvStateInfo> kvStateInformation,
	StateSerializerProvider<K> keySerializerProvider,
	File instanceBasePath,
	File instanceRocksDBPath,
	DBOptions dbOptions,
	Function<String, ColumnFamilyOptions> columnFamilyOptionsFactory,
	RocksDBNativeMetricOptions nativeMetricOptions,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	@Nonnull RocksDbTtlCompactFiltersManager ttlCompactFiltersManager
) {
	super(keyGroupRange,
		keyGroupPrefixBytes,
		numberOfTransferringThreads,
		cancelStreamRegistry,
		userCodeClassLoader,
		kvStateInformation,
		keySerializerProvider,
		instanceBasePath,
		instanceRocksDBPath,
		dbOptions,
		columnFamilyOptionsFactory,
		nativeMetricOptions,
		metricGroup,
		restoreStateHandles,
		ttlCompactFiltersManager);
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws IOException {
	throw new SuccessException();
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	throw new UnsupportedOperationException();
}