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

Example 41 with KvStateResponse

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

the class ClientTest method testSimpleRequests.

/**
 * Tests simple queries, of which half succeed and half fail.
 */
@Test
public void testSimpleRequests() throws Exception {
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    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);
        // Random result
        final byte[] expected = new byte[1024];
        ThreadLocalRandom.current().nextBytes(expected);
        final LinkedBlockingQueue<ByteBuf> received = new LinkedBlockingQueue<>();
        final AtomicReference<Channel> channel = new AtomicReference<>();
        serverChannel = createServerChannel(new ChannelDataCollectingHandler(channel, received));
        InetSocketAddress serverAddress = getKvStateServerAddress(serverChannel);
        long numQueries = 1024L;
        List<CompletableFuture<KvStateResponse>> futures = new ArrayList<>();
        for (long i = 0L; i < numQueries; i++) {
            KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
            futures.add(client.sendRequest(serverAddress, request));
        }
        // Respond to messages
        Exception testException = new RuntimeException("Expected test Exception");
        for (long i = 0L; i < numQueries; i++) {
            ByteBuf buf = received.take();
            assertNotNull("Receive timed out", buf);
            Channel ch = channel.get();
            assertNotNull("Channel not active", ch);
            assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
            long requestId = MessageSerializer.getRequestId(buf);
            KvStateInternalRequest deserRequest = serializer.deserializeRequest(buf);
            buf.release();
            if (i % 2L == 0L) {
                ByteBuf response = MessageSerializer.serializeResponse(serverChannel.alloc(), requestId, new KvStateResponse(expected));
                ch.writeAndFlush(response);
            } else {
                ByteBuf response = MessageSerializer.serializeRequestFailure(serverChannel.alloc(), requestId, testException);
                ch.writeAndFlush(response);
            }
        }
        for (long i = 0L; i < numQueries; i++) {
            if (i % 2L == 0L) {
                KvStateResponse serializedResult = futures.get((int) i).get();
                assertArrayEquals(expected, serializedResult.getContent());
            } else {
                try {
                    futures.get((int) i).get();
                    fail("Did not throw expected Exception");
                } catch (ExecutionException e) {
                    if (!(e.getCause() instanceof RuntimeException)) {
                        fail("Did not throw expected Exception");
                    }
                // else expected
                }
            }
        }
        assertEquals(numQueries, stats.getNumRequests());
        long expectedRequests = numQueries / 2L;
        // Counts can take some time to propagate
        while (stats.getNumSuccessful() != expectedRequests || stats.getNumFailed() != expectedRequests) {
            Thread.sleep(100L);
        }
        assertEquals(expectedRequests, stats.getNumSuccessful());
        assertEquals(expectedRequests, stats.getNumFailed());
    } finally {
        if (client != null) {
            Exception exc = null;
            try {
                // todo here we were seeing this problem:
                // https://github.com/netty/netty/issues/4357 if we do a get().
                // this is why we now simply wait a bit so that everything is
                // shut down and then we check
                client.shutdown().get();
            } catch (Exception e) {
                exc = e;
                LOG.error("An exception occurred while shutting down netty.", e);
            }
            Assert.assertTrue(ExceptionUtils.stringifyException(exc), client.isEventGroupShutdown());
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        assertEquals("Channel leak", 0L, stats.getNumConnections());
    }
}

Example 42 with KvStateResponse

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

the class KvStateServerHandlerTest method testSimpleQuery.

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

Example 43 with KvStateResponse

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

the class KvStateServerHandlerTest method testFailureOnGetSerializedValue.

/**
 * 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, getClass().getClassLoader());
    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);
    // 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("Expected test Exception"));
    assertEquals(1L, stats.getNumRequests());
    assertEquals(1L, stats.getNumFailed());
}

Example 44 with KvStateResponse

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

the class KvStateServerHandlerTest method testQueryUnknownKvStateID.

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

Example 45 with KvStateResponse

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

the class KvStateServerHandlerTest method testIncomingBufferIsRecycled.

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