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

Example 1 with KeyGroupRange

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

Example 2 with KeyGroupRange

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

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

the class QueryableStateClientTest method testIntegrationWithKvStateServer.

/**
	 * Tests queries against multiple servers.
	 *
	 * <p>The servers are populated with different keys and the client queries
	 * all available keys from all servers.
	 */
@Test
public void testIntegrationWithKvStateServer() throws Exception {
    // Config
    int numServers = 2;
    int numKeys = 1024;
    int numKeyGroups = 1;
    JobID jobId = new JobID();
    JobVertexID jobVertexId = new JobVertexID();
    KvStateServer[] servers = new KvStateServer[numServers];
    AtomicKvStateRequestStats[] serverStats = new AtomicKvStateRequestStats[numServers];
    QueryableStateClient client = null;
    KvStateClient networkClient = null;
    AtomicKvStateRequestStats networkClientStats = new AtomicKvStateRequestStats();
    MemoryStateBackend backend = new MemoryStateBackend();
    DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
    AbstractKeyedStateBackend<Integer> keyedStateBackend = backend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
    try {
        KvStateRegistry[] registries = new KvStateRegistry[numServers];
        KvStateID[] kvStateIds = new KvStateID[numServers];
        List<HeapValueState<Integer, VoidNamespace, Integer>> kvStates = new ArrayList<>();
        // Start the servers
        for (int i = 0; i < numServers; i++) {
            registries[i] = new KvStateRegistry();
            serverStats[i] = new AtomicKvStateRequestStats();
            servers[i] = new KvStateServer(InetAddress.getLocalHost(), 0, 1, 1, registries[i], serverStats[i]);
            servers[i].start();
            ValueStateDescriptor<Integer> descriptor = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
            RegisteredBackendStateMetaInfo<VoidNamespace, Integer> registeredBackendStateMetaInfo = new RegisteredBackendStateMetaInfo<>(descriptor.getType(), descriptor.getName(), VoidNamespaceSerializer.INSTANCE, IntSerializer.INSTANCE);
            // Register state
            HeapValueState<Integer, VoidNamespace, Integer> kvState = new HeapValueState<>(descriptor, new NestedMapsStateTable<Integer, VoidNamespace, Integer>(keyedStateBackend, registeredBackendStateMetaInfo), IntSerializer.INSTANCE, VoidNamespaceSerializer.INSTANCE);
            kvStates.add(kvState);
            kvStateIds[i] = registries[i].registerKvState(jobId, new JobVertexID(), new KeyGroupRange(i, i), "choco", kvState);
        }
        int[] expectedRequests = new int[numServers];
        for (int key = 0; key < numKeys; key++) {
            int targetKeyGroupIndex = MathUtils.murmurHash(key) % numServers;
            expectedRequests[targetKeyGroupIndex]++;
            HeapValueState<Integer, VoidNamespace, Integer> kvState = kvStates.get(targetKeyGroupIndex);
            keyedStateBackend.setCurrentKey(key);
            kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
            kvState.update(1337 + key);
        }
        // Location lookup service
        KvStateLocation location = new KvStateLocation(jobId, jobVertexId, numServers, "choco");
        for (int keyGroupIndex = 0; keyGroupIndex < numServers; keyGroupIndex++) {
            location.registerKvState(new KeyGroupRange(keyGroupIndex, keyGroupIndex), kvStateIds[keyGroupIndex], servers[keyGroupIndex].getAddress());
        }
        KvStateLocationLookupService lookupService = mock(KvStateLocationLookupService.class);
        when(lookupService.getKvStateLookupInfo(eq(jobId), eq("choco"))).thenReturn(Futures.successful(location));
        // The client
        networkClient = new KvStateClient(1, networkClientStats);
        client = new QueryableStateClient(lookupService, networkClient, testActorSystem.dispatcher());
        // Send all queries
        List<Future<byte[]>> futures = new ArrayList<>(numKeys);
        for (int key = 0; key < numKeys; key++) {
            byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(key, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
            futures.add(client.getKvState(jobId, "choco", key, serializedKeyAndNamespace));
        }
        // Verify results
        Future<Iterable<byte[]>> future = Futures.sequence(futures, testActorSystem.dispatcher());
        Iterable<byte[]> results = Await.result(future, timeout);
        int index = 0;
        for (byte[] buffer : results) {
            int deserializedValue = KvStateRequestSerializer.deserializeValue(buffer, IntSerializer.INSTANCE);
            assertEquals(1337 + index, deserializedValue);
            index++;
        }
        // Verify requests
        for (int i = 0; i < numServers; i++) {
            int numRetries = 10;
            for (int retry = 0; retry < numRetries; retry++) {
                try {
                    assertEquals("Unexpected number of requests", expectedRequests[i], serverStats[i].getNumRequests());
                    assertEquals("Unexpected success requests", expectedRequests[i], serverStats[i].getNumSuccessful());
                    assertEquals("Unexpected failed requests", 0, serverStats[i].getNumFailed());
                    break;
                } catch (Throwable t) {
                    // Retry
                    if (retry == numRetries - 1) {
                        throw t;
                    } else {
                        Thread.sleep(100);
                    }
                }
            }
        }
    } finally {
        if (client != null) {
            client.shutDown();
        }
        if (networkClient != null) {
            networkClient.shutDown();
        }
        for (KvStateServer server : servers) {
            if (server != null) {
                server.shutDown();
            }
        }
    }
}

Example 4 with KeyGroupRange

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

the class KvStateRequestSerializerTest method testListSerialization.

/**
	 * Tests list serialization utils.
	 */
@Test
public void testListSerialization() throws Exception {
    final long key = 0L;
    // objects for heap state list serialisation
    final HeapKeyedStateBackend<Long> longHeapKeyedStateBackend = new HeapKeyedStateBackend<>(mock(TaskKvStateRegistry.class), LongSerializer.INSTANCE, ClassLoader.getSystemClassLoader(), 1, new KeyGroupRange(0, 0), async, new ExecutionConfig());
    longHeapKeyedStateBackend.setCurrentKey(key);
    final InternalListState<VoidNamespace, Long> listState = longHeapKeyedStateBackend.createListState(VoidNamespaceSerializer.INSTANCE, new ListStateDescriptor<>("test", LongSerializer.INSTANCE));
    testListSerialization(key, listState);
}

Example 5 with KeyGroupRange

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

the class CopyOnWriteStateTableSnapshot method partitionEntriesByKeyGroup.

/**
	 * Partitions the snapshot data by key-group. The algorithm first builds a histogram for the distribution of keys
	 * into key-groups. Then, the histogram is accumulated to obtain the boundaries of each key-group in an array.
	 * Last, we use the accumulated counts as write position pointers for the key-group's bins when reordering the
	 * entries by key-group. This operation is lazily performed before the first writing of a key-group.
	 * <p>
	 * As a possible future optimization, we could perform the repartitioning in-place, using a scheme similar to the
	 * cuckoo cycles in cuckoo hashing. This can trade some performance for a smaller memory footprint.
	 */
@SuppressWarnings("unchecked")
private void partitionEntriesByKeyGroup() {
    // We only have to perform this step once before the first key-group is written
    if (null != keyGroupOffsets) {
        return;
    }
    final KeyGroupRange keyGroupRange = owningStateTable.keyContext.getKeyGroupRange();
    final int totalKeyGroups = owningStateTable.keyContext.getNumberOfKeyGroups();
    final int baseKgIdx = keyGroupRange.getStartKeyGroup();
    final int[] histogram = new int[keyGroupRange.getNumberOfKeyGroups() + 1];
    CopyOnWriteStateTable.StateTableEntry<K, N, S>[] unfold = new CopyOnWriteStateTable.StateTableEntry[stateTableSize];
    // 1) In this step we i) 'unfold' the linked list of entries to a flat array and ii) build a histogram for key-groups
    int unfoldIndex = 0;
    for (CopyOnWriteStateTable.StateTableEntry<K, N, S> entry : snapshotData) {
        while (null != entry) {
            int effectiveKgIdx = KeyGroupRangeAssignment.computeKeyGroupForKeyHash(entry.key.hashCode(), totalKeyGroups) - baseKgIdx + 1;
            ++histogram[effectiveKgIdx];
            unfold[unfoldIndex++] = entry;
            entry = entry.next;
        }
    }
    // 2) We accumulate the histogram bins to obtain key-group ranges in the final array
    for (int i = 1; i < histogram.length; ++i) {
        histogram[i] += histogram[i - 1];
    }
    // 3) We repartition the entries by key-group, using the histogram values as write indexes
    for (CopyOnWriteStateTable.StateTableEntry<K, N, S> t : unfold) {
        int effectiveKgIdx = KeyGroupRangeAssignment.computeKeyGroupForKeyHash(t.key.hashCode(), totalKeyGroups) - baseKgIdx;
        snapshotData[histogram[effectiveKgIdx]++] = t;
    }
    // 4) As byproduct, we also created the key-group offsets
    this.keyGroupOffsets = histogram;
}