Java Code Examples for org.apache.flink.api.common.state.ValueState#update()

/**
 * Processes element to deduplicate on keys, sends current element as last row, retracts previous element if
 * needed.
 *
 * @param currentRow latest row received by deduplicate function
 * @param generateRetraction whether need to send retract message to downstream
 * @param state state of function
 * @param out underlying collector
 * @throws Exception
 */
static void processLastRow(BaseRow currentRow, boolean generateRetraction, ValueState<BaseRow> state,
		Collector<BaseRow> out) throws Exception {
	// Check message should be accumulate
	Preconditions.checkArgument(BaseRowUtil.isAccumulateMsg(currentRow));
	if (generateRetraction) {
		// state stores complete row if generateRetraction is true
		BaseRow preRow = state.value();
		state.update(currentRow);
		if (preRow != null) {
			preRow.setHeader(BaseRowUtil.RETRACT_MSG);
			out.collect(preRow);
		}
	}
	out.collect(currentRow);
}
 

@Test
public void testValueStateWorkWithTtl() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	try {
		ValueStateDescriptor<MutableLong> kvId = new ValueStateDescriptor<>("id", MutableLong.class);
		kvId.enableTimeToLive(StateTtlConfig.newBuilder(Time.seconds(1)).build());

		ValueState<MutableLong> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
		backend.setCurrentKey(1);
		state.update(new MutableLong());
		state.value();
	} finally {
		backend.close();
		backend.dispose();
	}
}
 

@Test
public void testValueStateWorkWithTtl() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	try {
		ValueStateDescriptor<MutableLong> kvId = new ValueStateDescriptor<>("id", MutableLong.class);
		kvId.enableTimeToLive(StateTtlConfig.newBuilder(Time.seconds(1)).build());

		ValueState<MutableLong> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
		backend.setCurrentKey(1);
		state.update(new MutableLong());
		state.value();
	} finally {
		backend.close();
		backend.dispose();
	}
}
 

@Override
public Tuple2<Long, Long> map(Tuple2<Long, Long> value) throws Exception {
	ValueState<Long> state = getRuntimeContext().getState(
		new ValueStateDescriptor<>("testState", new TestSerializer()));

	state.update(value.f1);
	return value;
}
 

private void handleValue(
	Object value,
	Collector<String> out,
	TimerService timerService,
	int channel) throws IOException {
	final ValueState<String> state = getRuntimeContext().getState(this.state);
	if (state.value() == null) {
		out.collect("INPUT" + channel + ":" + value);
		state.update(String.valueOf(value));
		timerService.registerEventTimeTimer(timerService.currentWatermark() + 5);
	} else {
		state.clear();
		timerService.deleteEventTimeTimer(timerService.currentWatermark() + 4);
	}
}
 

@Test
public void testSetTtlTimeProvider() throws Exception {
	AbstractStreamOperator<Integer> operator = new AbstractStreamOperator<Integer>() {};
	try (AbstractStreamOperatorTestHarness<Integer> result = new AbstractStreamOperatorTestHarness<>(
			operator,
			1,
			1,
			0)) {

		result.config.setStateKeySerializer(IntSerializer.INSTANCE);

		Time timeToLive = Time.hours(1);
		result.initializeState(new OperatorSubtaskState());
		result.open();

		ValueStateDescriptor<Integer> stateDescriptor = new ValueStateDescriptor<>("test", IntSerializer.INSTANCE);
		stateDescriptor.enableTimeToLive(StateTtlConfig.newBuilder(timeToLive).build());
		KeyedStateBackend<Integer> keyedStateBackend = operator.getKeyedStateBackend();
		ValueState<Integer> state = keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);

		int expectedValue = 42;
		keyedStateBackend.setCurrentKey(1);
		result.setStateTtlProcessingTime(0L);
		state.update(expectedValue);
		Assert.assertEquals(expectedValue, (int) state.value());
		result.setStateTtlProcessingTime(timeToLive.toMilliseconds() + 1);
		Assert.assertNull(state.value());
	}
}
 

private void handleValue(
	Object value,
	Collector<String> out,
	TimerService timerService,
	int channel) throws IOException {
	final ValueState<String> state = getRuntimeContext().getState(this.state);
	if (state.value() == null) {
		out.collect("INPUT" + channel + ":" + value);
		state.update(String.valueOf(value));
		timerService.registerProcessingTimeTimer(timerService.currentProcessingTime() + 5);
	} else {
		state.clear();
		timerService.deleteProcessingTimeTimer(timerService.currentProcessingTime() + 4);
	}
}
 

