Java Code Examples for org.apache.flink.queryablestate.client.state.serialization.KvStateSerializer#serializeKeyAndNamespace()

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <V, K, N> List<V> getSerializedList(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
protected static <V, K, N> V getSerializedValue(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeValue(serializedValue, valueSerializer);
	}
}
 

/**
 * Tests key and namespace serialization utils.
 */
@Test
public void testKeyAndNamespaceSerialization() throws Exception {
	TypeSerializer<Long> keySerializer = LongSerializer.INSTANCE;
	TypeSerializer<String> namespaceSerializer = StringSerializer.INSTANCE;

	long expectedKey = Integer.MAX_VALUE + 12323L;
	String expectedNamespace = "knilf";

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			expectedKey, keySerializer, expectedNamespace, namespaceSerializer);

	Tuple2<Long, String> actual = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, keySerializer, namespaceSerializer);

	assertEquals(expectedKey, actual.f0.longValue());
	assertEquals(expectedNamespace, actual.f1);
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <V, K, N> List<V> getSerializedList(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
protected static <V, K, N> V getSerializedValue(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeValue(serializedValue, valueSerializer);
	}
}
 

/**
 * Tests key and namespace serialization utils.
 */
@Test
public void testKeyAndNamespaceSerialization() throws Exception {
	TypeSerializer<Long> keySerializer = LongSerializer.INSTANCE;
	TypeSerializer<String> namespaceSerializer = StringSerializer.INSTANCE;

	long expectedKey = Integer.MAX_VALUE + 12323L;
	String expectedNamespace = "knilf";

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			expectedKey, keySerializer, expectedNamespace, namespaceSerializer);

	Tuple2<Long, String> actual = KvStateSerializer.deserializeKeyAndNamespace(
			serializedKeyAndNamespace, keySerializer, namespaceSerializer);

	assertEquals(expectedKey, actual.f0.longValue());
	assertEquals(expectedNamespace, actual.f1);
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
protected static <V, K, N> V getSerializedValue(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeValue(serializedValue, valueSerializer);
	}
}
 

/**
 * Verifies that the serialization of a list using the given list state
 * matches the deserialization with {@link KvStateSerializer#deserializeList}.
 *
 * @param key
 * 		key of the list state
 * @param listState
 * 		list state using the {@link VoidNamespace}, must also be a {@link InternalKvState} instance
 *
 * @throws Exception
 */
public static void testListSerialization(
		final long key,
		final InternalListState<Long, VoidNamespace, Long> listState) throws Exception {

	TypeSerializer<Long> valueSerializer = LongSerializer.INSTANCE;
	listState.setCurrentNamespace(VoidNamespace.INSTANCE);

	// List
	final int numElements = 10;

	final List<Long> expectedValues = new ArrayList<>();
	for (int i = 0; i < numElements; i++) {
		final long value = ThreadLocalRandom.current().nextLong();
		expectedValues.add(value);
		listState.add(value);
	}

	final byte[] serializedKey =
		KvStateSerializer.serializeKeyAndNamespace(
			key, LongSerializer.INSTANCE,
			VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);

	final byte[] serializedValues = listState.getSerializedValue(
			serializedKey,
			listState.getKeySerializer(),
			listState.getNamespaceSerializer(),
			listState.getValueSerializer());

	List<Long> actualValues = KvStateSerializer.deserializeList(serializedValues, valueSerializer);
	assertEquals(expectedValues, actualValues);

	// Single value
	long expectedValue = ThreadLocalRandom.current().nextLong();
	byte[] serializedValue = KvStateSerializer.serializeValue(expectedValue, valueSerializer);
	List<Long> actualValue = KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	assertEquals(1, actualValue.size());
	assertEquals(expectedValue, actualValue.get(0).longValue());
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <UK, UV, K, N> Map<UK, UV> getSerializedMap(
		InternalKvState<K, N, Map<UK, UV>> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<UK> userKeySerializer,
		TypeSerializer<UV> userValueSerializer
) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeMap(serializedValue, userKeySerializer, userValueSerializer);
	}
}
 

