Java Code Examples for org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf#skipBytes()

@Test(expected = IndexOutOfBoundsException.class)
public void testReaderIndexLargerThanWriterIndex() {
    String content1 = "hello";
    String content2 = "world";
    int length = content1.length() + content2.length();
    ByteBuf buffer = newBuffer(length);
    buffer.setIndex(0, 0);
    buffer.writeCharSequence(content1, CharsetUtil.US_ASCII);
    buffer.markWriterIndex();
    buffer.skipBytes(content1.length());
    buffer.writeCharSequence(content2, CharsetUtil.US_ASCII);
    buffer.skipBytes(content2.length());
    assertTrue(buffer.readerIndex() <= buffer.writerIndex());

    try {
        buffer.resetWriterIndex();
    } finally {
        buffer.release();
    }
}
 

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

@Test(expected = IndexOutOfBoundsException.class)
public void testReaderIndexLargerThanWriterIndex() {
    String content1 = "hello";
    String content2 = "world";
    int length = content1.length() + content2.length();
    ByteBuf buffer = newBuffer(length);
    buffer.setIndex(0, 0);
    buffer.writeCharSequence(content1, CharsetUtil.US_ASCII);
    buffer.markWriterIndex();
    buffer.skipBytes(content1.length());
    buffer.writeCharSequence(content2, CharsetUtil.US_ASCII);
    buffer.skipBytes(content2.length());
    assertTrue(buffer.readerIndex() <= buffer.writerIndex());

    try {
        buffer.resetWriterIndex();
    } finally {
        buffer.release();
    }
}
 

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

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

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

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

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

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