Search in sources :

Example 1 with InternalKvState

use of org.apache.flink.runtime.state.internal.InternalKvState in project flink by apache.

the class StateBackendTestBase method testValueStateRace.

/**
	 * Tests {@link ValueState#value()} and {@link InternalKvState#getSerializedValue(byte[])}
	 * accessing the state concurrently. They should not get in the way of each
	 * other.
	 */
@Test
@SuppressWarnings("unchecked")
public void testValueStateRace() throws Exception {
    final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
    final Integer namespace = 1;
    final ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class);
    kvId.initializeSerializerUnlessSet(new ExecutionConfig());
    final TypeSerializer<Integer> keySerializer = IntSerializer.INSTANCE;
    final TypeSerializer<Integer> namespaceSerializer = IntSerializer.INSTANCE;
    final TypeSerializer<String> valueSerializer = kvId.getSerializer();
    final ValueState<String> state = backend.getPartitionedState(namespace, IntSerializer.INSTANCE, kvId);
    @SuppressWarnings("unchecked") final InternalKvState<Integer> kvState = (InternalKvState<Integer>) state;
    /**
		 * 1) Test that ValueState#value() before and after
		 * KvState#getSerializedValue(byte[]) return the same value.
		 */
    // set some key and namespace
    final int key1 = 1;
    backend.setCurrentKey(key1);
    kvState.setCurrentNamespace(2);
    state.update("2");
    assertEquals("2", state.value());
    // query another key and namespace
    assertNull(getSerializedValue(kvState, 3, keySerializer, namespace, IntSerializer.INSTANCE, valueSerializer));
    // the state should not have changed!
    assertEquals("2", state.value());
    // re-set values
    kvState.setCurrentNamespace(namespace);
    /**
		 * 2) Test two threads concurrently using ValueState#value() and
		 * KvState#getSerializedValue(byte[]).
		 */
    // some modifications to the state
    final int key2 = 10;
    backend.setCurrentKey(key2);
    assertNull(state.value());
    assertNull(getSerializedValue(kvState, key2, keySerializer, namespace, namespaceSerializer, valueSerializer));
    state.update("1");
    final CheckedThread getter = new CheckedThread("State getter") {

        @Override
        public void go() throws Exception {
            while (!isInterrupted()) {
                assertEquals("1", state.value());
            }
        }
    };
    final CheckedThread serializedGetter = new CheckedThread("Serialized state getter") {

        @Override
        public void go() throws Exception {
            while (!isInterrupted() && getter.isAlive()) {
                final String serializedValue = getSerializedValue(kvState, key2, keySerializer, namespace, namespaceSerializer, valueSerializer);
                assertEquals("1", serializedValue);
            }
        }
    };
    getter.start();
    serializedGetter.start();
    // run both threads for max 100ms
    Timer t = new Timer("stopper");
    t.schedule(new TimerTask() {

        @Override
        public void run() {
            getter.interrupt();
            serializedGetter.interrupt();
            this.cancel();
        }
    }, 100);
    // wait for both threads to finish
    try {
        // serializedGetter will finish if its assertion fails or if
        // getter is not alive any more
        serializedGetter.sync();
        // if serializedGetter crashed, getter will not know -> interrupt just in case
        getter.interrupt();
        getter.sync();
        // if not executed yet
        t.cancel();
    } finally {
        // clean up
        backend.dispose();
    }
}
Also used : ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) CheckedThread(org.apache.flink.core.testutils.CheckedThread) ValueStateDescriptor(org.apache.flink.api.common.state.ValueStateDescriptor) InternalKvState(org.apache.flink.runtime.state.internal.InternalKvState) Timer(java.util.Timer) TimerTask(java.util.TimerTask) Test(org.junit.Test)

Example 2 with InternalKvState

use of org.apache.flink.runtime.state.internal.InternalKvState in project flink by apache.

the class StateBackendTestBase method testListState.

