Search in sources :

Example 11 with AtomicKvStateRequestStats

use of org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats in project flink by apache.

the class ClientTest method testClientServerIntegration.

/**
 * Tests multiple clients querying multiple servers until 100k queries have been processed. At
 * this point, the client is shut down and its verified that all ongoing requests are failed.
 */
@Test
public void testClientServerIntegration() throws Throwable {
    // Config
    final int numServers = 2;
    final int numServerEventLoopThreads = 2;
    final int numServerQueryThreads = 2;
    final int numClientEventLoopThreads = 4;
    final int numClientsTasks = 8;
    final int batchSize = 16;
    final int numKeyGroups = 1;
    AbstractStateBackend abstractBackend = new MemoryStateBackend();
    KvStateRegistry dummyRegistry = new KvStateRegistry();
    DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
    dummyEnv.setKvStateRegistry(dummyRegistry);
    AbstractKeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), dummyRegistry.createTaskRegistry(new JobID(), new JobVertexID()), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
    AtomicKvStateRequestStats clientStats = new AtomicKvStateRequestStats();
    final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer = new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());
    Client<KvStateInternalRequest, KvStateResponse> client = null;
    ExecutorService clientTaskExecutor = null;
    final KvStateServerImpl[] server = new KvStateServerImpl[numServers];
    try {
        client = new Client<>("Test Client", numClientEventLoopThreads, serializer, clientStats);
        clientTaskExecutor = Executors.newFixedThreadPool(numClientsTasks);
        // Create state
        ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
        desc.setQueryable("any");
        // Create servers
        KvStateRegistry[] registry = new KvStateRegistry[numServers];
        AtomicKvStateRequestStats[] serverStats = new AtomicKvStateRequestStats[numServers];
        final KvStateID[] ids = new KvStateID[numServers];
        for (int i = 0; i < numServers; i++) {
            registry[i] = new KvStateRegistry();
            serverStats[i] = new AtomicKvStateRequestStats();
            server[i] = new KvStateServerImpl(InetAddress.getLocalHost().getHostName(), Collections.singletonList(0).iterator(), numServerEventLoopThreads, numServerQueryThreads, registry[i], serverStats[i]);
            server[i].start();
            backend.setCurrentKey(1010 + i);
            // Value per server
            ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
            state.update(201 + i);
            // we know it must be a KvState but this is not exposed to the user via State
            InternalKvState<Integer, ?, Integer> kvState = (InternalKvState<Integer, ?, Integer>) state;
            // Register KvState (one state instance for all server)
            ids[i] = registry[i].registerKvState(new JobID(), new JobVertexID(), new KeyGroupRange(0, 0), "any", kvState, getClass().getClassLoader());
        }
        final Client<KvStateInternalRequest, KvStateResponse> finalClient = client;
        Callable<Void> queryTask = () -> {
            while (true) {
                if (Thread.interrupted()) {
                    throw new InterruptedException();
                }
                // Random server permutation
                List<Integer> random = new ArrayList<>();
                for (int j = 0; j < batchSize; j++) {
                    random.add(j);
                }
                Collections.shuffle(random);
                // Dispatch queries
                List<CompletableFuture<KvStateResponse>> futures = new ArrayList<>(batchSize);
                for (int j = 0; j < batchSize; j++) {
                    int targetServer = random.get(j) % numServers;
                    byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(1010 + targetServer, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
                    KvStateInternalRequest request = new KvStateInternalRequest(ids[targetServer], serializedKeyAndNamespace);
                    futures.add(finalClient.sendRequest(server[targetServer].getServerAddress(), request));
                }
                // Verify results
                for (int j = 0; j < batchSize; j++) {
                    int targetServer = random.get(j) % numServers;
                    Future<KvStateResponse> future = futures.get(j);
                    byte[] buf = future.get().getContent();
                    int value = KvStateSerializer.deserializeValue(buf, IntSerializer.INSTANCE);
                    assertEquals(201L + targetServer, value);
                }
            }
        };
        // Submit tasks
        List<Future<Void>> taskFutures = new ArrayList<>();
        for (int i = 0; i < numClientsTasks; i++) {
            taskFutures.add(clientTaskExecutor.submit(queryTask));
        }
        long numRequests;
        while ((numRequests = clientStats.getNumRequests()) < 100_000L) {
            Thread.sleep(100L);
            LOG.info("Number of requests {}/100_000", numRequests);
        }
        try {
            client.shutdown().get();
        } catch (Exception e) {
            e.printStackTrace();
        }
        Assert.assertTrue(client.isEventGroupShutdown());
        final CombinableMatcher<Throwable> exceptionMatcher = either(FlinkMatchers.containsCause(ClosedChannelException.class)).or(FlinkMatchers.containsCause(IllegalStateException.class));
        for (Future<Void> future : taskFutures) {
            try {
                future.get();
                fail("Did not throw expected Exception after shut down");
            } catch (ExecutionException t) {
                assertThat(t, exceptionMatcher);
            }
        }
        assertEquals("Connection leak (client)", 0L, clientStats.getNumConnections());
        for (int i = 0; i < numServers; i++) {
            boolean success = false;
            int numRetries = 0;
            while (!success) {
                try {
                    assertEquals("Connection leak (server)", 0L, serverStats[i].getNumConnections());
                    success = true;
                } catch (Throwable t) {
                    if (numRetries < 10) {
                        LOG.info("Retrying connection leak check (server)");
                        Thread.sleep((numRetries + 1) * 50L);
                        numRetries++;
                    } else {
                        throw t;
                    }
                }
            }
        }
    } finally {
        if (client != null) {
            try {
                client.shutdown().get();
            } catch (Exception e) {
                e.printStackTrace();
            }
            Assert.assertTrue(client.isEventGroupShutdown());
        }
        for (int i = 0; i < numServers; i++) {
            if (server[i] != null) {
                server[i].shutdown();
            }
        }
        if (clientTaskExecutor != null) {
            clientTaskExecutor.shutdown();
        }
    }
}
Also used : KvStateRegistry(org.apache.flink.runtime.query.KvStateRegistry) MessageSerializer(org.apache.flink.queryablestate.network.messages.MessageSerializer) KeyGroupRange(org.apache.flink.runtime.state.KeyGroupRange) ArrayList(java.util.ArrayList) CloseableRegistry(org.apache.flink.core.fs.CloseableRegistry) KvStateServerImpl(org.apache.flink.queryablestate.server.KvStateServerImpl) ValueStateDescriptor(org.apache.flink.api.common.state.ValueStateDescriptor) KvStateID(org.apache.flink.queryablestate.KvStateID) List(java.util.List) ArrayList(java.util.ArrayList) KvStateInternalRequest(org.apache.flink.queryablestate.messages.KvStateInternalRequest) JobID(org.apache.flink.api.common.JobID) AtomicKvStateRequestStats(org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats) UnregisteredMetricsGroup(org.apache.flink.metrics.groups.UnregisteredMetricsGroup) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) JobVertexID(org.apache.flink.runtime.jobgraph.JobVertexID) KvStateResponse(org.apache.flink.queryablestate.messages.KvStateResponse) ExecutionException(java.util.concurrent.ExecutionException) AbstractStateBackend(org.apache.flink.runtime.state.AbstractStateBackend) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) ConnectException(java.net.ConnectException) ClosedChannelException(java.nio.channels.ClosedChannelException) UnknownHostException(java.net.UnknownHostException) ExecutionException(java.util.concurrent.ExecutionException) InternalKvState(org.apache.flink.runtime.state.internal.InternalKvState) ExecutorService(java.util.concurrent.ExecutorService) Future(java.util.concurrent.Future) CompletableFuture(java.util.concurrent.CompletableFuture) Test(org.junit.Test)