/**
 * Returns a future holding the request result.
 * @param jobId                     JobID of the job the queryable state belongs to.
 * @param queryableStateName        Name under which the state is queryable.
 * @param key			            The key that the state we request is associated with.
 * @param namespace					The namespace of the state.
 * @param keyTypeInfo				The {@link TypeInformation} of the keys.
 * @param namespaceTypeInfo			The {@link TypeInformation} of the namespace.
 * @param stateDescriptor			The {@link StateDescriptor} of the state we want to query.
 * @return Future holding the immutable {@link State} object containing the result.
 */
private <K, N, S extends State, V> CompletableFuture<S> getKvState(
		final JobID jobId,
		final String queryableStateName,
		final K key,
		final N namespace,
		final TypeInformation<K> keyTypeInfo,
		final TypeInformation<N> namespaceTypeInfo,
		final StateDescriptor<S, V> stateDescriptor) {

	Preconditions.checkNotNull(jobId);
	Preconditions.checkNotNull(queryableStateName);
	Preconditions.checkNotNull(key);
	Preconditions.checkNotNull(namespace);

	Preconditions.checkNotNull(keyTypeInfo);
	Preconditions.checkNotNull(namespaceTypeInfo);
	Preconditions.checkNotNull(stateDescriptor);

	TypeSerializer<K> keySerializer = keyTypeInfo.createSerializer(executionConfig);
	TypeSerializer<N> namespaceSerializer = namespaceTypeInfo.createSerializer(executionConfig);

	stateDescriptor.initializeSerializerUnlessSet(executionConfig);

	final byte[] serializedKeyAndNamespace;
	try {
		serializedKeyAndNamespace = KvStateSerializer
				.serializeKeyAndNamespace(key, keySerializer, namespace, namespaceSerializer);
	} catch (IOException e) {
		return FutureUtils.getFailedFuture(e);
	}

	return getKvState(jobId, queryableStateName, key.hashCode(), serializedKeyAndNamespace)
		.thenApply(stateResponse -> createState(stateResponse, stateDescriptor));
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <UK, UV, K, N> Map<UK, UV> getSerializedMap(
		InternalKvState<K, N, Map<UK, UV>> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<UK> userKeySerializer,
		TypeSerializer<UV> userValueSerializer
) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeMap(serializedValue, userKeySerializer, userValueSerializer);
	}
}
 

/**
 * Verifies that the serialization of a list using the given list state
 * matches the deserialization with {@link KvStateSerializer#deserializeList}.
 *
 * @param key
 * 		key of the list state
 * @param listState
 * 		list state using the {@link VoidNamespace}, must also be a {@link InternalKvState} instance
 *
 * @throws Exception
 */
public static void testListSerialization(
		final long key,
		final InternalListState<Long, VoidNamespace, Long> listState) throws Exception {

	TypeSerializer<Long> valueSerializer = LongSerializer.INSTANCE;
	listState.setCurrentNamespace(VoidNamespace.INSTANCE);

	// List
	final int numElements = 10;

	final List<Long> expectedValues = new ArrayList<>();
	for (int i = 0; i < numElements; i++) {
		final long value = ThreadLocalRandom.current().nextLong();
		expectedValues.add(value);
		listState.add(value);
	}

	final byte[] serializedKey =
		KvStateSerializer.serializeKeyAndNamespace(
			key, LongSerializer.INSTANCE,
			VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);

	final byte[] serializedValues = listState.getSerializedValue(
			serializedKey,
			listState.getKeySerializer(),
			listState.getNamespaceSerializer(),
			listState.getValueSerializer());

	List<Long> actualValues = KvStateSerializer.deserializeList(serializedValues, valueSerializer);
	assertEquals(expectedValues, actualValues);

	// Single value
	long expectedValue = ThreadLocalRandom.current().nextLong();
	byte[] serializedValue = KvStateSerializer.serializeValue(expectedValue, valueSerializer);
	List<Long> actualValue = KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	assertEquals(1, actualValue.size());
	assertEquals(expectedValue, actualValue.get(0).longValue());
}
 

