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

@Test
public void testIteratorRemovesFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 1);

	MultiStateKeyIterator<Integer> iterator =
		new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	int key = iterator.next();
	Assert.assertEquals("Unexpected keys pulled from state backend", 1, key);

	iterator.remove();
	Assert.assertFalse("Failed to drop key from all descriptors in state backend", iterator.hasNext());

	for (StateDescriptor<?, ?> descriptor : descriptors) {
		Assert.assertEquals(
			"Failed to drop key for state descriptor",
			0,
			keyedStateBackend.getKeys(descriptor.getName(), VoidNamespace.INSTANCE).count());
	}
}
 

@Test
public void testIteratorPullsKeyFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 2);

	MultiStateKeyIterator<Integer> iterator = new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	List<Integer> keys = new ArrayList<>();

	while (iterator.hasNext()) {
		keys.add(iterator.next());
	}

	Assert.assertEquals("Unexpected number of keys", 2, keys.size());
	Assert.assertEquals("Unexpected keys found", Arrays.asList(1, 2), keys);
}
 

public static KeyedStateBackend<ByteBuffer> createStateBackend() throws Exception {
  MemoryStateBackend backend = new MemoryStateBackend();

  AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend =
      backend.createKeyedStateBackend(
          new DummyEnvironment("test", 1, 0),
          new JobID(),
          "test_op",
          new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
          2,
          new KeyGroupRange(0, 1),
          new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
          TtlTimeProvider.DEFAULT,
          null,
          Collections.emptyList(),
          new CloseableRegistry());

  changeKey(keyedStateBackend);

  return keyedStateBackend;
}
 

@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) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 

@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 {
	return spy(delegate.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry));
}
 

private static AbstractKeyedStateBackend<Integer> createKeyedStateBackend() {
	MockStateBackend backend = new MockStateBackend();

	return backend.createKeyedStateBackend(
		new DummyEnvironment(),
		new JobID(),
		"mock-backend",
		IntSerializer.INSTANCE,
		129,
		KeyGroupRange.of(0, 128),
		null,
		TtlTimeProvider.DEFAULT,
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
		Collections.emptyList(),
		new CloseableRegistry());
}
 

@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) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 

@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 {
	return spy(delegate.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry));
}
 

@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 {
	return spy(delegate.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry));
}
 

@Test
public void testIteratorPullsKeyFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 2);

	MultiStateKeyIterator<Integer> iterator = new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	List<Integer> keys = new ArrayList<>();

	while (iterator.hasNext()) {
		keys.add(iterator.next());
	}

	Assert.assertEquals("Unexpected number of keys", 2, keys.size());
	Assert.assertEquals("Unexpected keys found", Arrays.asList(1, 2), keys);
}
 

@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) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 

@Test
public void testIteratorRemovesFromAllDescriptors() throws Exception {
	AbstractKeyedStateBackend<Integer> keyedStateBackend = createKeyedStateBackend();

	setKey(keyedStateBackend, descriptors.get(0), 1);
	setKey(keyedStateBackend, descriptors.get(1), 1);

	MultiStateKeyIterator<Integer> iterator =
		new MultiStateKeyIterator<>(descriptors, keyedStateBackend);

	int key = iterator.next();
	Assert.assertEquals("Unexpected keys pulled from state backend", 1, key);

	iterator.remove();
	Assert.assertFalse("Failed to drop key from all descriptors in state backend", iterator.hasNext());

	for (StateDescriptor<?, ?> descriptor : descriptors) {
		Assert.assertEquals(
			"Failed to drop key for state descriptor",
			0,
			keyedStateBackend.getKeys(descriptor.getName(), VoidNamespace.INSTANCE).count());
	}
}
 

protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
	AbstractKeyedStateBackend<K> keyedStatedBackend,
	KeyContext keyContext, //the operator
	Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {

	if (keyedStatedBackend == null) {
		return null;
	}

	final KeyGroupRange keyGroupRange = keyedStatedBackend.getKeyGroupRange();

	final InternalTimeServiceManager<K> timeServiceManager = new InternalTimeServiceManager<>(
		keyGroupRange,
		keyContext,
		keyedStatedBackend,
		processingTimeService,
		keyedStatedBackend.requiresLegacySynchronousTimerSnapshots());

	// and then initialize the timer services
	for (KeyGroupStatePartitionStreamProvider streamProvider : rawKeyedStates) {
		int keyGroupIdx = streamProvider.getKeyGroupId();

		Preconditions.checkArgument(keyGroupRange.contains(keyGroupIdx),
			"Key Group " + keyGroupIdx + " does not belong to the local range.");

		timeServiceManager.restoreStateForKeyGroup(
			streamProvider.getStream(),
			keyGroupIdx, environment.getUserClassLoader());
	}

	return timeServiceManager;
}
 

