Examples with KvStateServerAddress org.apache.flink.runtime.query.KvStateServerAddress used on opensource projects

Example 6 with KvStateServerAddress

use of org.apache.flink.runtime.query.KvStateServerAddress in project flink by apache.

the class KvStateClientTest method testConcurrentQueries.

/**
	 * Multiple threads concurrently fire queries.
	 */
@Test
public void testConcurrentQueries() throws Exception {
    Deadline deadline = TEST_TIMEOUT.fromNow();
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    ExecutorService executor = null;
    KvStateClient client = null;
    Channel serverChannel = null;
    final byte[] serializedResult = new byte[1024];
    ThreadLocalRandom.current().nextBytes(serializedResult);
    try {
        int numQueryTasks = 4;
        final int numQueriesPerTask = 1024;
        executor = Executors.newFixedThreadPool(numQueryTasks);
        client = new KvStateClient(1, stats);
        serverChannel = createServerChannel(new ChannelInboundHandlerAdapter() {

            @Override
            public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                ByteBuf buf = (ByteBuf) msg;
                assertEquals(KvStateRequestType.REQUEST, KvStateRequestSerializer.deserializeHeader(buf));
                KvStateRequest request = KvStateRequestSerializer.deserializeKvStateRequest(buf);
                buf.release();
                ByteBuf response = KvStateRequestSerializer.serializeKvStateRequestResult(ctx.alloc(), request.getRequestId(), serializedResult);
                ctx.channel().writeAndFlush(response);
            }
        });
        final KvStateServerAddress serverAddress = getKvStateServerAddress(serverChannel);
        final KvStateClient finalClient = client;
        Callable<List<Future<byte[]>>> queryTask = new Callable<List<Future<byte[]>>>() {

            @Override
            public List<Future<byte[]>> call() throws Exception {
                List<Future<byte[]>> results = new ArrayList<>(numQueriesPerTask);
                for (int i = 0; i < numQueriesPerTask; i++) {
                    results.add(finalClient.getKvState(serverAddress, new KvStateID(), new byte[0]));
                }
                return results;
            }
        };
        // Submit query tasks
        List<java.util.concurrent.Future<List<Future<byte[]>>>> futures = new ArrayList<>();
        for (int i = 0; i < numQueryTasks; i++) {
            futures.add(executor.submit(queryTask));
        }
        // Verify results
        for (java.util.concurrent.Future<List<Future<byte[]>>> future : futures) {
            List<Future<byte[]>> results = future.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
            for (Future<byte[]> result : results) {
                byte[] actual = Await.result(result, deadline.timeLeft());
                assertArrayEquals(serializedResult, actual);
            }
        }
        int totalQueries = numQueryTasks * numQueriesPerTask;
        // Counts can take some time to propagate
        while (deadline.hasTimeLeft() && stats.getNumSuccessful() != totalQueries) {
            Thread.sleep(100);
        }
        assertEquals(totalQueries, stats.getNumRequests());
        assertEquals(totalQueries, stats.getNumSuccessful());
    } finally {
        if (executor != null) {
            executor.shutdown();
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        if (client != null) {
            client.shutDown();
        }
        assertEquals("Channel leak", 0, stats.getNumConnections());
    }
}

Example 7 with KvStateServerAddress

use of org.apache.flink.runtime.query.KvStateServerAddress in project flink by apache.

the class KvStateClientTest method testSimpleRequests.

/**
	 * Tests simple queries, of which half succeed and half fail.
	 */