private void handleValue(
	Object value,
	Collector<String> out,
	TimerService timerService,
	int channel) throws IOException {
	final ValueState<String> state = getRuntimeContext().getState(this.state);
	if (state.value() == null) {
		out.collect("INPUT" + channel + ":" + value);
		state.update(String.valueOf(value));
		timerService.registerProcessingTimeTimer(timerService.currentProcessingTime() + 5);
	} else {
		state.clear();
		timerService.deleteProcessingTimeTimer(timerService.currentProcessingTime() + 4);
	}
}
 

@Override
public Tuple2<Long, Long> map(Tuple2<Long, Long> value) throws Exception {
	ValueState<Long> state = getRuntimeContext().getState(
		new ValueStateDescriptor<>("testState", new TestSerializer()));

	state.update(value.f1);
	return value;
}
 

@Test
@SuppressWarnings("unchecked")
public void testMultipleValueStates() throws Exception {
	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(
			IntSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			new DummyEnvironment());

	ValueStateDescriptor<String> desc1 = new ValueStateDescriptor<>("a-string", StringSerializer.INSTANCE);
	ValueStateDescriptor<Integer> desc2 = new ValueStateDescriptor<>("an-integer", IntSerializer.INSTANCE);

	ValueState<String> state1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc1);
	ValueState<Integer> state2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc2);

	// some modifications to the state
	backend.setCurrentKey(1);
	assertNull(state1.value());
	assertNull(state2.value());
	state1.update("1");

	// state2 should still have nothing
	assertEquals("1", state1.value());
	assertNull(state2.value());
	state2.update(13);

	// both have some state now
	assertEquals("1", state1.value());
	assertEquals(13, (int) state2.value());

	// draw a snapshot
	KeyedStateHandle snapshot1 =
		runSnapshot(
			backend.snapshot(682375462378L, 2, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);

	backend.dispose();
	backend = restoreKeyedBackend(
			IntSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			Collections.singletonList(snapshot1),
			new DummyEnvironment());

	snapshot1.discardState();

	backend.setCurrentKey(1);

	state1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc1);
	state2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc2);

	// verify that they are still the same
	assertEquals("1", state1.value());
	assertEquals(13, (int) state2.value());

	backend.dispose();
}
 

private void testKeySerializerUpgrade(
		TypeSerializer<TestType> initialKeySerializer,
		TypeSerializer<TestType> newKeySerializer) throws Exception {

	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();

	AbstractKeyedStateBackend<TestType> backend = createKeyedBackend(initialKeySerializer);

	final String stateName = "test-name";
	try {
		ValueStateDescriptor<Integer> kvId = new ValueStateDescriptor<>(stateName, Integer.class);
		ValueState<Integer> valueState = backend
			.getPartitionedState(VoidNamespace.INSTANCE, CustomVoidNamespaceSerializer.INSTANCE, kvId);

		backend.setCurrentKey(new TestType("foo", 123));
		valueState.update(1);
		backend.setCurrentKey(new TestType("bar", 456));
		valueState.update(5);

		KeyedStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);
		backend.dispose();

		backend = restoreKeyedBackend(newKeySerializer, snapshot);

		valueState = backend
			.getPartitionedState(VoidNamespace.INSTANCE, CustomVoidNamespaceSerializer.INSTANCE, kvId);

		// access and check previous state
		backend.setCurrentKey(new TestType("foo", 123));
		assertEquals(1, valueState.value().intValue());
		backend.setCurrentKey(new TestType("bar", 456));
		assertEquals(5, valueState.value().intValue());

		// do another snapshot to verify the snapshot logic after migration
		snapshot = runSnapshot(
			backend.snapshot(2L, 3L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);
		snapshot.discardState();
	} finally {
		backend.dispose();
	}
}
 