/**
 * Returns a future holding the request result.
 * @param jobId                     JobID of the job the queryable state belongs to.
 * @param queryableStateName        Name under which the state is queryable.
 * @param key			            The key that the state we request is associated with.
 * @param namespace					The namespace of the state.
 * @param keyTypeInfo				The {@link TypeInformation} of the keys.
 * @param namespaceTypeInfo			The {@link TypeInformation} of the namespace.
 * @param stateDescriptor			The {@link StateDescriptor} of the state we want to query.
 * @return Future holding the immutable {@link State} object containing the result.
 */
private <K, N, S extends State, V> CompletableFuture<S> getKvState(
		final JobID jobId,
		final String queryableStateName,
		final K key,
		final N namespace,
		final TypeInformation<K> keyTypeInfo,
		final TypeInformation<N> namespaceTypeInfo,
		final StateDescriptor<S, V> stateDescriptor) {

	Preconditions.checkNotNull(jobId);
	Preconditions.checkNotNull(queryableStateName);
	Preconditions.checkNotNull(key);
	Preconditions.checkNotNull(namespace);

	Preconditions.checkNotNull(keyTypeInfo);
	Preconditions.checkNotNull(namespaceTypeInfo);
	Preconditions.checkNotNull(stateDescriptor);

	TypeSerializer<K> keySerializer = keyTypeInfo.createSerializer(executionConfig);
	TypeSerializer<N> namespaceSerializer = namespaceTypeInfo.createSerializer(executionConfig);

	stateDescriptor.initializeSerializerUnlessSet(executionConfig);

	final byte[] serializedKeyAndNamespace;
	try {
		serializedKeyAndNamespace = KvStateSerializer
				.serializeKeyAndNamespace(key, keySerializer, namespace, namespaceSerializer);
	} catch (IOException e) {
		return FutureUtils.getFailedFuture(e);
	}

	return getKvState(jobId, queryableStateName, key.hashCode(), serializedKeyAndNamespace)
		.thenApply(stateResponse -> createState(stateResponse, stateDescriptor));
}
 

/**
 * Verifies that the serialization of a map using the given map state
 * matches the deserialization with {@link KvStateSerializer#deserializeList}.
 *
 * @param key
 * 		key of the map state
 * @param mapState
 * 		map state using the {@link VoidNamespace}, must also be a {@link InternalKvState} instance
 *
 * @throws Exception
 */
public static void testMapSerialization(
		final long key,
		final InternalMapState<Long, VoidNamespace, Long, String> mapState) throws Exception {

	TypeSerializer<Long> userKeySerializer = LongSerializer.INSTANCE;
	TypeSerializer<String> userValueSerializer = StringSerializer.INSTANCE;
	mapState.setCurrentNamespace(VoidNamespace.INSTANCE);

	// Map
	final int numElements = 10;

	final Map<Long, String> expectedValues = new HashMap<>();
	for (int i = 1; i <= numElements; i++) {
		final long value = ThreadLocalRandom.current().nextLong();
		expectedValues.put(value, Long.toString(value));
		mapState.put(value, Long.toString(value));
	}

	expectedValues.put(0L, null);
	mapState.put(0L, null);

	final byte[] serializedKey =
		KvStateSerializer.serializeKeyAndNamespace(
			key, LongSerializer.INSTANCE,
			VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);

	final byte[] serializedValues = mapState.getSerializedValue(
			serializedKey,
			mapState.getKeySerializer(),
			mapState.getNamespaceSerializer(),
			mapState.getValueSerializer());

	Map<Long, String> actualValues = KvStateSerializer.deserializeMap(serializedValues, userKeySerializer, userValueSerializer);
	assertEquals(expectedValues.size(), actualValues.size());
	for (Map.Entry<Long, String> actualEntry : actualValues.entrySet()) {
		assertEquals(expectedValues.get(actualEntry.getKey()), actualEntry.getValue());
	}

	// Single value
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	long expectedKey = ThreadLocalRandom.current().nextLong();
	String expectedValue = Long.toString(expectedKey);
	byte[] isNull = {0};

	baos.write(KvStateSerializer.serializeValue(expectedKey, userKeySerializer));
	baos.write(isNull);
	baos.write(KvStateSerializer.serializeValue(expectedValue, userValueSerializer));
	byte[] serializedValue = baos.toByteArray();

	Map<Long, String> actualValue = KvStateSerializer.deserializeMap(serializedValue, userKeySerializer, userValueSerializer);
	assertEquals(1, actualValue.size());
	assertEquals(expectedValue, actualValue.get(expectedKey));
}
 

