Examples with AbstractStateBackend org.apache.flink.runtime.state.AbstractStateBackend used on opensource projects

Example 1 with AbstractStateBackend

use of org.apache.flink.runtime.state.AbstractStateBackend in project flink by apache.

the class KvStateServerTest method testSimpleRequest.

/**
	 * Tests a simple successful query via a SocketChannel.
	 */
@Test
public void testSimpleRequest() throws Exception {
    KvStateServer server = null;
    Bootstrap bootstrap = null;
    try {
        KvStateRegistry registry = new KvStateRegistry();
        KvStateRequestStats stats = new AtomicKvStateRequestStats();
        server = new KvStateServer(InetAddress.getLocalHost(), 0, 1, 1, registry, stats);
        server.start();
        KvStateServerAddress serverAddress = server.getAddress();
        int numKeyGroups = 1;
        AbstractStateBackend abstractBackend = new MemoryStateBackend();
        DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
        dummyEnv.setKvStateRegistry(registry);
        AbstractKeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), registry.createTaskRegistry(new JobID(), new JobVertexID()));
        final KvStateServerHandlerTest.TestRegistryListener registryListener = new KvStateServerHandlerTest.TestRegistryListener();
        registry.registerListener(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 = KvStateRequestSerializer.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.getHost(), serverAddress.getPort()).sync().channel();
        long requestId = Integer.MAX_VALUE + 182828L;
        assertTrue(registryListener.registrationName.equals("vanilla"));
        ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(channel.alloc(), requestId, registryListener.kvStateId, serializedKeyAndNamespace);
        channel.writeAndFlush(request);
        ByteBuf buf = responses.poll(TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
        assertEquals(KvStateRequestType.REQUEST_RESULT, KvStateRequestSerializer.deserializeHeader(buf));
        KvStateRequestResult response = KvStateRequestSerializer.deserializeKvStateRequestResult(buf);
        assertEquals(requestId, response.getRequestId());
        int actualValue = KvStateRequestSerializer.deserializeValue(response.getSerializedResult(), IntSerializer.INSTANCE);
        assertEquals(expectedValue, actualValue);
    } finally {
        if (server != null) {
            server.shutDown();
        }
        if (bootstrap != null) {
            EventLoopGroup group = bootstrap.group();
            if (group != null) {
                group.shutdownGracefully();
            }
        }
    }
}

Example 2 with AbstractStateBackend

use of org.apache.flink.runtime.state.AbstractStateBackend 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 3 with AbstractStateBackend

use of org.apache.flink.runtime.state.AbstractStateBackend 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();
}

Example 4 with AbstractStateBackend

use of org.apache.flink.runtime.state.AbstractStateBackend 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 5 with AbstractStateBackend

use of org.apache.flink.runtime.state.AbstractStateBackend 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());
}