Examples with DummyEnvironment org.apache.flink.runtime.operators.testutils.DummyEnvironment used on opensource projects

Example 11 with DummyEnvironment

use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project flink by apache.

the class RocksDBAsyncSnapshotTest method testCleanupOfSnapshotsInFailureCase.

/**
	 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
	 * procedure.
	 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
    long checkpointId = 1L;
    long timestamp = 42L;
    Environment env = new DummyEnvironment("test task", 1, 0);
    CheckpointStreamFactory.CheckpointStateOutputStream outputStream = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    CheckpointStreamFactory checkpointStreamFactory = mock(CheckpointStreamFactory.class);
    AbstractStateBackend stateBackend = mock(AbstractStateBackend.class);
    final IOException testException = new IOException("Test exception");
    doReturn(checkpointStreamFactory).when(stateBackend).createStreamFactory(any(JobID.class), anyString());
    doThrow(testException).when(outputStream).write(anyInt());
    doReturn(outputStream).when(checkpointStreamFactory).createCheckpointStateOutputStream(eq(checkpointId), eq(timestamp));
    RocksDBStateBackend backend = new RocksDBStateBackend(stateBackend);
    backend.setDbStoragePath("file:///tmp/foobar");
    AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(env, new JobID(), "test operator", VoidSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), null);
    // register a state so that the state backend has to checkpoint something
    keyedStateBackend.getPartitionedState("namespace", StringSerializer.INSTANCE, new ValueStateDescriptor<>("foobar", String.class));
    RunnableFuture<KeyGroupsStateHandle> snapshotFuture = keyedStateBackend.snapshot(checkpointId, timestamp, checkpointStreamFactory, CheckpointOptions.forFullCheckpoint());
    try {
        FutureUtil.runIfNotDoneAndGet(snapshotFuture);
        fail("Expected an exception to be thrown here.");
    } catch (ExecutionException e) {
        Assert.assertEquals(testException, e.getCause());
    }
    verify(outputStream).close();
}

Example 12 with DummyEnvironment

use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project flink by apache.

the class StreamingRuntimeContextTest method createMapPlainMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() 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<MapState<Integer, String>>() {

        @Override
        public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
            MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, 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(MapStateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    return operatorMock;
}

Example 13 with DummyEnvironment

use of org.apache.flink.runtime.operators.testutils.DummyEnvironment 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();
        }
    }
}

Example 14 with DummyEnvironment

use of org.apache.flink.runtime.operators.testutils.DummyEnvironment in project flink by apache.

the class KvStateServerHandlerTest method testQueryUnknownKey.

/**
	 * Tests the failure response with {@link UnknownKeyOrNamespace} as cause
	 * on queries for non-existing keys.
	 */
@Test
public void testQueryUnknownKey() 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);
    int numKeyGroups = 1;
    AbstractStateBackend abstractBackend = new MemoryStateBackend();
    DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
    dummyEnv.setKvStateRegistry(registry);
    KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
    final TestRegistryListener registryListener = new TestRegistryListener();
    registry.registerListener(registryListener);
    // Register state
    ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
    desc.setQueryable("vanilla");
    backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
    byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(1238283, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
    long requestId = Integer.MAX_VALUE + 22982L;
    assertTrue(registryListener.registrationName.equals("vanilla"));
    ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(channel.alloc(), requestId, registryListener.kvStateId, serializedKeyAndNamespace);
    // 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 UnknownKeyOrNamespace);
    assertEquals(1, stats.getNumRequests());
    assertEquals(1, stats.getNumFailed());
}

Example 15 with DummyEnvironment

use of org.apache.flink.runtime.operators.testutils.DummyEnvironment 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();
    KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, 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 = abstractBackend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, numKeyGroups, new KeyGroupRange(0, 0), registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
    final TestRegistryListener registryListener = new TestRegistryListener();
    registry.registerListener(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 = KvStateRequestSerializer.serializeKeyAndNamespace(key, IntSerializer.INSTANCE, VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE);
    long requestId = Integer.MAX_VALUE + 182828L;
    assertTrue(registryListener.registrationName.equals("vanilla"));
    ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(channel.alloc(), requestId, registryListener.kvStateId, serializedKeyAndNamespace);
    // Write the request and wait for the response
    channel.writeInbound(request);
    Object msg = readInboundBlocking(channel);
    assertTrue("Not ChunkedByteBuf", msg instanceof ChunkedByteBuf);
}