@Test
public void testSimpleRequests() throws Exception {
    Deadline deadline = TEST_TIMEOUT.fromNow();
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    KvStateClient client = null;
    Channel serverChannel = null;
    try {
        client = new KvStateClient(1, 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 ChannelInboundHandlerAdapter() {

            @Override
            public void channelActive(ChannelHandlerContext ctx) throws Exception {
                channel.set(ctx.channel());
            }

            @Override
            public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                received.add((ByteBuf) msg);
            }
        });
        KvStateServerAddress serverAddress = getKvStateServerAddress(serverChannel);
        List<Future<byte[]>> futures = new ArrayList<>();
        int numQueries = 1024;
        for (int i = 0; i < numQueries; i++) {
            futures.add(client.getKvState(serverAddress, new KvStateID(), new byte[0]));
        }
        // Respond to messages
        Exception testException = new RuntimeException("Expected test Exception");
        for (int i = 0; i < numQueries; i++) {
            ByteBuf buf = received.poll(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
            assertNotNull("Receive timed out", buf);
            Channel ch = channel.get();
            assertNotNull("Channel not active", ch);
            assertEquals(KvStateRequestType.REQUEST, KvStateRequestSerializer.deserializeHeader(buf));
            KvStateRequest request = KvStateRequestSerializer.deserializeKvStateRequest(buf);
            buf.release();
            if (i % 2 == 0) {
                ByteBuf response = KvStateRequestSerializer.serializeKvStateRequestResult(serverChannel.alloc(), request.getRequestId(), expected);
                ch.writeAndFlush(response);
            } else {
                ByteBuf response = KvStateRequestSerializer.serializeKvStateRequestFailure(serverChannel.alloc(), request.getRequestId(), testException);
                ch.writeAndFlush(response);
            }
        }
        for (int i = 0; i < numQueries; i++) {
            if (i % 2 == 0) {
                byte[] serializedResult = Await.result(futures.get(i), deadline.timeLeft());
                assertArrayEquals(expected, serializedResult);
            } else {
                try {
                    Await.result(futures.get(i), deadline.timeLeft());
                    fail("Did not throw expected Exception");
                } catch (RuntimeException ignored) {
                // Expected
                }
            }
        }
        assertEquals(numQueries, stats.getNumRequests());
        int expectedRequests = numQueries / 2;
        // Counts can take some time to propagate
        while (deadline.hasTimeLeft() && (stats.getNumSuccessful() != expectedRequests || stats.getNumFailed() != expectedRequests)) {
            Thread.sleep(100);
        }
        assertEquals(expectedRequests, stats.getNumSuccessful());
        assertEquals(expectedRequests, stats.getNumFailed());
    } finally {
        if (client != null) {
            client.shutDown();
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        assertEquals("Channel leak", 0, stats.getNumConnections());
    }
}

Example 8 with KvStateServerAddress

use of org.apache.flink.runtime.query.KvStateServerAddress in project flink by apache.

the class KvStateClientTest method testFailureClosesChannel.

/**
	 * Tests that a server failure closes the connection and removes it from
	 * the established connections.
	 */
@Test
public void testFailureClosesChannel() throws Exception {
    Deadline deadline = TEST_TIMEOUT.fromNow();
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    KvStateClient client = null;
    Channel serverChannel = null;
    try {
        client = new KvStateClient(1, stats);
        final LinkedBlockingQueue<ByteBuf> received = new LinkedBlockingQueue<>();
        final AtomicReference<Channel> channel = new AtomicReference<>();
        serverChannel = createServerChannel(new ChannelInboundHandlerAdapter() {

            @Override
            public void channelActive(ChannelHandlerContext ctx) throws Exception {
                channel.set(ctx.channel());
            }

            @Override
            public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                received.add((ByteBuf) msg);
            }
        });
        KvStateServerAddress serverAddress = getKvStateServerAddress(serverChannel);
        // Requests
        List<Future<byte[]>> futures = new ArrayList<>();
        futures.add(client.getKvState(serverAddress, new KvStateID(), new byte[0]));
        futures.add(client.getKvState(serverAddress, new KvStateID(), new byte[0]));
        ByteBuf buf = received.poll(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
        assertNotNull("Receive timed out", buf);
        buf.release();
        buf = received.poll(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
        assertNotNull("Receive timed out", buf);
        buf.release();
        assertEquals(1, stats.getNumConnections());
        Channel ch = channel.get();
        assertNotNull("Channel not active", ch);
        // Respond with failure
        ch.writeAndFlush(KvStateRequestSerializer.serializeServerFailure(serverChannel.alloc(), new RuntimeException("Expected test server failure")));
        try {
            Await.result(futures.remove(0), deadline.timeLeft());
            fail("Did not throw expected server failure");
        } catch (RuntimeException ignored) {
        // Expected
        }
        try {
            Await.result(futures.remove(0), deadline.timeLeft());
            fail("Did not throw expected server failure");
        } catch (RuntimeException ignored) {
        // Expected
        }
        assertEquals(0, stats.getNumConnections());
        // Counts can take some time to propagate
        while (deadline.hasTimeLeft() && (stats.getNumSuccessful() != 0 || stats.getNumFailed() != 2)) {
            Thread.sleep(100);
        }
        assertEquals(2, stats.getNumRequests());
        assertEquals(0, stats.getNumSuccessful());
        assertEquals(2, stats.getNumFailed());
    } finally {
        if (client != null) {
            client.shutDown();
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        assertEquals("Channel leak", 0, stats.getNumConnections());
    }
}