@Test
@SuppressWarnings("unchecked,rawtypes")
public void testListState() {
    try {
        CheckpointStreamFactory streamFactory = createStreamFactory();
        AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
        ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", String.class);
        kvId.initializeSerializerUnlessSet(new ExecutionConfig());
        TypeSerializer<Integer> keySerializer = IntSerializer.INSTANCE;
        TypeSerializer<VoidNamespace> namespaceSerializer = VoidNamespaceSerializer.INSTANCE;
        TypeSerializer<String> valueSerializer = kvId.getElementSerializer();
        ListState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        @SuppressWarnings("unchecked") InternalKvState<VoidNamespace> kvState = (InternalKvState<VoidNamespace>) state;
        Joiner joiner = Joiner.on(",");
        // some modifications to the state
        backend.setCurrentKey(1);
        assertEquals(null, state.get());
        assertEquals(null, getSerializedList(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        state.add("1");
        backend.setCurrentKey(2);
        assertEquals(null, state.get());
        assertEquals(null, getSerializedList(kvState, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        state.add("2");
        backend.setCurrentKey(1);
        assertEquals("1", joiner.join(state.get()));
        assertEquals("1", joiner.join(getSerializedList(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        // draw a snapshot
        KeyGroupsStateHandle snapshot1 = FutureUtil.runIfNotDoneAndGet(backend.snapshot(682375462378L, 2, streamFactory, CheckpointOptions.forFullCheckpoint()));
        // make some more modifications
        backend.setCurrentKey(1);
        state.add("u1");
        backend.setCurrentKey(2);
        state.add("u2");
        backend.setCurrentKey(3);
        state.add("u3");
        // draw another snapshot
        KeyGroupsStateHandle snapshot2 = FutureUtil.runIfNotDoneAndGet(backend.snapshot(682375462379L, 4, streamFactory, CheckpointOptions.forFullCheckpoint()));
        // validate the original state
        backend.setCurrentKey(1);
        assertEquals("1,u1", joiner.join(state.get()));
        assertEquals("1,u1", joiner.join(getSerializedList(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.setCurrentKey(2);
        assertEquals("2,u2", joiner.join(state.get()));
        assertEquals("2,u2", joiner.join(getSerializedList(kvState, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.setCurrentKey(3);
        assertEquals("u3", joiner.join(state.get()));
        assertEquals("u3", joiner.join(getSerializedList(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.dispose();
        // restore the first snapshot and validate it
        backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
        snapshot1.discardState();
        ListState<String> restored1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        @SuppressWarnings("unchecked") InternalKvState<VoidNamespace> restoredKvState1 = (InternalKvState<VoidNamespace>) restored1;
        backend.setCurrentKey(1);
        assertEquals("1", joiner.join(restored1.get()));
        assertEquals("1", joiner.join(getSerializedList(restoredKvState1, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.setCurrentKey(2);
        assertEquals("2", joiner.join(restored1.get()));
        assertEquals("2", joiner.join(getSerializedList(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.dispose();
        // restore the second snapshot and validate it
        backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot2);
        snapshot2.discardState();
        ListState<String> restored2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        @SuppressWarnings("unchecked") InternalKvState<VoidNamespace> restoredKvState2 = (InternalKvState<VoidNamespace>) restored2;
        backend.setCurrentKey(1);
        assertEquals("1,u1", joiner.join(restored2.get()));
        assertEquals("1,u1", joiner.join(getSerializedList(restoredKvState2, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.setCurrentKey(2);
        assertEquals("2,u2", joiner.join(restored2.get()));
        assertEquals("2,u2", joiner.join(getSerializedList(restoredKvState2, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.setCurrentKey(3);
        assertEquals("u3", joiner.join(restored2.get()));
        assertEquals("u3", joiner.join(getSerializedList(restoredKvState2, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
        backend.dispose();
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}
Also used : Joiner(com.google.common.base.Joiner) BlockerCheckpointStreamFactory(org.apache.flink.runtime.util.BlockerCheckpointStreamFactory) ListStateDescriptor(org.apache.flink.api.common.state.ListStateDescriptor) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException) InternalKvState(org.apache.flink.runtime.state.internal.InternalKvState) Test(org.junit.Test)

Example 3 with InternalKvState

use of org.apache.flink.runtime.state.internal.InternalKvState in project flink by apache.

the class StateBackendTestBase method testFoldingState.

@Test
@SuppressWarnings("unchecked,rawtypes")
public void testFoldingState() {
    try {
        CheckpointStreamFactory streamFactory = createStreamFactory();
        AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
        FoldingStateDescriptor<Integer, String> kvId = new FoldingStateDescriptor<>("id", "Fold-Initial:", new AppendingFold(), String.class);
        kvId.initializeSerializerUnlessSet(new ExecutionConfig());
        TypeSerializer<Integer> keySerializer = IntSerializer.INSTANCE;
        TypeSerializer<VoidNamespace> namespaceSerializer = VoidNamespaceSerializer.INSTANCE;
        TypeSerializer<String> valueSerializer = kvId.getSerializer();
        FoldingState<Integer, String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        @SuppressWarnings("unchecked") InternalKvState<VoidNamespace> kvState = (InternalKvState<VoidNamespace>) state;
        // some modifications to the state
        backend.setCurrentKey(1);
        assertEquals(null, state.get());
        assertEquals(null, getSerializedValue(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        state.add(1);
        backend.setCurrentKey(2);
        assertEquals(null, state.get());
        assertEquals(null, getSerializedValue(kvState, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        state.add(2);
        backend.setCurrentKey(1);
        assertEquals("Fold-Initial:,1", state.get());
        assertEquals("Fold-Initial:,1", getSerializedValue(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        // draw a snapshot
        KeyGroupsStateHandle snapshot1 = FutureUtil.runIfNotDoneAndGet(backend.snapshot(682375462378L, 2, streamFactory, CheckpointOptions.forFullCheckpoint()));
        // make some more modifications
        backend.setCurrentKey(1);
        state.clear();
        state.add(101);
        backend.setCurrentKey(2);
        state.add(102);
        backend.setCurrentKey(3);
        state.add(103);
        // draw another snapshot
        KeyGroupsStateHandle snapshot2 = FutureUtil.runIfNotDoneAndGet(backend.snapshot(682375462379L, 4, streamFactory, CheckpointOptions.forFullCheckpoint()));
        // validate the original state
        backend.setCurrentKey(1);
        assertEquals("Fold-Initial:,101", state.get());
        assertEquals("Fold-Initial:,101", getSerializedValue(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.setCurrentKey(2);
        assertEquals("Fold-Initial:,2,102", state.get());
        assertEquals("Fold-Initial:,2,102", getSerializedValue(kvState, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.setCurrentKey(3);
        assertEquals("Fold-Initial:,103", state.get());
        assertEquals("Fold-Initial:,103", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.dispose();
        // restore the first snapshot and validate it
        backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
        snapshot1.discardState();
        FoldingState<Integer, String> restored1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        @SuppressWarnings("unchecked") InternalKvState<VoidNamespace> restoredKvState1 = (InternalKvState<VoidNamespace>) restored1;
        backend.setCurrentKey(1);
        assertEquals("Fold-Initial:,1", restored1.get());
        assertEquals("Fold-Initial:,1", getSerializedValue(restoredKvState1, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.setCurrentKey(2);
        assertEquals("Fold-Initial:,2", restored1.get());
        assertEquals("Fold-Initial:,2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.dispose();
        // restore the second snapshot and validate it
        backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot2);
        snapshot1.discardState();
        @SuppressWarnings("unchecked") FoldingState<Integer, String> restored2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
        @SuppressWarnings("unchecked") InternalKvState<VoidNamespace> restoredKvState2 = (InternalKvState<VoidNamespace>) restored2;
        backend.setCurrentKey(1);
        assertEquals("Fold-Initial:,101", restored2.get());
        assertEquals("Fold-Initial:,101", getSerializedValue(restoredKvState2, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.setCurrentKey(2);
        assertEquals("Fold-Initial:,2,102", restored2.get());
        assertEquals("Fold-Initial:,2,102", getSerializedValue(restoredKvState2, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.setCurrentKey(3);
        assertEquals("Fold-Initial:,103", restored2.get());
        assertEquals("Fold-Initial:,103", getSerializedValue(restoredKvState2, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
        backend.dispose();
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}
Also used : BlockerCheckpointStreamFactory(org.apache.flink.runtime.util.BlockerCheckpointStreamFactory) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) FoldingStateDescriptor(org.apache.flink.api.common.state.FoldingStateDescriptor) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException) InternalKvState(org.apache.flink.runtime.state.internal.InternalKvState) Test(org.junit.Test)

Example 4 with InternalKvState

use of org.apache.flink.runtime.state.internal.InternalKvState in project flink by apache.

the class AbstractKeyedStateBackend method getOrCreateKeyedState.

/**
	 * @see KeyedStateBackend
	 */
@Override
public <N, S extends State, V> S getOrCreateKeyedState(final TypeSerializer<N> namespaceSerializer, StateDescriptor<S, V> stateDescriptor) throws Exception {
    checkNotNull(namespaceSerializer, "Namespace serializer");
    if (keySerializer == null) {
        throw new UnsupportedOperationException("State key serializer has not been configured in the config. " + "This operation cannot use partitioned state.");
    }
    if (!stateDescriptor.isSerializerInitialized()) {
        throw new IllegalStateException("The serializer of the descriptor has not been initialized!");
    }
    InternalKvState<?> existing = keyValueStatesByName.get(stateDescriptor.getName());
    if (existing != null) {
        @SuppressWarnings("unchecked") S typedState = (S) existing;
        return typedState;
    }
    // create a new blank key/value state
    S state = stateDescriptor.bind(new StateBinder() {

        @Override
        public <T> ValueState<T> createValueState(ValueStateDescriptor<T> stateDesc) throws Exception {
            return AbstractKeyedStateBackend.this.createValueState(namespaceSerializer, stateDesc);
        }

        @Override
        public <T> ListState<T> createListState(ListStateDescriptor<T> stateDesc) throws Exception {
            return AbstractKeyedStateBackend.this.createListState(namespaceSerializer, stateDesc);
        }

        @Override
        public <T> ReducingState<T> createReducingState(ReducingStateDescriptor<T> stateDesc) throws Exception {
            return AbstractKeyedStateBackend.this.createReducingState(namespaceSerializer, stateDesc);
        }

        @Override
        public <T, ACC, R> AggregatingState<T, R> createAggregatingState(AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
            return AbstractKeyedStateBackend.this.createAggregatingState(namespaceSerializer, stateDesc);
        }

        @Override
        public <T, ACC> FoldingState<T, ACC> createFoldingState(FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
            return AbstractKeyedStateBackend.this.createFoldingState(namespaceSerializer, stateDesc);
        }

        @Override
        public <UK, UV> MapState<UK, UV> createMapState(MapStateDescriptor<UK, UV> stateDesc) throws Exception {
            return AbstractKeyedStateBackend.this.createMapState(namespaceSerializer, stateDesc);
        }
    });
    @SuppressWarnings("unchecked") InternalKvState<N> kvState = (InternalKvState<N>) state;
    keyValueStatesByName.put(stateDescriptor.getName(), kvState);
    // Publish queryable state
    if (stateDescriptor.isQueryable()) {
        if (kvStateRegistry == null) {
            throw new IllegalStateException("State backend has not been initialized for job.");
        }
        String name = stateDescriptor.getQueryableStateName();
        kvStateRegistry.registerKvState(keyGroupRange, name, kvState);
    }
    return state;
}
Also used : ListState(org.apache.flink.api.common.state.ListState) InternalListState(org.apache.flink.runtime.state.internal.InternalListState) StateBinder(org.apache.flink.api.common.state.StateBinder) InternalMapState(org.apache.flink.runtime.state.internal.InternalMapState) MapState(org.apache.flink.api.common.state.MapState) IOException(java.io.IOException) InternalReducingState(org.apache.flink.runtime.state.internal.InternalReducingState) ReducingState(org.apache.flink.api.common.state.ReducingState) InternalValueState(org.apache.flink.runtime.state.internal.InternalValueState) ValueState(org.apache.flink.api.common.state.ValueState) InternalKvState(org.apache.flink.runtime.state.internal.InternalKvState) FoldingState(org.apache.flink.api.common.state.FoldingState) InternalFoldingState(org.apache.flink.runtime.state.internal.InternalFoldingState) InternalAggregatingState(org.apache.flink.runtime.state.internal.InternalAggregatingState) AggregatingState(org.apache.flink.api.common.state.AggregatingState)

Example 5 with InternalKvState

use of org.apache.flink.runtime.state.internal.InternalKvState in project flink by apache.

the class KvStateClientTest method testClientServerIntegration.

/**
	 * Tests multiple clients querying multiple servers until 100k queries have
	 * been processed. At this point, the client is shut down and its verified
	 * that all ongoing requests are failed.
	 */
@Test
public void testClientServerIntegration() throws Exception {
    // Config
    final int numServers = 2;
    final int numServerEventLoopThreads = 2;
    final int numServerQueryThreads = 2;
    final int numClientEventLoopThreads = 4;
    final int numClientsTasks = 8;
    final int batchSize = 16;
    final int numKeyGroups = 1;
    AbstractStateBackend abstractBackend = new MemoryStateBackend();
    KvStateRegistry dummyRegistry = new KvStateRegistry();
    DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
    dummyEnv.setKvStateRegistry(dummyRegistry);
    AbstractKeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), dummyRegistry.createTaskRegistry(new JobID(), new JobVertexID()));
    final FiniteDuration timeout = new FiniteDuration(10, TimeUnit.SECONDS);
    AtomicKvStateRequestStats clientStats = new AtomicKvStateRequestStats();
    KvStateClient client = null;
    ExecutorService clientTaskExecutor = null;
    final KvStateServer[] server = new KvStateServer[numServers];
    try {
        client = new KvStateClient(numClientEventLoopThreads, clientStats);
        clientTaskExecutor = Executors.newFixedThreadPool(numClientsTasks);
        // Create state
        ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
        desc.setQueryable("any");
        // Create servers
        KvStateRegistry[] registry = new KvStateRegistry[numServers];
        AtomicKvStateRequestStats[] serverStats = new AtomicKvStateRequestStats[numServers];
        final KvStateID[] ids = new KvStateID[numServers];
        for (int i = 0; i < numServers; i++) {
            registry[i] = new KvStateRegistry();
            serverStats[i] = new AtomicKvStateRequestStats();
            server[i] = new KvStateServer(InetAddress.getLocalHost(), 0, numServerEventLoopThreads, numServerQueryThreads, registry[i], serverStats[i]);
            server[i].start();
            backend.setCurrentKey(1010 + i);
            // Value per server
            ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
            state.update(201 + i);
            // we know it must be a KvStat but this is not exposed to the user via State
            InternalKvState<?> kvState = (InternalKvState<?>) state;
            // Register KvState (one state instance for all server)
            ids[i] = registry[i].registerKvState(new JobID(), new JobVertexID(), new KeyGroupRange(0, 0), "any", kvState);
        }
        final KvStateClient finalClient = client;
        Callable<Void> queryTask = new Callable<Void>() {

            @Override
            public Void call() throws Exception {
                while (true) {
                    if (Thread.interrupted()) {
                        throw new InterruptedException();
                    }
                    // Random server permutation
                    List<Integer> random = new ArrayList<>();
                    for (int j = 0; j < batchSize; j++) {
                        random.add(j);
                    }
                    Collections.shuffle(random);
                    // Dispatch queries
                    List<Future<byte[]>> futures = new ArrayList<>(batchSize);
                    for (int j = 0; j < batchSize; j++) {
                        int targetServer = random.get(j) % numServers;
                        byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(1010 + targetServer, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
                        futures.add(finalClient.getKvState(server[targetServer].getAddress(), ids[targetServer], serializedKeyAndNamespace));
                    }
                    // Verify results
                    for (int j = 0; j < batchSize; j++) {
                        int targetServer = random.get(j) % numServers;
                        Future<byte[]> future = futures.get(j);
                        byte[] buf = Await.result(future, timeout);
                        int value = KvStateRequestSerializer.deserializeValue(buf, IntSerializer.INSTANCE);
                        assertEquals(201 + targetServer, value);
                    }
                }
            }
        };
        // Submit tasks
        List<java.util.concurrent.Future<Void>> taskFutures = new ArrayList<>();
        for (int i = 0; i < numClientsTasks; i++) {
            taskFutures.add(clientTaskExecutor.submit(queryTask));
        }
        long numRequests;
        while ((numRequests = clientStats.getNumRequests()) < 100_000) {
            Thread.sleep(100);
            LOG.info("Number of requests {}/100_000", numRequests);
        }
        // Shut down
        client.shutDown();
        for (java.util.concurrent.Future<Void> future : taskFutures) {
            try {
                future.get();
                fail("Did not throw expected Exception after shut down");
            } catch (ExecutionException t) {
                if (t.getCause() instanceof ClosedChannelException || t.getCause() instanceof IllegalStateException) {
                // Expected
                } else {
                    t.printStackTrace();
                    fail("Failed with unexpected Exception type: " + t.getClass().getName());
                }
            }
        }
        assertEquals("Connection leak (client)", 0, clientStats.getNumConnections());
        for (int i = 0; i < numServers; i++) {
            boolean success = false;
            int numRetries = 0;
            while (!success) {
                try {
                    assertEquals("Connection leak (server)", 0, serverStats[i].getNumConnections());
                    success = true;
                } catch (Throwable t) {
                    if (numRetries < 10) {
                        LOG.info("Retrying connection leak check (server)");
                        Thread.sleep((numRetries + 1) * 50);
                        numRetries++;
                    } else {
                        throw t;
                    }
                }
            }
        }
    } finally {
        if (client != null) {
            client.shutDown();
        }
        for (int i = 0; i < numServers; i++) {
            if (server[i] != null) {
                server[i].shutDown();
            }
        }
        if (clientTaskExecutor != null) {
            clientTaskExecutor.shutdown();
        }
    }
}
Also used : KvStateRegistry(org.apache.flink.runtime.query.KvStateRegistry) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) JobVertexID(org.apache.flink.runtime.jobgraph.JobVertexID) KeyGroupRange(org.apache.flink.runtime.state.KeyGroupRange) ArrayList(java.util.ArrayList) Callable(java.util.concurrent.Callable) ValueStateDescriptor(org.apache.flink.api.common.state.ValueStateDescriptor) KvStateID(org.apache.flink.runtime.query.KvStateID) ExecutionException(java.util.concurrent.ExecutionException) AbstractStateBackend(org.apache.flink.runtime.state.AbstractStateBackend) ClosedChannelException(java.nio.channels.ClosedChannelException) FiniteDuration(scala.concurrent.duration.FiniteDuration) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) InternalKvState(org.apache.flink.runtime.state.internal.InternalKvState) ExecutorService(java.util.concurrent.ExecutorService) Future(scala.concurrent.Future) JobID(org.apache.flink.api.common.JobID) Test(org.junit.Test)

Aggregations

InternalKvState (org.apache.flink.runtime.state.internal.InternalKvState)9 ExecutionConfig (org.apache.flink.api.common.ExecutionConfig)7 Test (org.junit.Test)7 IOException (java.io.IOException)5 ExecutionException (java.util.concurrent.ExecutionException)5 BlockerCheckpointStreamFactory (org.apache.flink.runtime.util.BlockerCheckpointStreamFactory)5 ValueStateDescriptor (org.apache.flink.api.common.state.ValueStateDescriptor)4 ArrayList (java.util.ArrayList)2 FoldingState (org.apache.flink.api.common.state.FoldingState)2 FoldingStateDescriptor (org.apache.flink.api.common.state.FoldingStateDescriptor)2 ListState (org.apache.flink.api.common.state.ListState)2 ListStateDescriptor (org.apache.flink.api.common.state.ListStateDescriptor)2 MapState (org.apache.flink.api.common.state.MapState)2 MapStateDescriptor (org.apache.flink.api.common.state.MapStateDescriptor)2 ReducingState (org.apache.flink.api.common.state.ReducingState)2 ReducingStateDescriptor (org.apache.flink.api.common.state.ReducingStateDescriptor)2 ValueState (org.apache.flink.api.common.state.ValueState)2 DummyEnvironment (org.apache.flink.runtime.operators.testutils.DummyEnvironment)2 InternalValueState (org.apache.flink.runtime.state.internal.InternalValueState)2 Joiner (com.google.common.base.Joiner)1