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

Example 6 with MessageSerializer

use of org.apache.flink.queryablestate.network.messages.MessageSerializer in project flink by apache.

the class KvStateServerHandlerTest method testQueryUnknownKey.

/**
 * 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);
    // 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 UnknownKeyOrNamespaceException);
    assertEquals(1L, stats.getNumRequests());
    assertEquals(1L, stats.getNumFailed());
}

Example 7 with MessageSerializer

use of org.apache.flink.queryablestate.network.messages.MessageSerializer in project flink by apache.

the class KvStateServerHandlerTest method testSerializerMismatch.

/**
 * Tests the failure response if the serializers don't match.
 */
@Test
public void testSerializerMismatch() 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);
    AbstractKeyedStateBackend<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");
    ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
    int key = 99812822;
    // Update the KvState
    backend.setCurrentKey(key);
    state.update(712828289);
    byte[] wrongKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace("wrong-key-type", StringSerializer.INSTANCE, "wrong-namespace-type", StringSerializer.INSTANCE);
    byte[] wrongNamespace = KvStateSerializer.serializeKeyAndNamespace(key, IntSerializer.INSTANCE, "wrong-namespace-type", StringSerializer.INSTANCE);
    assertTrue(registryListener.registrationName.equals("vanilla"));
    KvStateInternalRequest request = new KvStateInternalRequest(registryListener.kvStateId, wrongKeyAndNamespace);
    ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), 182828L, 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(182828L, response.getRequestId());
    assertTrue(response.getCause().getMessage().contains("IOException"));
    // Repeat with wrong namespace only
    request = new KvStateInternalRequest(registryListener.kvStateId, wrongNamespace);
    serRequest = MessageSerializer.serializeRequest(channel.alloc(), 182829L, request);
    // Write the request and wait for the response
    channel.writeInbound(serRequest);
    buf = (ByteBuf) readInboundBlocking(channel);
    // skip frame length
    buf.skipBytes(4);
    // Verify the response
    assertEquals(MessageType.REQUEST_FAILURE, MessageSerializer.deserializeHeader(buf));
    response = MessageSerializer.deserializeRequestFailure(buf);
    buf.release();
    assertEquals(182829L, response.getRequestId());
    assertTrue(response.getCause().getMessage().contains("IOException"));
    assertEquals(2L, stats.getNumRequests());
    assertEquals(2L, stats.getNumFailed());
}

Example 8 with MessageSerializer

use of org.apache.flink.queryablestate.network.messages.MessageSerializer in project flink by apache.

the class KvStateServerHandlerTest method testUnexpectedMessage.

/**
 * Tests response on unexpected messages.
 */
@Test
public void testUnexpectedMessage() 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);
    // Write the request and wait for the response
    ByteBuf unexpectedMessage = Unpooled.buffer(8);
    unexpectedMessage.writeInt(4);
    unexpectedMessage.writeInt(123238213);
    channel.writeInbound(unexpectedMessage);
    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();
    assertEquals(0L, stats.getNumRequests());
    assertEquals(0L, stats.getNumFailed());
    KvStateResponse stateResponse = new KvStateResponse(new byte[0]);
    unexpectedMessage = MessageSerializer.serializeResponse(channel.alloc(), 192L, stateResponse);
    channel.writeInbound(unexpectedMessage);
    buf = (ByteBuf) readInboundBlocking(channel);
    // skip frame length
    buf.skipBytes(4);
    // Verify the response
    assertEquals(MessageType.SERVER_FAILURE, MessageSerializer.deserializeHeader(buf));
    response = MessageSerializer.deserializeServerFailure(buf);
    buf.release();
    assertTrue("Unexpected failure cause " + response.getClass().getName(), response instanceof IllegalArgumentException);
    assertEquals(0L, stats.getNumRequests());
    assertEquals(0L, stats.getNumFailed());
}

Example 9 with MessageSerializer

use of org.apache.flink.queryablestate.network.messages.MessageSerializer in project flink by apache.

the class ClientTest method testConcurrentQueries.