Example 12 with AtomicKvStateRequestStats

use of org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats in project flink by apache.

the class AbstractServerTest method testPortRangeSuccess.

/**
 * Tests that in case of port collision and big enough port range, the server will try to bind
 * to the next port in the range.
 */
@Test
public void testPortRangeSuccess() throws Throwable {
    AtomicKvStateRequestStats serverStats1 = new AtomicKvStateRequestStats();
    AtomicKvStateRequestStats serverStats2 = new AtomicKvStateRequestStats();
    AtomicKvStateRequestStats clientStats = new AtomicKvStateRequestStats();
    final int portRangeStart = 7777;
    final int portRangeEnd = 7900;
    List<Integer> portList = IntStream.range(portRangeStart, portRangeEnd + 1).boxed().collect(Collectors.toList());
    try (TestServer server1 = new TestServer("Test Server 1", serverStats1, portList.iterator());
        TestServer server2 = new TestServer("Test Server 2", serverStats2, portList.iterator());
        TestClient client = new TestClient("Test Client", 1, new MessageSerializer<>(new TestMessage.TestMessageDeserializer(), new TestMessage.TestMessageDeserializer()), clientStats)) {
        server1.start();
        Assert.assertTrue(server1.getServerAddress().getPort() >= portRangeStart && server1.getServerAddress().getPort() <= portRangeEnd);
        server2.start();
        Assert.assertTrue(server2.getServerAddress().getPort() >= portRangeStart && server2.getServerAddress().getPort() <= portRangeEnd);
        TestMessage response1 = client.sendRequest(server1.getServerAddress(), new TestMessage("ping")).join();
        Assert.assertEquals(server1.getServerName() + "-ping", response1.getMessage());
        TestMessage response2 = client.sendRequest(server2.getServerAddress(), new TestMessage("pong")).join();
        Assert.assertEquals(server2.getServerName() + "-pong", response2.getMessage());
        Assert.assertEquals(1L, serverStats1.getNumConnections());
        Assert.assertEquals(1L, serverStats2.getNumConnections());
        Assert.assertEquals(2L, clientStats.getNumConnections());
        Assert.assertEquals(0L, clientStats.getNumFailed());
        Assert.assertEquals(2L, clientStats.getNumSuccessful());
        Assert.assertEquals(2L, clientStats.getNumRequests());
    }
    Assert.assertEquals(0L, serverStats1.getNumConnections());
    Assert.assertEquals(0L, serverStats2.getNumConnections());
    Assert.assertEquals(0L, clientStats.getNumConnections());
    Assert.assertEquals(0L, clientStats.getNumFailed());
    Assert.assertEquals(2L, clientStats.getNumSuccessful());
    Assert.assertEquals(2L, clientStats.getNumRequests());
}
Also used : AtomicKvStateRequestStats(org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats) Test(org.junit.Test)

