Java Code Examples for org.apache.flink.queryablestate.messages.KvStateInternalRequest#KvStateInternalRequestDeserializer

/**
 * Tests that the channel is closed if an Exception reaches the channel handler.
 */
@Test
public void testCloseChannelOnExceptionCaught() 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(handler);

	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));

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

	// Verify the response
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable response = MessageSerializer.deserializeServerFailure(buf);

	assertTrue(response.getMessage().contains("Expected test Exception"));

	channel.closeFuture().await(READ_TIMEOUT_MILLIS);
	assertFalse(channel.isActive());
}
 

/**
 * Tests response serialization.
 */
@Test
public void testResponseSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 72727278L;
	byte[] serializedResult = randomByteArray(1024);

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

	ByteBuf buf = MessageSerializer.serializeResponse(alloc, requestId, response);

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateResponse responseDeser = serializer.deserializeResponse(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertArrayEquals(serializedResult, responseDeser.getContent());
}
 

/**
 * Tests response serialization with zero-length serialized result.
 */
@Test
public void testResponseSerializationWithZeroLengthSerializedResult() throws Exception {
	byte[] serializedResult = new byte[0];

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

	ByteBuf buf = MessageSerializer.serializeResponse(alloc, 72727278L, response);

	int frameLength = buf.readInt();

	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	assertEquals(72727278L, MessageSerializer.getRequestId(buf));
	KvStateResponse responseDeser = serializer.deserializeResponse(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertArrayEquals(serializedResult, responseDeser.getContent());
}
 

/**
 * Tests request serialization with zero-length serialized key and namespace.
 */
@Test
public void testRequestSerializationWithZeroLengthKeyAndNamespace() throws Exception {

	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = new byte[0];

	final KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, serializedKeyAndNamespace);
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

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

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateInternalRequest requestDeser = serializer.deserializeRequest(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(kvStateId, requestDeser.getKvStateId());
	assertArrayEquals(serializedKeyAndNamespace, requestDeser.getSerializedKeyAndNamespace());
}
 

/**
 * Tests that the channel is closed if an Exception reaches the channel handler.
 */
@Test
public void testCloseChannelOnExceptionCaught() 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(handler);

	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));

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

	// Verify the response
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable response = MessageSerializer.deserializeServerFailure(buf);

	assertTrue(response.getMessage().contains("Expected test Exception"));

	channel.closeFuture().await(READ_TIMEOUT_MILLIS);
	assertFalse(channel.isActive());
}
 

/**
 * Tests response serialization.
 */
@Test
public void testResponseSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 72727278L;
	byte[] serializedResult = randomByteArray(1024);

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

	ByteBuf buf = MessageSerializer.serializeResponse(alloc, requestId, response);

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	assertEquals(requestId, MessageSerializer.getRequestId(buf));
	KvStateResponse responseDeser = serializer.deserializeResponse(buf);

	assertEquals(buf.readerIndex(), frameLength + 4);

	assertArrayEquals(serializedResult, responseDeser.getContent());
}
 

/**
 * Tests response serialization with zero-length serialized result.
 */
@Test
public void testResponseSerializationWithZeroLengthSerializedResult() throws Exception {
	byte[] serializedResult = new byte[0];

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

	ByteBuf buf = MessageSerializer.serializeResponse(alloc, 72727278L, response);

	int frameLength = buf.readInt();

	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	assertEquals(72727278L, MessageSerializer.getRequestId(buf));
	KvStateResponse responseDeser = serializer.deserializeResponse(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertArrayEquals(serializedResult, responseDeser.getContent());
}
 

/**
 * Tests that the channel is closed if an Exception reaches the channel handler.
 */
@Test
public void testCloseChannelOnExceptionCaught() 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(handler);

	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));

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

	// Verify the response
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable response = MessageSerializer.deserializeServerFailure(buf);
	buf.release();

	assertTrue(response.getMessage().contains("Expected test Exception"));

	channel.closeFuture().await(READ_TIMEOUT_MILLIS);
	assertFalse(channel.isActive());
}
 

/**
 * Tests that incoming buffer instances are recycled.
 */