/**
 * Multiple threads concurrently fire queries.
 */
@Test
public void testConcurrentQueries() throws Exception {
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    final MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer = new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());
    ExecutorService executor = null;
    Client<KvStateInternalRequest, KvStateResponse> 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 Client<>("Test Client", 1, serializer, stats);
        serverChannel = createServerChannel(new RespondingChannelHandler(serializer, serializedResult));
        final InetSocketAddress serverAddress = getKvStateServerAddress(serverChannel);
        final Client<KvStateInternalRequest, KvStateResponse> finalClient = client;
        Callable<List<CompletableFuture<KvStateResponse>>> queryTask = () -> {
            List<CompletableFuture<KvStateResponse>> results = new ArrayList<>(numQueriesPerTask);
            for (int i = 0; i < numQueriesPerTask; i++) {
                KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
                results.add(finalClient.sendRequest(serverAddress, request));
            }
            return results;
        };
        // Submit query tasks
        List<Future<List<CompletableFuture<KvStateResponse>>>> futures = new ArrayList<>();
        for (int i = 0; i < numQueryTasks; i++) {
            futures.add(executor.submit(queryTask));
        }
        // Verify results
        for (Future<List<CompletableFuture<KvStateResponse>>> future : futures) {
            List<CompletableFuture<KvStateResponse>> results = future.get();
            for (CompletableFuture<KvStateResponse> result : results) {
                KvStateResponse actual = result.get();
                assertArrayEquals(serializedResult, actual.getContent());
            }
        }
        int totalQueries = numQueryTasks * numQueriesPerTask;
        // Counts can take some time to propagate
        while (stats.getNumSuccessful() != totalQueries) {
            Thread.sleep(100L);
        }
        assertEquals(totalQueries, stats.getNumRequests());
        assertEquals(totalQueries, stats.getNumSuccessful());
    } finally {
        if (executor != null) {
            executor.shutdown();
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        if (client != null) {
            try {
                client.shutdown().get();
            } catch (Exception e) {
                e.printStackTrace();
            }
            Assert.assertTrue(client.isEventGroupShutdown());
        }
        assertEquals("Channel leak", 0L, stats.getNumConnections());
    }
}

Example 10 with MessageSerializer

use of org.apache.flink.queryablestate.network.messages.MessageSerializer in project flink by apache.

the class ClientTest method testFailureClosesChannel.

/**
 * Tests that a server failure closes the connection and removes it from the established
 * connections.
 */
@Test
public void testFailureClosesChannel() 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
        List<Future<KvStateResponse>> futures = new ArrayList<>();
        KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
        futures.add(client.sendRequest(serverAddress, request));
        futures.add(client.sendRequest(serverAddress, request));
        ByteBuf buf = received.take();
        assertNotNull("Receive timed out", buf);
        buf.release();
        buf = received.take();
        assertNotNull("Receive timed out", buf);
        buf.release();
        assertEquals(1L, stats.getNumConnections());
        Channel ch = channel.get();
        assertNotNull("Channel not active", ch);
        // Respond with failure
        ch.writeAndFlush(MessageSerializer.serializeServerFailure(serverChannel.alloc(), new RuntimeException("Expected test server failure")));
        try {
            futures.remove(0).get();
            fail("Did not throw expected server failure");
        } catch (ExecutionException e) {
            if (!(e.getCause() instanceof RuntimeException)) {
                fail("Did not throw expected Exception");
            }
        // Expected
        }
        try {
            futures.remove(0).get();
            fail("Did not throw expected server failure");
        } catch (ExecutionException e) {
            if (!(e.getCause() instanceof RuntimeException)) {
                fail("Did not throw expected Exception");
            }
        // Expected
        }
        assertEquals(0L, stats.getNumConnections());
        // Counts can take some time to propagate
        while (stats.getNumSuccessful() != 0L || stats.getNumFailed() != 2L) {
            Thread.sleep(100L);
        }
        assertEquals(2L, stats.getNumRequests());
        assertEquals(0L, stats.getNumSuccessful());
        assertEquals(2L, 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());
    }
}