Example 13 with AtomicKvStateRequestStats

use of org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats in project flink by apache.

the class KvStateServerTest method testSimpleRequest.

/**
 * Tests a simple successful query via a SocketChannel.
 */
@Test
public void testSimpleRequest() throws Throwable {
    KvStateServerImpl server = null;
    Bootstrap bootstrap = null;
    try {
        KvStateRegistry registry = new KvStateRegistry();
        KvStateRequestStats stats = new AtomicKvStateRequestStats();
        server = new KvStateServerImpl(InetAddress.getLocalHost().getHostName(), Collections.singletonList(0).iterator(), 1, 1, registry, stats);
        server.start();
        InetSocketAddress serverAddress = server.getServerAddress();
        int numKeyGroups = 1;
        AbstractStateBackend abstractBackend = new MemoryStateBackend();
        DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
        dummyEnv.setKvStateRegistry(registry);
        final JobID jobId = new JobID();
        AbstractKeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(dummyEnv, jobId, "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), registry.createTaskRegistry(jobId, new JobVertexID()), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
        final KvStateServerHandlerTest.TestRegistryListener registryListener = new KvStateServerHandlerTest.TestRegistryListener();
        registry.registerListener(jobId, registryListener);
        ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
        desc.setQueryable("vanilla");
        ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
        // Update KvState
        int expectedValue = 712828289;
        int key = 99812822;
        backend.setCurrentKey(key);
        state.update(expectedValue);
        // Request
        byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(key, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
        // Connect to the server
        final BlockingQueue<ByteBuf> responses = new LinkedBlockingQueue<>();
        bootstrap = createBootstrap(new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE, 0, 4, 0, 4), new ChannelInboundHandlerAdapter() {

            @Override
            public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                responses.add((ByteBuf) msg);
            }
        });
        Channel channel = bootstrap.connect(serverAddress.getAddress(), serverAddress.getPort()).sync().channel();
        long requestId = Integer.MAX_VALUE + 182828L;
        assertTrue(registryListener.registrationName.equals("vanilla"));
        final KvStateInternalRequest request = new KvStateInternalRequest(registryListener.kvStateId, serializedKeyAndNamespace);
        ByteBuf serializeRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);
        channel.writeAndFlush(serializeRequest);
        ByteBuf buf = responses.poll(TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
        assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
        assertEquals(requestId, MessageSerializer.getRequestId(buf));
        KvStateResponse response = server.getSerializer().deserializeResponse(buf);
        int actualValue = KvStateSerializer.deserializeValue(response.getContent(), IntSerializer.INSTANCE);
        assertEquals(expectedValue, actualValue);
    } finally {
        if (server != null) {
            server.shutdown();
        }
        if (bootstrap != null) {
            EventLoopGroup group = bootstrap.group();
            if (group != null) {
                // note: no "quiet period" to not trigger Netty#4357
                group.shutdownGracefully(0, 10, TimeUnit.SECONDS);
            }
        }
    }
}
Also used : KvStateRegistry(org.apache.flink.runtime.query.KvStateRegistry) UnregisteredMetricsGroup(org.apache.flink.metrics.groups.UnregisteredMetricsGroup) InetSocketAddress(java.net.InetSocketAddress) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) JobVertexID(org.apache.flink.runtime.jobgraph.JobVertexID) KeyGroupRange(org.apache.flink.runtime.state.KeyGroupRange) ChannelHandlerContext(org.apache.flink.shaded.netty4.io.netty.channel.ChannelHandlerContext) CloseableRegistry(org.apache.flink.core.fs.CloseableRegistry) ByteBuf(org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) AtomicKvStateRequestStats(org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats) KvStateRequestStats(org.apache.flink.queryablestate.network.stats.KvStateRequestStats) KvStateServerImpl(org.apache.flink.queryablestate.server.KvStateServerImpl) ValueStateDescriptor(org.apache.flink.api.common.state.ValueStateDescriptor) Bootstrap(org.apache.flink.shaded.netty4.io.netty.bootstrap.Bootstrap) KvStateResponse(org.apache.flink.queryablestate.messages.KvStateResponse) LengthFieldBasedFrameDecoder(org.apache.flink.shaded.netty4.io.netty.handler.codec.LengthFieldBasedFrameDecoder) AbstractStateBackend(org.apache.flink.runtime.state.AbstractStateBackend) NioSocketChannel(org.apache.flink.shaded.netty4.io.netty.channel.socket.nio.NioSocketChannel) SocketChannel(org.apache.flink.shaded.netty4.io.netty.channel.socket.SocketChannel) Channel(org.apache.flink.shaded.netty4.io.netty.channel.Channel) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) EventLoopGroup(org.apache.flink.shaded.netty4.io.netty.channel.EventLoopGroup) NioEventLoopGroup(org.apache.flink.shaded.netty4.io.netty.channel.nio.NioEventLoopGroup) KvStateInternalRequest(org.apache.flink.queryablestate.messages.KvStateInternalRequest) AtomicKvStateRequestStats(org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats) JobID(org.apache.flink.api.common.JobID) ChannelInboundHandlerAdapter(org.apache.flink.shaded.netty4.io.netty.channel.ChannelInboundHandlerAdapter) Test(org.junit.Test)

