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

Example 1 with KvStateRegistry

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

the class TaskManagerServices method createNetworkEnvironment.

/**
	 * Creates the {@link NetworkEnvironment} from the given {@link TaskManagerServicesConfiguration}.
	 *
	 * @param taskManagerServicesConfiguration to construct the network environment from
	 * @return Network environment
	 * @throws IOException
	 */
private static NetworkEnvironment createNetworkEnvironment(TaskManagerServicesConfiguration taskManagerServicesConfiguration) throws IOException {
    NetworkEnvironmentConfiguration networkEnvironmentConfiguration = taskManagerServicesConfiguration.getNetworkConfig();
    NetworkBufferPool networkBufferPool = new NetworkBufferPool(networkEnvironmentConfiguration.numNetworkBuffers(), networkEnvironmentConfiguration.networkBufferSize(), networkEnvironmentConfiguration.memoryType());
    ConnectionManager connectionManager;
    if (networkEnvironmentConfiguration.nettyConfig() != null) {
        connectionManager = new NettyConnectionManager(networkEnvironmentConfiguration.nettyConfig());
    } else {
        connectionManager = new LocalConnectionManager();
    }
    ResultPartitionManager resultPartitionManager = new ResultPartitionManager();
    TaskEventDispatcher taskEventDispatcher = new TaskEventDispatcher();
    KvStateRegistry kvStateRegistry = new KvStateRegistry();
    KvStateServer kvStateServer;
    if (taskManagerServicesConfiguration.getQueryableStateConfig().enabled()) {
        QueryableStateConfiguration qsConfig = taskManagerServicesConfiguration.getQueryableStateConfig();
        int numNetworkThreads = qsConfig.numServerThreads() == 0 ? taskManagerServicesConfiguration.getNumberOfSlots() : qsConfig.numServerThreads();
        int numQueryThreads = qsConfig.numQueryThreads() == 0 ? taskManagerServicesConfiguration.getNumberOfSlots() : qsConfig.numQueryThreads();
        kvStateServer = new KvStateServer(taskManagerServicesConfiguration.getTaskManagerAddress(), qsConfig.port(), numNetworkThreads, numQueryThreads, kvStateRegistry, new DisabledKvStateRequestStats());
    } else {
        kvStateServer = null;
    }
    // we start the network first, to make sure it can allocate its buffers first
    return new NetworkEnvironment(networkBufferPool, connectionManager, resultPartitionManager, taskEventDispatcher, kvStateRegistry, kvStateServer, networkEnvironmentConfiguration.ioMode(), networkEnvironmentConfiguration.partitionRequestInitialBackoff(), networkEnvironmentConfiguration.partitionRequestMaxBackoff(), networkEnvironmentConfiguration.networkBuffersPerChannel(), networkEnvironmentConfiguration.extraNetworkBuffersPerGate());
}

Example 2 with KvStateRegistry

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

the class NetworkEnvironmentTest method testRegisterTaskUsesBoundedBuffers.

/**
	 * Verifies that {@link NetworkEnvironment#registerTask(Task)} sets up (un)bounded buffer pool
	 * instances for various types of input and output channels.
	 */
@Test
public void testRegisterTaskUsesBoundedBuffers() throws Exception {
    final NetworkEnvironment network = new NetworkEnvironment(new NetworkBufferPool(numBuffers, memorySegmentSize, MemoryType.HEAP), new LocalConnectionManager(), new ResultPartitionManager(), new TaskEventDispatcher(), new KvStateRegistry(), null, IOManager.IOMode.SYNC, 0, 0, 2, 8);
    // result partitions
    ResultPartition rp1 = createResultPartition(ResultPartitionType.PIPELINED, 2);
    ResultPartition rp2 = createResultPartition(ResultPartitionType.BLOCKING, 2);
    ResultPartition rp3 = createResultPartition(ResultPartitionType.PIPELINED_BOUNDED, 2);
    ResultPartition rp4 = createResultPartition(ResultPartitionType.PIPELINED_BOUNDED, 8);
    final ResultPartition[] resultPartitions = new ResultPartition[] { rp1, rp2, rp3, rp4 };
    final ResultPartitionWriter[] resultPartitionWriters = new ResultPartitionWriter[] { new ResultPartitionWriter(rp1), new ResultPartitionWriter(rp2), new ResultPartitionWriter(rp3), new ResultPartitionWriter(rp4) };
    // input gates
    final SingleInputGate[] inputGates = new SingleInputGate[] { createSingleInputGateMock(ResultPartitionType.PIPELINED, 2), createSingleInputGateMock(ResultPartitionType.BLOCKING, 2), createSingleInputGateMock(ResultPartitionType.PIPELINED_BOUNDED, 2), createSingleInputGateMock(ResultPartitionType.PIPELINED_BOUNDED, 8) };
    // overall task to register
    Task task = mock(Task.class);
    when(task.getProducedPartitions()).thenReturn(resultPartitions);
    when(task.getAllWriters()).thenReturn(resultPartitionWriters);
    when(task.getAllInputGates()).thenReturn(inputGates);
    network.registerTask(task);
    assertEquals(Integer.MAX_VALUE, rp1.getBufferPool().getMaxNumberOfMemorySegments());
    assertEquals(Integer.MAX_VALUE, rp2.getBufferPool().getMaxNumberOfMemorySegments());
    assertEquals(2 * 2 + 8, rp3.getBufferPool().getMaxNumberOfMemorySegments());
    assertEquals(8 * 2 + 8, rp4.getBufferPool().getMaxNumberOfMemorySegments());
    network.shutdown();
}