@Test
public void testIncomingBufferIsRecycled() 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);

	KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), 282872L, request);

	assertEquals(1L, serRequest.refCnt());

	// Write regular request
	channel.writeInbound(serRequest);
	assertEquals("Buffer not recycled", 0L, serRequest.refCnt());

	// Write unexpected msg
	ByteBuf unexpected = channel.alloc().buffer(8);
	unexpected.writeInt(4);
	unexpected.writeInt(4);

	assertEquals(1L, unexpected.refCnt());

	channel.writeInbound(unexpected);
	assertEquals("Buffer not recycled", 0L, unexpected.refCnt());
}
 

private static Client<KvStateInternalRequest, KvStateResponse> createInternalClient(int threads) {
	final MessageSerializer<KvStateInternalRequest, KvStateResponse> messageSerializer =
			new MessageSerializer<>(
					new KvStateInternalRequest.KvStateInternalRequestDeserializer(),
					new KvStateResponse.KvStateResponseDeserializer());

	return new Client<>(
			"Queryable State Proxy Client",
			threads,
			messageSerializer,
			new DisabledKvStateRequestStats());
}
 

/**
 * Tests the failure response with {@link UnknownKvStateIdException} as cause on
 * queries for unregistered KvStateIDs.
 */
@Test
public void testQueryUnknownKvStateID() 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);

	long requestId = Integer.MAX_VALUE + 182828L;

	KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);

	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);
	buf.release();

	assertEquals(requestId, response.getRequestId());

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

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

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

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

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

/**
 * Tests the failure response on a failure on the {@link InternalKvState#getSerializedValue(byte[], TypeSerializer, TypeSerializer, TypeSerializer)} call.
 */
@Test
public void testFailureOnGetSerializedValue() 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);

	// Failing KvState
	InternalKvState<Integer, VoidNamespace, Long> kvState =
			new InternalKvState<Integer, VoidNamespace, Long>() {
				@Override
				public TypeSerializer<Integer> getKeySerializer() {
					return IntSerializer.INSTANCE;
				}

				@Override
				public TypeSerializer<VoidNamespace> getNamespaceSerializer() {
					return VoidNamespaceSerializer.INSTANCE;
				}

				@Override
				public TypeSerializer<Long> getValueSerializer() {
					return LongSerializer.INSTANCE;
				}

				@Override
				public void setCurrentNamespace(VoidNamespace namespace) {
					// do nothing
				}

				@Override
				public byte[] getSerializedValue(
						final byte[] serializedKeyAndNamespace,
						final TypeSerializer<Integer> safeKeySerializer,
						final TypeSerializer<VoidNamespace> safeNamespaceSerializer,
						final TypeSerializer<Long> safeValueSerializer) throws Exception {
					throw new RuntimeException("Expected test Exception");
				}

				@Override
				public StateIncrementalVisitor<Integer, VoidNamespace, Long> getStateIncrementalVisitor(int recommendedMaxNumberOfReturnedRecords) {
					throw new UnsupportedOperationException();
				}

				@Override
				public void clear() {

				}
			};

	KvStateID kvStateId = registry.registerKvState(
			new JobID(),
			new JobVertexID(),
			new KeyGroupRange(0, 0),
			"vanilla",
			kvState);

	KvStateInternalRequest request = new KvStateInternalRequest(kvStateId, new byte[0]);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), 282872L, 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);

	assertTrue(response.getCause().getMessage().contains("Expected test Exception"));

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

/**
 * Tests the failure response on a rejected execution, because the query executor has been closed.
 */
@Test
public void testQueryExecutorShutDown() throws Throwable {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	KvStateServerImpl localTestServer = new KvStateServerImpl(
			InetAddress.getLocalHost(),
			Collections.singletonList(0).iterator(),
			1,
			1,
			new KvStateRegistry(),
			new DisabledKvStateRequestStats());

	localTestServer.start();
	localTestServer.shutdown();
	assertTrue(localTestServer.getQueryExecutor().isTerminated());

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

	KvStateServerHandler handler = new KvStateServerHandler(localTestServer, 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);

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

	KvStateInternalRequest request = new KvStateInternalRequest(registryListener.kvStateId, new byte[0]);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), 282872L, 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);

	assertTrue(response.getCause().getMessage().contains("RejectedExecutionException"));

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

	localTestServer.shutdown();
}
 

