Examples with KvStateInternalRequest org.apache.flink.queryablestate.messages.KvStateInternalRequest used on opensource projects

Example 1 with KvStateInternalRequest

use of org.apache.flink.queryablestate.messages.KvStateInternalRequest in project flink by apache.

the class ClientTest method testServerClosesChannel.

/**
 * Tests that a server channel close, closes the connection and removes it from the established
 * connections.
 */
@Test
public void testServerClosesChannel() throws Exception {
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer = new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());
    Client<KvStateInternalRequest, KvStateResponse> client = null;
    Channel serverChannel = null;
    try {
        client = new Client<>("Test Client", 1, serializer, stats);
        final LinkedBlockingQueue<ByteBuf> received = new LinkedBlockingQueue<>();
        final AtomicReference<Channel> channel = new AtomicReference<>();
        serverChannel = createServerChannel(new ChannelDataCollectingHandler(channel, received));
        InetSocketAddress serverAddress = getKvStateServerAddress(serverChannel);
        // Requests
        KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
        Future<KvStateResponse> future = client.sendRequest(serverAddress, request);
        received.take();
        assertEquals(1, stats.getNumConnections());
        channel.get().close().await();
        try {
            future.get();
            fail("Did not throw expected server failure");
        } catch (ExecutionException e) {
            if (!(e.getCause() instanceof ClosedChannelException)) {
                fail("Did not throw expected Exception");
            }
        // Expected
        }
        assertEquals(0L, stats.getNumConnections());
        // Counts can take some time to propagate
        while (stats.getNumSuccessful() != 0L || stats.getNumFailed() != 1L) {
            Thread.sleep(100L);
        }
        assertEquals(1L, stats.getNumRequests());
        assertEquals(0L, stats.getNumSuccessful());
        assertEquals(1L, stats.getNumFailed());
    } finally {
        if (client != null) {
            try {
                client.shutdown().get();
            } catch (Exception e) {
                e.printStackTrace();
            }
            Assert.assertTrue(client.isEventGroupShutdown());
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        assertEquals("Channel leak", 0L, stats.getNumConnections());
    }
}

Example 2 with KvStateInternalRequest

use of org.apache.flink.queryablestate.messages.KvStateInternalRequest in project flink by apache.

the class KvStateClientHandlerTest method testReadCallbacksAndBufferRecycling.

/**
 * 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);
    // skip frame length
    buf.skipBytes(4);
    // 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"));
    // skip frame length
    buf.skipBytes(4);
    // 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"));
    // skip frame length
    buf.skipBytes(4);
    // 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));
}

Example 3 with KvStateInternalRequest

use of org.apache.flink.queryablestate.messages.KvStateInternalRequest in project flink by apache.

the class KvStateServerHandlerTest method testChunkedResponse.

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

Example 4 with KvStateInternalRequest

use of org.apache.flink.queryablestate.messages.KvStateInternalRequest in project flink by apache.

the class KvStateServerHandlerTest method testCloseChannelOnExceptionCaught.

/**
 * 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);
    // skip frame length
    buf.skipBytes(4);
    // 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());
}

Example 5 with KvStateInternalRequest

use of org.apache.flink.queryablestate.messages.KvStateInternalRequest in project flink by apache.

the class KvStateServerHandlerTest method testQueryExecutorShutDown.

/**
 * 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);
    // skip frame length
    buf.skipBytes(4);
    // 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();
}