Example 3 with KvStateRegistry

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

the class StreamingRuntimeContextTest method createListPlainMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {
    AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
    ExecutionConfig config = new ExecutionConfig();
    KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);
    DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);
    when(operatorMock.getExecutionConfig()).thenReturn(config);
    doAnswer(new Answer<ListState<String>>() {

        @Override
        public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
            ListStateDescriptor<String> descr = (ListStateDescriptor<String>) invocationOnMock.getArguments()[2];
            AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
            backend.setCurrentKey(0);
            return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
        }
    }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    return operatorMock;
}

Example 4 with KvStateRegistry

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

the class KvStateServerHandlerTest method testQueryUnknownKvStateID.

/**
	 * Tests the failure response with {@link UnknownKvStateID} as cause on
	 * queries for unregistered KvStateIDs.
	 */
@Test
public void testQueryUnknownKvStateID() throws Exception {
    KvStateRegistry registry = new KvStateRegistry();
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
    EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
    long requestId = Integer.MAX_VALUE + 182828L;
    ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(channel.alloc(), requestId, new KvStateID(), new byte[0]);
    // Write the request and wait for the response
    channel.writeInbound(request);
    ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
    // skip frame length
    buf.skipBytes(4);
    // Verify the response
    assertEquals(KvStateRequestType.REQUEST_FAILURE, KvStateRequestSerializer.deserializeHeader(buf));
    KvStateRequestFailure response = KvStateRequestSerializer.deserializeKvStateRequestFailure(buf);
    assertEquals(requestId, response.getRequestId());
    assertTrue("Did not respond with expected failure cause", response.getCause() instanceof UnknownKvStateID);
    assertEquals(1, stats.getNumRequests());
    assertEquals(1, stats.getNumFailed());
}

Example 5 with KvStateRegistry

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

the class KvStateServerHandlerTest method testFailureOnGetSerializedValue.

/**
	 * Tests the failure response on a failure on the {@link InternalKvState#getSerializedValue(byte[])}
	 * call.
	 */
@Test
public void testFailureOnGetSerializedValue() throws Exception {
    KvStateRegistry registry = new KvStateRegistry();
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
    EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
    // Failing KvState
    InternalKvState<?> kvState = mock(InternalKvState.class);
    when(kvState.getSerializedValue(any(byte[].class))).thenThrow(new RuntimeException("Expected test Exception"));
    KvStateID kvStateId = registry.registerKvState(new JobID(), new JobVertexID(), new KeyGroupRange(0, 0), "vanilla", kvState);
    ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(channel.alloc(), 282872, kvStateId, new byte[0]);
    // Write the request and wait for the response
    channel.writeInbound(request);
    ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
    // skip frame length
    buf.skipBytes(4);
    // Verify the response
    assertEquals(KvStateRequestType.REQUEST_FAILURE, KvStateRequestSerializer.deserializeHeader(buf));
    KvStateRequestFailure response = KvStateRequestSerializer.deserializeKvStateRequestFailure(buf);
    assertTrue(response.getCause().getMessage().contains("Expected test Exception"));
    assertEquals(1, stats.getNumRequests());
    assertEquals(1, stats.getNumFailed());
}