Example 14 with AtomicKvStateRequestStats

use of org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats 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());
    }
}
Also used : MessageSerializer(org.apache.flink.queryablestate.network.messages.MessageSerializer) InetSocketAddress(java.net.InetSocketAddress) ArrayList(java.util.ArrayList) ByteBuf(org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) CompletableFuture(java.util.concurrent.CompletableFuture) KvStateResponse(org.apache.flink.queryablestate.messages.KvStateResponse) KvStateID(org.apache.flink.queryablestate.KvStateID) ExecutionException(java.util.concurrent.ExecutionException) SocketChannel(org.apache.flink.shaded.netty4.io.netty.channel.socket.SocketChannel) NioServerSocketChannel(org.apache.flink.shaded.netty4.io.netty.channel.socket.nio.NioServerSocketChannel) Channel(org.apache.flink.shaded.netty4.io.netty.channel.Channel) AtomicReference(java.util.concurrent.atomic.AtomicReference) ConnectException(java.net.ConnectException) ClosedChannelException(java.nio.channels.ClosedChannelException) UnknownHostException(java.net.UnknownHostException) ExecutionException(java.util.concurrent.ExecutionException) KvStateInternalRequest(org.apache.flink.queryablestate.messages.KvStateInternalRequest) AtomicKvStateRequestStats(org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats) Test(org.junit.Test)