@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 BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		taskStateManager.createLocalRecoveryConfig(),
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@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 {

	return backend.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		this.ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry);
}
 

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
	ExecutionConfig config = new ExecutionConfig();

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<ListState<String>>() {

		@Override
		public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			ListStateDescriptor<String> descr =
					(ListStateDescriptor<String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	return operatorMock;
}
 

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
	ExecutionConfig config = new ExecutionConfig();

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<ListState<String>>() {

		@Override
		public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			ListStateDescriptor<String> descr =
					(ListStateDescriptor<String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	return operatorMock;
}
 

private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				ProcessingTimeService processingTimeService,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) context.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) context.getRawKeyedStateInputs();
	super.initializeState(context);
}
 

@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

		ValueState<String> state =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}
 

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception {

	AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
	ExecutionConfig config = new ExecutionConfig();

	KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);

	DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);

	when(operatorMock.getExecutionConfig()).thenReturn(config);

	doAnswer(new Answer<MapState<Integer, String>>() {

		@Override
		public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
			MapStateDescriptor<Integer, String> descr =
					(MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2];

			AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
				new DummyEnvironment("test_task", 1, 0),
				new JobID(),
				"test_op",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());
			backend.setCurrentKey(0);
			return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
		}
	}).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));

	when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
	when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
	return operatorMock;
}
 

StreamOperatorStateContextImpl(
	boolean restored,
	OperatorStateBackend operatorStateBackend,
	AbstractKeyedStateBackend<?> keyedStateBackend,
	InternalTimeServiceManager<?> internalTimeServiceManager,
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs,
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs) {

	this.restored = restored;
	this.operatorStateBackend = operatorStateBackend;
	this.keyedStateBackend = keyedStateBackend;
	this.internalTimeServiceManager = internalTimeServiceManager;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

@SuppressWarnings({"unchecked", "rawtypes"})
public void setCurrentKey(Object key) {
	if (keyedStateBackend != null) {
		try {
			// need to work around type restrictions
			@SuppressWarnings("unchecked,rawtypes")
			AbstractKeyedStateBackend rawBackend = (AbstractKeyedStateBackend) keyedStateBackend;

			rawBackend.setCurrentKey(key);
		} catch (Exception e) {
			throw new RuntimeException("Exception occurred while setting the current key context.", e);
		}
	}
}
 

@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();
}
 

@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

		ValueState<String> state =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}
 

private static void setKey(
	AbstractKeyedStateBackend<Integer> backend,
	ValueStateDescriptor<Integer> descriptor,
	int key) throws Exception {
	backend.setCurrentKey(key);
	backend
		.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descriptor)
		.update(0);
}
 

@Override
public void initializeState(StateInitializationContext controller) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) controller.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) controller.getRawKeyedStateInputs();
	super.initializeState(controller);
}
 

@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final Environment environment = new SavepointEnvironment
		.Builder(getRuntimeContext(), split.getNumKeyGroups())
		.setSubtaskIndex(split.getSplitNumber())
		.setPrioritizedOperatorSubtaskState(split.getPrioritizedOperatorSubtaskState())
		.build();

	final StreamOperatorStateContext context = getStreamOperatorStateContext(environment);

	keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();

	final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
	SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);
	FunctionUtils.setFunctionRuntimeContext(userFunction, ctx);

	keys = getKeyIterator(ctx);

	final InternalTimerService<VoidNamespace> timerService = restoreTimerService(context);
	try {
		this.ctx = new Context(keyedStateBackend, timerService);
	} catch (Exception e) {
		throw new IOException("Failed to restore timer state", e);
	}
}
 

@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 {

	return backend.createKeyedStateBackend(
		env,
		jobID,
		operatorIdentifier,
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		kvStateRegistry,
		this.ttlTimeProvider,
		metricGroup,
		stateHandles,
		cancelStreamRegistry);
}
 

protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
	AbstractKeyedStateBackend<K> keyedStatedBackend,
	KeyContext keyContext, //the operator
	ProcessingTimeService processingTimeService,
	Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {

	if (keyedStatedBackend == null) {
		return null;
	}

	final KeyGroupRange keyGroupRange = keyedStatedBackend.getKeyGroupRange();

	final InternalTimeServiceManager<K> timeServiceManager = new InternalTimeServiceManager<>(
		keyGroupRange,
		keyContext,
		keyedStatedBackend,
		processingTimeService,
		keyedStatedBackend.requiresLegacySynchronousTimerSnapshots());

	// and then initialize the timer services
	for (KeyGroupStatePartitionStreamProvider streamProvider : rawKeyedStates) {
		int keyGroupIdx = streamProvider.getKeyGroupId();

		Preconditions.checkArgument(keyGroupRange.contains(keyGroupIdx),
			"Key Group " + keyGroupIdx + " does not belong to the local range.");

		timeServiceManager.restoreStateForKeyGroup(
			streamProvider.getStream(),
			keyGroupIdx, environment.getUserClassLoader());
	}

	return timeServiceManager;
}