private void testKeyedValueStateUpgrade(
		ValueStateDescriptor<TestType> initialAccessDescriptor,
		ValueStateDescriptor<TestType> newAccessDescriptorAfterRestore) throws Exception {

	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();

	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {
		ValueState<TestType> valueState = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			CustomVoidNamespaceSerializer.INSTANCE,
			initialAccessDescriptor);

		backend.setCurrentKey(1);
		valueState.update(new TestType("foo", 1456));
		backend.setCurrentKey(2);
		valueState.update(new TestType("bar", 478));
		backend.setCurrentKey(3);
		valueState.update(new TestType("hello", 189));

		KeyedStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);
		backend.dispose();

		backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot);

		valueState = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			CustomVoidNamespaceSerializer.INSTANCE,
			newAccessDescriptorAfterRestore);

		// make sure that reading and writing each key state works with the new serializer
		backend.setCurrentKey(1);
		assertEquals(new TestType("foo", 1456), valueState.value());
		valueState.update(new TestType("newValue1", 751));

		backend.setCurrentKey(2);
		assertEquals(new TestType("bar", 478), valueState.value());
		valueState.update(new TestType("newValue2", 167));

		backend.setCurrentKey(3);
		assertEquals(new TestType("hello", 189), valueState.value());
		valueState.update(new TestType("newValue3", 444));

		// do another snapshot to verify the snapshot logic after migration
		snapshot = runSnapshot(
			backend.snapshot(2L, 3L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);
		snapshot.discardState();
	} finally {
		backend.dispose();
	}
}
 

<K> void testIteratorHelper(
	TypeSerializer<K> keySerializer,
	TypeSerializer namespaceSerializer,
	int maxKeyGroupNumber,
	Function<Integer, K> getKeyFunc) throws Exception {

	String testStateName = "aha";
	String namespace = "ns";

	String dbPath = tmp.newFolder().getAbsolutePath();
	String checkpointPath = tmp.newFolder().toURI().toString();
	RocksDBStateBackend backend = new RocksDBStateBackend(new FsStateBackend(checkpointPath), true);
	backend.setDbStoragePath(dbPath);

	MockEnvironment env = MockEnvironment.builder().build();
	RocksDBKeyedStateBackend<K> keyedStateBackend = (RocksDBKeyedStateBackend<K>) backend.createKeyedStateBackend(
		env,
		new JobID(),
		"Test",
		keySerializer,
		maxKeyGroupNumber,
		new KeyGroupRange(0, maxKeyGroupNumber - 1),
		mock(TaskKvStateRegistry.class),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());

	try {
		ValueState<String> testState = keyedStateBackend.getPartitionedState(
			namespace,
			namespaceSerializer,
			new ValueStateDescriptor<String>(testStateName, String.class));

		// insert record
		for (int i = 0; i < 1000; ++i) {
			keyedStateBackend.setCurrentKey(getKeyFunc.apply(i));
			testState.update(String.valueOf(i));
		}

		DataOutputSerializer outputStream = new DataOutputSerializer(8);
		boolean ambiguousKeyPossible = RocksDBKeySerializationUtils.isAmbiguousKeyPossible(keySerializer, namespaceSerializer);
		RocksDBKeySerializationUtils.writeNameSpace(
			namespace,
			namespaceSerializer,
			outputStream,
			ambiguousKeyPossible);

		byte[] nameSpaceBytes = outputStream.getCopyOfBuffer();

		// already created with the state, should be closed with the backend
		ColumnFamilyHandle handle = keyedStateBackend.getColumnFamilyHandle(testStateName);

		try (
			RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(keyedStateBackend.db, handle, keyedStateBackend.getReadOptions());
			RocksStateKeysIterator<K> iteratorWrapper =
				new RocksStateKeysIterator<>(
					iterator,
					testStateName,
					keySerializer,
					keyedStateBackend.getKeyGroupPrefixBytes(),
					ambiguousKeyPossible,
					nameSpaceBytes)) {

			iterator.seekToFirst();

			// valid record
			List<Integer> fetchedKeys = new ArrayList<>(1000);
			while (iteratorWrapper.hasNext()) {
				fetchedKeys.add(Integer.parseInt(iteratorWrapper.next().toString()));
			}

			fetchedKeys.sort(Comparator.comparingInt(a -> a));
			Assert.assertEquals(1000, fetchedKeys.size());

			for (int i = 0; i < 1000; ++i) {
				Assert.assertEquals(i, fetchedKeys.get(i).intValue());
			}
		}
	} finally {
		if (keyedStateBackend != null) {
			keyedStateBackend.dispose();
		}
		IOUtils.closeQuietly(env);
	}
}
 