Example 15 with AtomicKvStateRequestStats

use of org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats 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());
}
Also used : KvStateRegistry(org.apache.flink.runtime.query.KvStateRegistry) MessageSerializer(org.apache.flink.queryablestate.network.messages.MessageSerializer) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) EmbeddedChannel(org.apache.flink.shaded.netty4.io.netty.channel.embedded.EmbeddedChannel) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) ByteBuf(org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf) ValueStateDescriptor(org.apache.flink.api.common.state.ValueStateDescriptor) KvStateServerHandler(org.apache.flink.queryablestate.server.KvStateServerHandler) KvStateInternalRequest(org.apache.flink.queryablestate.messages.KvStateInternalRequest) KvStateResponse(org.apache.flink.queryablestate.messages.KvStateResponse) AtomicKvStateRequestStats(org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats) AbstractStateBackend(org.apache.flink.runtime.state.AbstractStateBackend) Test(org.junit.Test)

Aggregations

AtomicKvStateRequestStats (org.apache.flink.queryablestate.network.stats.AtomicKvStateRequestStats)18 Test (org.junit.Test)18 KvStateInternalRequest (org.apache.flink.queryablestate.messages.KvStateInternalRequest)17 KvStateResponse (org.apache.flink.queryablestate.messages.KvStateResponse)17 MessageSerializer (org.apache.flink.queryablestate.network.messages.MessageSerializer)16 ByteBuf (org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf)14 KvStateRegistry (org.apache.flink.runtime.query.KvStateRegistry)12 KvStateServerHandler (org.apache.flink.queryablestate.server.KvStateServerHandler)10 EmbeddedChannel (org.apache.flink.shaded.netty4.io.netty.channel.embedded.EmbeddedChannel)10 KvStateID (org.apache.flink.queryablestate.KvStateID)9 ValueStateDescriptor (org.apache.flink.api.common.state.ValueStateDescriptor)7 DummyEnvironment (org.apache.flink.runtime.operators.testutils.DummyEnvironment)7 AbstractStateBackend (org.apache.flink.runtime.state.AbstractStateBackend)7 MemoryStateBackend (org.apache.flink.runtime.state.memory.MemoryStateBackend)7 ConnectException (java.net.ConnectException)6 InetSocketAddress (java.net.InetSocketAddress)6 UnknownHostException (java.net.UnknownHostException)6 ClosedChannelException (java.nio.channels.ClosedChannelException)6 ExecutionException (java.util.concurrent.ExecutionException)6 RequestFailure (org.apache.flink.queryablestate.network.messages.RequestFailure)5