org.apache.flink.queryablestate.network.messages.MessageType Java Examples

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

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

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

/**
 * Tests request serialization.
 */
@Test
public void testRequestSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = randomByteArray(1024);

	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 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 request failure serialization.
 */
@Test
public void testKvStateRequestFailureSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1111222L;
	IllegalStateException cause = new IllegalStateException("Expected test");

	ByteBuf buf = MessageSerializer.serializeRequestFailure(alloc, requestId, cause);

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST_FAILURE, MessageSerializer.deserializeHeader(buf));
	RequestFailure requestFailure = MessageSerializer.deserializeRequestFailure(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(requestId, requestFailure.getRequestId());
	assertEquals(cause.getClass(), requestFailure.getCause().getClass());
	assertEquals(cause.getMessage(), requestFailure.getCause().getMessage());
}
 

/**
 * 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 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 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 request serialization.
 */
@Test
public void testRequestSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = randomByteArray(1024);

	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 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 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 request failure serialization.
 */
@Test
public void testKvStateRequestFailureSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1111222L;
	IllegalStateException cause = new IllegalStateException("Expected test");

	ByteBuf buf = MessageSerializer.serializeRequestFailure(alloc, requestId, cause);

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST_FAILURE, MessageSerializer.deserializeHeader(buf));
	RequestFailure requestFailure = MessageSerializer.deserializeRequestFailure(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(requestId, requestFailure.getRequestId());
	assertEquals(cause.getClass(), requestFailure.getCause().getClass());
	assertEquals(cause.getMessage(), requestFailure.getCause().getMessage());
}
 

/**
 * Tests request serialization.
 */
@Test
public void testRequestSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1337L;
	KvStateID kvStateId = new KvStateID();
	byte[] serializedKeyAndNamespace = randomByteArray(1024);

	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 request failure serialization.
 */
@Test
public void testKvStateRequestFailureSerialization() throws Exception {
	long requestId = Integer.MAX_VALUE + 1111222L;
	IllegalStateException cause = new IllegalStateException("Expected test");

	ByteBuf buf = MessageSerializer.serializeRequestFailure(alloc, requestId, cause);

	int frameLength = buf.readInt();
	assertEquals(MessageType.REQUEST_FAILURE, MessageSerializer.deserializeHeader(buf));
	RequestFailure requestFailure = MessageSerializer.deserializeRequestFailure(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(requestId, requestFailure.getRequestId());
	assertEquals(cause.getClass(), requestFailure.getCause().getClass());
	assertEquals(cause.getMessage(), requestFailure.getCause().getMessage());
}
 

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

@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
	try {
		ByteBuf buf = (ByteBuf) msg;
		MessageType msgType = MessageSerializer.deserializeHeader(buf);

		if (msgType == MessageType.REQUEST_RESULT) {
			long requestId = MessageSerializer.getRequestId(buf);
			RESP result = serializer.deserializeResponse(buf);
			callback.onRequestResult(requestId, result);
		} else if (msgType == MessageType.REQUEST_FAILURE) {
			RequestFailure failure = MessageSerializer.deserializeRequestFailure(buf);
			callback.onRequestFailure(failure.getRequestId(), failure.getCause());
		} else if (msgType == MessageType.SERVER_FAILURE) {
			throw MessageSerializer.deserializeServerFailure(buf);
		} else {
			throw new IllegalStateException("Unexpected response type '" + msgType + "'");
		}
	} catch (Throwable t1) {
		try {
			callback.onFailure(t1);
		} catch (Throwable t2) {
			LOG.error("Failed to notify callback about failure", t2);
		}
	} finally {
		ReferenceCountUtil.release(msg);
	}
}
 

/**
 * Tests server failure serialization.
 */
@Test
public void testServerFailureSerialization() throws Exception {
	IllegalStateException cause = new IllegalStateException("Expected test");

	ByteBuf buf = MessageSerializer.serializeServerFailure(alloc, cause);

	int frameLength = buf.readInt();
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable request = MessageSerializer.deserializeServerFailure(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(cause.getClass(), request.getClass());
	assertEquals(cause.getMessage(), request.getMessage());
}
 

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

@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
	try {
		ByteBuf buf = (ByteBuf) msg;
		MessageType msgType = MessageSerializer.deserializeHeader(buf);

		if (msgType == MessageType.REQUEST_RESULT) {
			long requestId = MessageSerializer.getRequestId(buf);
			RESP result = serializer.deserializeResponse(buf);
			callback.onRequestResult(requestId, result);
		} else if (msgType == MessageType.REQUEST_FAILURE) {
			RequestFailure failure = MessageSerializer.deserializeRequestFailure(buf);
			callback.onRequestFailure(failure.getRequestId(), failure.getCause());
		} else if (msgType == MessageType.SERVER_FAILURE) {
			throw MessageSerializer.deserializeServerFailure(buf);
		} else {
			throw new IllegalStateException("Unexpected response type '" + msgType + "'");
		}
	} catch (Throwable t1) {
		try {
			callback.onFailure(t1);
		} catch (Throwable t2) {
			LOG.error("Failed to notify callback about failure", t2);
		}
	} finally {
		ReferenceCountUtil.release(msg);
	}
}
 

/**
 * Tests server failure serialization.
 */
@Test
public void testServerFailureSerialization() throws Exception {
	IllegalStateException cause = new IllegalStateException("Expected test");

	ByteBuf buf = MessageSerializer.serializeServerFailure(alloc, cause);

	int frameLength = buf.readInt();
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable request = MessageSerializer.deserializeServerFailure(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(cause.getClass(), request.getClass());
	assertEquals(cause.getMessage(), request.getMessage());
}
 

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

	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 server failure serialization.
 */
@Test
public void testServerFailureSerialization() throws Exception {
	IllegalStateException cause = new IllegalStateException("Expected test");

	ByteBuf buf = MessageSerializer.serializeServerFailure(alloc, cause);

	int frameLength = buf.readInt();
	assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
	Throwable request = MessageSerializer.deserializeServerFailure(buf);
	assertEquals(buf.readerIndex(), frameLength + 4);

	assertEquals(cause.getClass(), request.getClass());
	assertEquals(cause.getMessage(), request.getMessage());
}
 

@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
	try {
		ByteBuf buf = (ByteBuf) msg;
		MessageType msgType = MessageSerializer.deserializeHeader(buf);

		if (msgType == MessageType.REQUEST_RESULT) {
			long requestId = MessageSerializer.getRequestId(buf);
			RESP result = serializer.deserializeResponse(buf);
			callback.onRequestResult(requestId, result);
		} else if (msgType == MessageType.REQUEST_FAILURE) {
			RequestFailure failure = MessageSerializer.deserializeRequestFailure(buf);
			callback.onRequestFailure(failure.getRequestId(), failure.getCause());
		} else if (msgType == MessageType.SERVER_FAILURE) {
			throw MessageSerializer.deserializeServerFailure(buf);
		} else {
			throw new IllegalStateException("Unexpected response type '" + msgType + "'");
		}
	} catch (Throwable t1) {
		try {
			callback.onFailure(t1);
		} catch (Throwable t2) {
			LOG.error("Failed to notify callback about failure", t2);
		}
	} finally {
		ReferenceCountUtil.release(msg);
	}
}
 

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

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

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

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

	localTestServer.shutdown();
}
 

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