/**
 * Verifies that the serialization of a map using the given map state
 * matches the deserialization with {@link KvStateSerializer#deserializeList}.
 *
 * @param key
 * 		key of the map state
 * @param mapState
 * 		map state using the {@link VoidNamespace}, must also be a {@link InternalKvState} instance
 *
 * @throws Exception
 */
public static void testMapSerialization(
		final long key,
		final InternalMapState<Long, VoidNamespace, Long, String> mapState) throws Exception {

	TypeSerializer<Long> userKeySerializer = LongSerializer.INSTANCE;
	TypeSerializer<String> userValueSerializer = StringSerializer.INSTANCE;
	mapState.setCurrentNamespace(VoidNamespace.INSTANCE);

	// Map
	final int numElements = 10;

	final Map<Long, String> expectedValues = new HashMap<>();
	for (int i = 1; i <= numElements; i++) {
		final long value = ThreadLocalRandom.current().nextLong();
		expectedValues.put(value, Long.toString(value));
		mapState.put(value, Long.toString(value));
	}

	expectedValues.put(0L, null);
	mapState.put(0L, null);

	final byte[] serializedKey =
		KvStateSerializer.serializeKeyAndNamespace(
			key, LongSerializer.INSTANCE,
			VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);

	final byte[] serializedValues = mapState.getSerializedValue(
			serializedKey,
			mapState.getKeySerializer(),
			mapState.getNamespaceSerializer(),
			mapState.getValueSerializer());

	Map<Long, String> actualValues = KvStateSerializer.deserializeMap(serializedValues, userKeySerializer, userValueSerializer);
	assertEquals(expectedValues.size(), actualValues.size());
	for (Map.Entry<Long, String> actualEntry : actualValues.entrySet()) {
		assertEquals(expectedValues.get(actualEntry.getKey()), actualEntry.getValue());
	}

	// Single value
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	long expectedKey = ThreadLocalRandom.current().nextLong();
	String expectedValue = Long.toString(expectedKey);
	byte[] isNull = {0};

	baos.write(KvStateSerializer.serializeValue(expectedKey, userKeySerializer));
	baos.write(isNull);
	baos.write(KvStateSerializer.serializeValue(expectedValue, userValueSerializer));
	byte[] serializedValue = baos.toByteArray();

	Map<Long, String> actualValue = KvStateSerializer.deserializeMap(serializedValue, userKeySerializer, userValueSerializer);
	assertEquals(1, actualValue.size());
	assertEquals(expectedValue, actualValue.get(expectedKey));
}
 

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

/**
 * Tests that large responses are chunked.
 */
@Test
public void testChunkedResponse() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	KvStateRequestStats 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);

	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);

	// Register state
	ValueStateDescriptor<byte[]> desc = new ValueStateDescriptor<>("any", BytePrimitiveArraySerializer.INSTANCE);
	desc.setQueryable("vanilla");

	ValueState<byte[]> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	// Update KvState
	byte[] bytes = new byte[2 * channel.config().getWriteBufferHighWaterMark()];

	byte current = 0;
	for (int i = 0; i < bytes.length; i++) {
		bytes[i] = current++;
	}

	int key = 99812822;
	backend.setCurrentKey(key);
	state.update(bytes);

	// Request
	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);

	Object msg = readInboundBlocking(channel);
	assertTrue("Not ChunkedByteBuf", msg instanceof ChunkedByteBuf);
	((ChunkedByteBuf) msg).close();
}
 

