Examples with KvStateID org.apache.flink.queryablestate.KvStateID used on opensource projects

Example 6 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateRegistryTest method testKvStateEntry.

@Test
public void testKvStateEntry() throws InterruptedException {
    final int threads = 10;
    final CountDownLatch latch1 = new CountDownLatch(threads);
    final CountDownLatch latch2 = new CountDownLatch(1);
    final List<KvStateInfo<?, ?, ?>> infos = Collections.synchronizedList(new ArrayList<>());
    final JobID jobID = new JobID();
    final JobVertexID jobVertexId = new JobVertexID();
    final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
    final String registrationName = "foobar";
    final KvStateRegistry kvStateRegistry = new KvStateRegistry();
    final KvStateID stateID = kvStateRegistry.registerKvState(jobID, jobVertexId, keyGroupRange, registrationName, new DummyKvState(), getClass().getClassLoader());
    final AtomicReference<Throwable> exceptionHolder = new AtomicReference<>();
    for (int i = 0; i < threads; i++) {
        new Thread(() -> {
            final KvStateEntry<?, ?, ?> kvState = kvStateRegistry.getKvState(stateID);
            final KvStateInfo<?, ?, ?> stateInfo = kvState.getInfoForCurrentThread();
            infos.add(stateInfo);
            latch1.countDown();
            try {
                latch2.await();
            } catch (InterruptedException e) {
                // compare and set, so that we do not overwrite an exception
                // that was (potentially) already encountered.
                exceptionHolder.compareAndSet(null, e);
            }
        }).start();
    }
    latch1.await();
    final KvStateEntry<?, ?, ?> kvState = kvStateRegistry.getKvState(stateID);
    // verify that all the threads are done correctly.
    Assert.assertEquals(threads, infos.size());
    Assert.assertEquals(threads, kvState.getCacheSize());
    latch2.countDown();
    for (KvStateInfo<?, ?, ?> infoA : infos) {
        boolean instanceAlreadyFound = false;
        for (KvStateInfo<?, ?, ?> infoB : infos) {
            if (infoA == infoB) {
                if (instanceAlreadyFound) {
                    Assert.fail("More than one thread sharing the same serializer instance.");
                }
                instanceAlreadyFound = true;
            } else {
                Assert.assertEquals(infoA, infoB);
            }
        }
    }
    kvStateRegistry.unregisterKvState(jobID, jobVertexId, keyGroupRange, registrationName, stateID);
    Assert.assertEquals(0L, kvState.getCacheSize());
    Throwable t = exceptionHolder.get();
    if (t != null) {
        fail(t.getMessage());
    }
}

Example 7 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateRegistryTest method testLegacyCodePathPreference.

/**
 * Tests that {@link KvStateRegistryListener} registered under {@link
 * HighAvailabilityServices#DEFAULT_JOB_ID} will be used for all notifications.
 */
@Test
public void testLegacyCodePathPreference() {
    final KvStateRegistry kvStateRegistry = new KvStateRegistry();
    final ArrayDeque<JobID> stateRegistrationNotifications = new ArrayDeque<>(2);
    final ArrayDeque<JobID> stateDeregistrationNotifications = new ArrayDeque<>(2);
    final TestingKvStateRegistryListener testingListener = new TestingKvStateRegistryListener(stateRegistrationNotifications, stateDeregistrationNotifications);
    final ArrayDeque<JobID> anotherQueue = new ArrayDeque<>(2);
    final TestingKvStateRegistryListener anotherListener = new TestingKvStateRegistryListener(anotherQueue, anotherQueue);
    final JobID jobId = new JobID();
    kvStateRegistry.registerListener(HighAvailabilityServices.DEFAULT_JOB_ID, testingListener);
    kvStateRegistry.registerListener(jobId, anotherListener);
    final JobVertexID jobVertexId = new JobVertexID();
    final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
    final String registrationName = "registrationName";
    final KvStateID kvStateID = kvStateRegistry.registerKvState(jobId, jobVertexId, keyGroupRange, registrationName, new DummyKvState(), getClass().getClassLoader());
    assertThat(stateRegistrationNotifications.poll(), equalTo(jobId));
    // another listener should not have received any notifications
    assertThat(anotherQueue.isEmpty(), is(true));
    kvStateRegistry.unregisterKvState(jobId, jobVertexId, keyGroupRange, registrationName, kvStateID);
    assertThat(stateDeregistrationNotifications.poll(), equalTo(jobId));
    // another listener should not have received any notifications
    assertThat(anotherQueue.isEmpty(), is(true));
}

Example 8 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateLocationRegistryTest method testRegisterDuplicateName.

/**
 * Tests that registrations with duplicate names throw an Exception.
 */
@Test
public void testRegisterDuplicateName() throws Exception {
    ExecutionJobVertex[] vertices = new ExecutionJobVertex[] { createJobVertex(32), createJobVertex(13) };
    Map<JobVertexID, ExecutionJobVertex> vertexMap = createVertexMap(vertices);
    String registrationName = "duplicated-name";
    KvStateLocationRegistry registry = new KvStateLocationRegistry(new JobID(), vertexMap);
    // First operator registers
    registry.notifyKvStateRegistered(vertices[0].getJobVertexId(), new KeyGroupRange(0, 0), registrationName, new KvStateID(), new InetSocketAddress(InetAddress.getLocalHost(), 12328));
    try {
        // Second operator registers same name
        registry.notifyKvStateRegistered(vertices[1].getJobVertexId(), new KeyGroupRange(0, 0), registrationName, new KvStateID(), new InetSocketAddress(InetAddress.getLocalHost(), 12032));
        fail("Did not throw expected Exception after duplicated name");
    } catch (IllegalStateException ignored) {
    // Expected
    }
}