/**
 * Verify state restore resilience when:
 *  - snapshot was taken without any Kryo registrations, specific serializers or default serializers for the state type
 *  - restored with the state type registered (no specific serializer)
 *
 * This test should not fail, because de- / serialization of the state should not be performed with Kryo's default
 * {@link com.esotericsoftware.kryo.serializers.FieldSerializer}.
 */
@Test
public void testKryoRegisteringRestoreResilienceWithRegisteredType() throws Exception {
	CheckpointStreamFactory streamFactory = createStreamFactory();
	Environment env = new DummyEnvironment();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE, env);

	TypeInformation<TestPojo> pojoType = new GenericTypeInfo<>(TestPojo.class);

	// make sure that we are in fact using the KryoSerializer
	assertTrue(pojoType.createSerializer(env.getExecutionConfig()) instanceof KryoSerializer);

	ValueStateDescriptor<TestPojo> kvId = new ValueStateDescriptor<>("id", pojoType);

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

	// ============== create snapshot - no Kryo registration or specific / default serializers ==============

	// make some more modifications
	backend.setCurrentKey(1);
	state.update(new TestPojo("u1", 1));

	backend.setCurrentKey(2);
	state.update(new TestPojo("u2", 2));

	KeyedStateHandle snapshot = runSnapshot(
		backend.snapshot(
			682375462378L,
			2,
			streamFactory,
			CheckpointOptions.forCheckpointWithDefaultLocation()),
		sharedStateRegistry);

	backend.dispose();

	// ====================================== restore snapshot  ======================================

	env.getExecutionConfig().registerKryoType(TestPojo.class);

	backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot, env);

	snapshot.discardState();

	state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
	backend.setCurrentKey(1);
	assertEquals(state.value(), new TestPojo("u1", 1));

	backend.setCurrentKey(2);
	assertEquals(state.value(), new TestPojo("u2", 2));

	backend.dispose();
}
 

/**
 * Tests a simple successful query via an EmbeddedChannel.
 */
@Test
public void testSimpleQuery() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	// Register state
	ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
	desc.setQueryable("vanilla");

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Update the KvState and request it
	int expectedValue = 712828289;

	int key = 99812822;
	backend.setCurrentKey(key);
	ValueState<Integer> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	state.update(expectedValue);

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(
			registryListener.kvStateId, serializedKeyAndNamespace);

	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	long deserRequestId = MessageSerializer.getRequestId(buf);
	KvStateResponse response = serializer.deserializeResponse(buf);
	buf.release();

	assertEquals(requestId, deserRequestId);

	int actualValue = KvStateSerializer.deserializeValue(response.getContent(), IntSerializer.INSTANCE);
	assertEquals(expectedValue, actualValue);

	assertEquals(stats.toString(), 1, stats.getNumRequests());

	// Wait for async successful request report
	long deadline = System.nanoTime() + TimeUnit.NANOSECONDS.convert(30, TimeUnit.SECONDS);
	while (stats.getNumSuccessful() != 1L && System.nanoTime() <= deadline) {
		Thread.sleep(10L);
	}

	assertEquals(stats.toString(), 1L, stats.getNumSuccessful());
}
 

/**
 * Verify state restore resilience when:
 *  - snapshot was taken without any Kryo registrations, specific serializers or default serializers for the state type
 *  - restored with the state type registered (no specific serializer)
 *
 * This test should not fail, because de- / serialization of the state should not be performed with Kryo's default
 * {@link com.esotericsoftware.kryo.serializers.FieldSerializer}.
 */
@Test
public void testKryoRegisteringRestoreResilienceWithRegisteredType() throws Exception {
	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE, env);

	TypeInformation<TestPojo> pojoType = new GenericTypeInfo<>(TestPojo.class);

	// make sure that we are in fact using the KryoSerializer
	assertTrue(pojoType.createSerializer(env.getExecutionConfig()) instanceof KryoSerializer);

	ValueStateDescriptor<TestPojo> kvId = new ValueStateDescriptor<>("id", pojoType);

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

	// ============== create snapshot - no Kryo registration or specific / default serializers ==============

	// make some more modifications
	backend.setCurrentKey(1);
	state.update(new TestPojo("u1", 1));

	backend.setCurrentKey(2);
	state.update(new TestPojo("u2", 2));

	KeyedStateHandle snapshot = runSnapshot(
		backend.snapshot(
			682375462378L,
			2,
			streamFactory,
			CheckpointOptions.forCheckpointWithDefaultLocation()),
		sharedStateRegistry);

	backend.dispose();

	// ====================================== restore snapshot  ======================================

	env.getExecutionConfig().registerKryoType(TestPojo.class);

	backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot, env);

	snapshot.discardState();

	state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
	backend.setCurrentKey(1);
	assertEquals(state.value(), new TestPojo("u1", 1));

	backend.setCurrentKey(2);
	assertEquals(state.value(), new TestPojo("u2", 2));

	backend.dispose();
}
 