/**
 * Tests that on reads the expected callback methods are called and read
 * buffers are recycled.
 */
@Test
public void testReadCallbacksAndBufferRecycling() throws Exception {
	final ClientHandlerCallback<KvStateResponse> callback = mock(ClientHandlerCallback.class);

	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());
	final EmbeddedChannel channel = new EmbeddedChannel(new ClientHandler<>("Test Client", serializer, callback));

	final byte[] content = new byte[0];
	final KvStateResponse response = new KvStateResponse(content);

	//
	// Request success
	//
	ByteBuf buf = MessageSerializer.serializeResponse(channel.alloc(), 1222112277L, response);
	buf.skipBytes(4); // skip frame length

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onRequestResult(eq(1222112277L), any(KvStateResponse.class));
	assertEquals("Buffer not recycled", 0, buf.refCnt());

	//
	// Request failure
	//
	buf = MessageSerializer.serializeRequestFailure(
			channel.alloc(),
			1222112278,
			new RuntimeException("Expected test Exception"));
	buf.skipBytes(4); // skip frame length

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onRequestFailure(eq(1222112278L), isA(RuntimeException.class));
	assertEquals("Buffer not recycled", 0, buf.refCnt());

	//
	// Server failure
	//
	buf = MessageSerializer.serializeServerFailure(
			channel.alloc(),
			new RuntimeException("Expected test Exception"));
	buf.skipBytes(4); // skip frame length

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onFailure(isA(RuntimeException.class));

	//
	// Unexpected messages
	//
	buf = channel.alloc().buffer(4).writeInt(1223823);

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onFailure(isA(IllegalStateException.class));
	assertEquals("Buffer not recycled", 0, buf.refCnt());

	//
	// Exception caught
	//
	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));
	verify(callback, times(3)).onFailure(isA(RuntimeException.class));

	//
	// Channel inactive
	//
	channel.pipeline().fireChannelInactive();
	verify(callback, times(1)).onFailure(isA(ClosedChannelException.class));
}
 

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

/**
 * Tests that on reads the expected callback methods are called and read
 * buffers are recycled.
 */
@Test
public void testReadCallbacksAndBufferRecycling() throws Exception {
	final ClientHandlerCallback<KvStateResponse> callback = mock(ClientHandlerCallback.class);

	final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());
	final EmbeddedChannel channel = new EmbeddedChannel(new ClientHandler<>("Test Client", serializer, callback));

	final byte[] content = new byte[0];
	final KvStateResponse response = new KvStateResponse(content);

	//
	// Request success
	//
	ByteBuf buf = MessageSerializer.serializeResponse(channel.alloc(), 1222112277L, response);
	buf.skipBytes(4); // skip frame length

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onRequestResult(eq(1222112277L), any(KvStateResponse.class));
	assertEquals("Buffer not recycled", 0, buf.refCnt());

	//
	// Request failure
	//
	buf = MessageSerializer.serializeRequestFailure(
			channel.alloc(),
			1222112278,
			new RuntimeException("Expected test Exception"));
	buf.skipBytes(4); // skip frame length

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onRequestFailure(eq(1222112278L), isA(RuntimeException.class));
	assertEquals("Buffer not recycled", 0, buf.refCnt());

	//
	// Server failure
	//
	buf = MessageSerializer.serializeServerFailure(
			channel.alloc(),
			new RuntimeException("Expected test Exception"));
	buf.skipBytes(4); // skip frame length

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onFailure(isA(RuntimeException.class));

	//
	// Unexpected messages
	//
	buf = channel.alloc().buffer(4).writeInt(1223823);

	// Verify callback
	channel.writeInbound(buf);
	verify(callback, times(1)).onFailure(isA(IllegalStateException.class));
	assertEquals("Buffer not recycled", 0, buf.refCnt());

	//
	// Exception caught
	//
	channel.pipeline().fireExceptionCaught(new RuntimeException("Expected test Exception"));
	verify(callback, times(3)).onFailure(isA(RuntimeException.class));

	//
	// Channel inactive
	//
	channel.pipeline().fireChannelInactive();
	verify(callback, times(1)).onFailure(isA(ClosedChannelException.class));
}