Example 9 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class ClientTest method testSimpleRequests.

/**
 * Tests simple queries, of which half succeed and half fail.
 */
@Test
public void testSimpleRequests() throws Exception {
    AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();
    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);
        // 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 ChannelDataCollectingHandler(channel, received));
        InetSocketAddress serverAddress = getKvStateServerAddress(serverChannel);
        long numQueries = 1024L;
        List<CompletableFuture<KvStateResponse>> futures = new ArrayList<>();
        for (long i = 0L; i < numQueries; i++) {
            KvStateInternalRequest request = new KvStateInternalRequest(new KvStateID(), new byte[0]);
            futures.add(client.sendRequest(serverAddress, request));
        }
        // Respond to messages
        Exception testException = new RuntimeException("Expected test Exception");
        for (long i = 0L; i < numQueries; i++) {
            ByteBuf buf = received.take();
            assertNotNull("Receive timed out", buf);
            Channel ch = channel.get();
            assertNotNull("Channel not active", ch);
            assertEquals(MessageType.REQUEST, MessageSerializer.deserializeHeader(buf));
            long requestId = MessageSerializer.getRequestId(buf);
            KvStateInternalRequest deserRequest = serializer.deserializeRequest(buf);
            buf.release();
            if (i % 2L == 0L) {
                ByteBuf response = MessageSerializer.serializeResponse(serverChannel.alloc(), requestId, new KvStateResponse(expected));
                ch.writeAndFlush(response);
            } else {
                ByteBuf response = MessageSerializer.serializeRequestFailure(serverChannel.alloc(), requestId, testException);
                ch.writeAndFlush(response);
            }
        }
        for (long i = 0L; i < numQueries; i++) {
            if (i % 2L == 0L) {
                KvStateResponse serializedResult = futures.get((int) i).get();
                assertArrayEquals(expected, serializedResult.getContent());
            } else {
                try {
                    futures.get((int) i).get();
                    fail("Did not throw expected Exception");
                } catch (ExecutionException e) {
                    if (!(e.getCause() instanceof RuntimeException)) {
                        fail("Did not throw expected Exception");
                    }
                // else expected
                }
            }
        }
        assertEquals(numQueries, stats.getNumRequests());
        long expectedRequests = numQueries / 2L;
        // Counts can take some time to propagate
        while (stats.getNumSuccessful() != expectedRequests || stats.getNumFailed() != expectedRequests) {
            Thread.sleep(100L);
        }
        assertEquals(expectedRequests, stats.getNumSuccessful());
        assertEquals(expectedRequests, stats.getNumFailed());
    } finally {
        if (client != null) {
            Exception exc = null;
            try {
                // todo here we were seeing this problem:
                // https://github.com/netty/netty/issues/4357 if we do a get().
                // this is why we now simply wait a bit so that everything is
                // shut down and then we check
                client.shutdown().get();
            } catch (Exception e) {
                exc = e;
                LOG.error("An exception occurred while shutting down netty.", e);
            }
            Assert.assertTrue(ExceptionUtils.stringifyException(exc), client.isEventGroupShutdown());
        }
        if (serverChannel != null) {
            serverChannel.close();
        }
        assertEquals("Channel leak", 0L, stats.getNumConnections());
    }
}

Example 10 with KvStateID

use of org.apache.flink.queryablestate.KvStateID in project flink by apache.

the class KvStateServerHandlerTest method testFailureOnGetSerializedValue.

/**
 * Tests the failure response on a failure on the {@link
 * InternalKvState#getSerializedValue(byte[], TypeSerializer, TypeSerializer, TypeSerializer)}
 * call.
 */
@Test
public void testFailureOnGetSerializedValue() 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);
    // Failing KvState
    InternalKvState<Integer, VoidNamespace, Long> kvState = new InternalKvState<Integer, VoidNamespace, Long>() {

        @Override
        public TypeSerializer<Integer> getKeySerializer() {
            return IntSerializer.INSTANCE;
        }

        @Override
        public TypeSerializer<VoidNamespace> getNamespaceSerializer() {
            return VoidNamespaceSerializer.INSTANCE;
        }

        @Override
        public TypeSerializer<Long> getValueSerializer() {
            return LongSerializer.INSTANCE;
        }

        @Override
        public void setCurrentNamespace(VoidNamespace namespace) {
        // do nothing
        }

        @Override
        public byte[] getSerializedValue(final byte[] serializedKeyAndNamespace, final TypeSerializer<Integer> safeKeySerializer, final TypeSerializer<VoidNamespace> safeNamespaceSerializer, final TypeSerializer<Long> safeValueSerializer) throws Exception {
            throw new RuntimeException("Expected test Exception");
        }

        @Override
        public StateIncrementalVisitor<Integer, VoidNamespace, Long> getStateIncrementalVisitor(int recommendedMaxNumberOfReturnedRecords) {
            throw new UnsupportedOperationException();
        }

        @Override
        public void clear() {
        }
    };
    KvStateID kvStateId = registry.registerKvState(new JobID(), new JobVertexID(), new KeyGroupRange(0, 0), "vanilla", kvState, getClass().getClassLoader());
    KvStateInternalRequest request = new KvStateInternalRequest(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("Expected test Exception"));
    assertEquals(1L, stats.getNumRequests());
    assertEquals(1L, stats.getNumFailed());
}