private void testNamespaceSerializerUpgrade(
		TypeSerializer<TestType> initialNamespaceSerializer,
		TypeSerializer<TestType> newNamespaceSerializerAfterRestore) throws Exception {

	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();

	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	final String stateName = "test-name";
	try {
		ValueStateDescriptor<Integer> kvId = new ValueStateDescriptor<>(stateName, Integer.class);
		ValueState<Integer> valueState = backend
			.getPartitionedState(
				new TestType("namespace", 123),
				initialNamespaceSerializer,
				kvId);

		backend.setCurrentKey(1);
		valueState.update(10);
		backend.setCurrentKey(5);
		valueState.update(50);

		KeyedStateHandle snapshot = runSnapshot(
			backend.snapshot(1L, 2L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);

		// test incompatible namespace serializer; start with a freshly restored backend
		backend.dispose();
		backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot);

		valueState = backend.getPartitionedState(
			new TestType("namespace", 123),
			newNamespaceSerializerAfterRestore,
			kvId);

		// access and check previous state
		backend.setCurrentKey(1);
		assertEquals(10, valueState.value().intValue());
		valueState.update(10);
		backend.setCurrentKey(5);
		assertEquals(50, valueState.value().intValue());

		// do another snapshot to verify the snapshot logic after migration
		snapshot = runSnapshot(
			backend.snapshot(2L, 3L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
			sharedStateRegistry);
		snapshot.discardState();
	} finally {
		backend.dispose();
	}
}
 

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
	if (enableIncrementalCheckpointing) {
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
		try {
			ValueStateDescriptor<String> kvId =
				new ValueStateDescriptor<>("id", String.class, null);

			kvId.initializeSerializerUnlessSet(new ExecutionConfig());

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

			Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
			SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistry());
			for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {

				reset(sharedStateRegistry);

				backend.setCurrentKey(checkpointId);
				state.update("Hello-" + checkpointId);

				RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(
					checkpointId,
					checkpointId,
					createStreamFactory(),
					CheckpointOptions.forCheckpointWithDefaultLocation());

				snapshot.run();

				SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();

				IncrementalRemoteKeyedStateHandle stateHandle =
					(IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();

				Map<StateHandleID, StreamStateHandle> sharedState =
					new HashMap<>(stateHandle.getSharedState());

				stateHandle.registerSharedStates(sharedStateRegistry);

				for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
					verify(sharedStateRegistry).registerReference(
						stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()),
						e.getValue());
				}

				previousStateHandles.add(stateHandle);
				backend.notifyCheckpointComplete(checkpointId);

				//-----------------------------------------------------------------

				if (previousStateHandles.size() > 1) {
					checkRemove(previousStateHandles.remove(), sharedStateRegistry);
				}
			}

			while (!previousStateHandles.isEmpty()) {

				reset(sharedStateRegistry);

				checkRemove(previousStateHandles.remove(), sharedStateRegistry);
			}
		} finally {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

@Override
void updateInternal(@Nonnull ValueState<String> state, String update) throws Exception {
	state.update(update);
}
 

@Test
@SuppressWarnings("unchecked")
public void testNumStateEntries() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class);

	assertEquals(0, backend.numKeyValueStateEntries());

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

	backend.setCurrentKey(0);
	state.update("hello");
	state.update("ciao");

	assertEquals(1, backend.numKeyValueStateEntries());

	backend.setCurrentKey(42);
	state.update("foo");

	assertEquals(2, backend.numKeyValueStateEntries());

	backend.setCurrentKey(0);
	state.clear();

	assertEquals(1, backend.numKeyValueStateEntries());

	backend.setCurrentKey(42);
	state.clear();

	assertEquals(0, backend.numKeyValueStateEntries());

	backend.dispose();
}