/**
 * Verifies that the serialization of a map using the given map state
 * matches the deserialization with {@link KvStateSerializer#deserializeList}.
 *
 * @param key
 * 		key of the map state
 * @param mapState
 * 		map state using the {@link VoidNamespace}, must also be a {@link InternalKvState} instance
 *
 * @throws Exception
 */
public static void testMapSerialization(
		final long key,
		final InternalMapState<Long, VoidNamespace, Long, String> mapState) throws Exception {

	TypeSerializer<Long> userKeySerializer = LongSerializer.INSTANCE;
	TypeSerializer<String> userValueSerializer = StringSerializer.INSTANCE;
	mapState.setCurrentNamespace(VoidNamespace.INSTANCE);

	// Map
	final int numElements = 10;

	final Map<Long, String> expectedValues = new HashMap<>();
	for (int i = 1; i <= numElements; i++) {
		final long value = ThreadLocalRandom.current().nextLong();
		expectedValues.put(value, Long.toString(value));
		mapState.put(value, Long.toString(value));
	}

	expectedValues.put(0L, null);
	mapState.put(0L, null);

	final byte[] serializedKey =
		KvStateSerializer.serializeKeyAndNamespace(
			key, LongSerializer.INSTANCE,
			VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);

	final byte[] serializedValues = mapState.getSerializedValue(
			serializedKey,
			mapState.getKeySerializer(),
			mapState.getNamespaceSerializer(),
			mapState.getValueSerializer());

	Map<Long, String> actualValues = KvStateSerializer.deserializeMap(serializedValues, userKeySerializer, userValueSerializer);
	assertEquals(expectedValues.size(), actualValues.size());
	for (Map.Entry<Long, String> actualEntry : actualValues.entrySet()) {
		assertEquals(expectedValues.get(actualEntry.getKey()), actualEntry.getValue());
	}

	// Single value
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	long expectedKey = ThreadLocalRandom.current().nextLong();
	String expectedValue = Long.toString(expectedKey);
	byte[] isNull = {0};

	baos.write(KvStateSerializer.serializeValue(expectedKey, userKeySerializer));
	baos.write(isNull);
	baos.write(KvStateSerializer.serializeValue(expectedValue, userValueSerializer));
	byte[] serializedValue = baos.toByteArray();

	Map<Long, String> actualValue = KvStateSerializer.deserializeMap(serializedValue, userKeySerializer, userValueSerializer);
	assertEquals(1, actualValue.size());
	assertEquals(expectedValue, actualValue.get(expectedKey));
}
 

/**
 * Tests the failure response with {@link UnknownKeyOrNamespaceException} as cause
 * on queries for non-existing keys.
 */
@Test
public void testQueryUnknownKey() 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);

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

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

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

	backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);

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

	long requestId = Integer.MAX_VALUE + 22982L;

	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_FAILURE, MessageSerializer.deserializeHeader(buf));
	RequestFailure response = MessageSerializer.deserializeRequestFailure(buf);

	assertEquals(requestId, response.getRequestId());

	assertTrue("Did not respond with expected failure cause", response.getCause() instanceof UnknownKeyOrNamespaceException);

	assertEquals(1L, stats.getNumRequests());
	assertEquals(1L, stats.getNumFailed());
}
 

/**
 * Tests that large responses are chunked.
 */
@Test
public void testChunkedResponse() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	KvStateRequestStats 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);

	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);

	// Register state
	ValueStateDescriptor<byte[]> desc = new ValueStateDescriptor<>("any", BytePrimitiveArraySerializer.INSTANCE);
	desc.setQueryable("vanilla");

	ValueState<byte[]> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	// Update KvState
	byte[] bytes = new byte[2 * channel.config().getWriteBufferHighWaterMark()];

	byte current = 0;
	for (int i = 0; i < bytes.length; i++) {
		bytes[i] = current++;
	}

	int key = 99812822;
	backend.setCurrentKey(key);
	state.update(bytes);

	// Request
	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);

	Object msg = readInboundBlocking(channel);
	assertTrue("Not ChunkedByteBuf", msg instanceof ChunkedByteBuf);
}