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

Example 1 with MockInputSplitProvider

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

the class OneInputStreamTaskTest method testSnapshottingAndRestoring.

/**
	 * Tests that the stream operator can snapshot and restore the operator state of chained
	 * operators
	 */
@Test
public void testSnapshottingAndRestoring() throws Exception {
    final Deadline deadline = new FiniteDuration(2, TimeUnit.MINUTES).fromNow();
    final OneInputStreamTask<String, String> streamTask = new OneInputStreamTask<String, String>();
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<String, String>(streamTask, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setupOutputForSingletonOperatorChain();
    IdentityKeySelector<String> keySelector = new IdentityKeySelector<>();
    testHarness.configureForKeyedStream(keySelector, BasicTypeInfo.STRING_TYPE_INFO);
    long checkpointId = 1L;
    long checkpointTimestamp = 1L;
    long recoveryTimestamp = 3L;
    long seed = 2L;
    int numberChainedTasks = 11;
    StreamConfig streamConfig = testHarness.getStreamConfig();
    configureChainedTestingStreamOperator(streamConfig, numberChainedTasks, seed, recoveryTimestamp);
    AcknowledgeStreamMockEnvironment env = new AcknowledgeStreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.executionConfig, testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize);
    // reset number of restore calls
    TestingStreamOperator.numberRestoreCalls = 0;
    testHarness.invoke(env);
    testHarness.waitForTaskRunning(deadline.timeLeft().toMillis());
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, checkpointTimestamp);
    while (!streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint())) ;
    // since no state was set, there shouldn't be restore calls
    assertEquals(0, TestingStreamOperator.numberRestoreCalls);
    env.getCheckpointLatch().await();
    assertEquals(checkpointId, env.getCheckpointId());
    testHarness.endInput();
    testHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
    final OneInputStreamTask<String, String> restoredTask = new OneInputStreamTask<String, String>();
    restoredTask.setInitialState(new TaskStateHandles(env.getCheckpointStateHandles()));
    final OneInputStreamTaskTestHarness<String, String> restoredTaskHarness = new OneInputStreamTaskTestHarness<String, String>(restoredTask, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    restoredTaskHarness.configureForKeyedStream(keySelector, BasicTypeInfo.STRING_TYPE_INFO);
    StreamConfig restoredTaskStreamConfig = restoredTaskHarness.getStreamConfig();
    configureChainedTestingStreamOperator(restoredTaskStreamConfig, numberChainedTasks, seed, recoveryTimestamp);
    TestingStreamOperator.numberRestoreCalls = 0;
    restoredTaskHarness.invoke();
    restoredTaskHarness.endInput();
    restoredTaskHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
    // restore of every chained operator should have been called
    assertEquals(numberChainedTasks, TestingStreamOperator.numberRestoreCalls);
    TestingStreamOperator.numberRestoreCalls = 0;
}

Example 2 with MockInputSplitProvider

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

the class AsyncWaitOperatorTest method testStateSnapshotAndRestore.

@Test
public void testStateSnapshotAndRestore() throws Exception {
    final OneInputStreamTask<Integer, Integer> task = new OneInputStreamTask<>();
    final OneInputStreamTaskTestHarness<Integer, Integer> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, BasicTypeInfo.INT_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO);
    testHarness.setupOutputForSingletonOperatorChain();
    AsyncWaitOperator<Integer, Integer> operator = new AsyncWaitOperator<>(new LazyAsyncFunction(), TIMEOUT, 3, AsyncDataStream.OutputMode.ORDERED);
    final StreamConfig streamConfig = testHarness.getStreamConfig();
    streamConfig.setStreamOperator(operator);
    final AcknowledgeStreamMockEnvironment env = new AcknowledgeStreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.getExecutionConfig(), testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize);
    testHarness.invoke(env);
    testHarness.waitForTaskRunning();
    final long initialTime = 0L;
    testHarness.processElement(new StreamRecord<>(1, initialTime + 1));
    testHarness.processElement(new StreamRecord<>(2, initialTime + 2));
    testHarness.processElement(new StreamRecord<>(3, initialTime + 3));
    testHarness.processElement(new StreamRecord<>(4, initialTime + 4));
    testHarness.waitForInputProcessing();
    final long checkpointId = 1L;
    final long checkpointTimestamp = 1L;
    final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, checkpointTimestamp);
    task.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint());
    env.getCheckpointLatch().await();
    assertEquals(checkpointId, env.getCheckpointId());
    LazyAsyncFunction.countDown();
    testHarness.endInput();
    testHarness.waitForTaskCompletion();
    // set the operator state from previous attempt into the restored one
    final OneInputStreamTask<Integer, Integer> restoredTask = new OneInputStreamTask<>();
    restoredTask.setInitialState(new TaskStateHandles(env.getCheckpointStateHandles()));
    final OneInputStreamTaskTestHarness<Integer, Integer> restoredTaskHarness = new OneInputStreamTaskTestHarness<>(restoredTask, BasicTypeInfo.INT_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO);
    restoredTaskHarness.setupOutputForSingletonOperatorChain();
    AsyncWaitOperator<Integer, Integer> restoredOperator = new AsyncWaitOperator<>(new MyAsyncFunction(), TIMEOUT, 6, AsyncDataStream.OutputMode.ORDERED);
    restoredTaskHarness.getStreamConfig().setStreamOperator(restoredOperator);
    restoredTaskHarness.invoke();
    restoredTaskHarness.waitForTaskRunning();
    restoredTaskHarness.processElement(new StreamRecord<>(5, initialTime + 5));
    restoredTaskHarness.processElement(new StreamRecord<>(6, initialTime + 6));
    restoredTaskHarness.processElement(new StreamRecord<>(7, initialTime + 7));
    // trigger the checkpoint while processing stream elements
    restoredTask.triggerCheckpoint(new CheckpointMetaData(checkpointId, checkpointTimestamp), CheckpointOptions.forFullCheckpoint());
    restoredTaskHarness.processElement(new StreamRecord<>(8, initialTime + 8));
    restoredTaskHarness.endInput();
    restoredTaskHarness.waitForTaskCompletion();
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
    expectedOutput.add(new StreamRecord<>(2, initialTime + 1));
    expectedOutput.add(new StreamRecord<>(4, initialTime + 2));
    expectedOutput.add(new StreamRecord<>(6, initialTime + 3));
    expectedOutput.add(new StreamRecord<>(8, initialTime + 4));
    expectedOutput.add(new StreamRecord<>(10, initialTime + 5));
    expectedOutput.add(new StreamRecord<>(12, initialTime + 6));
    expectedOutput.add(new StreamRecord<>(14, initialTime + 7));
    expectedOutput.add(new StreamRecord<>(16, initialTime + 8));
    // remove CheckpointBarrier which is not expected
    Iterator<Object> iterator = restoredTaskHarness.getOutput().iterator();
    while (iterator.hasNext()) {
        if (iterator.next() instanceof CheckpointBarrier) {
            iterator.remove();
        }
    }
    TestHarnessUtil.assertOutputEquals("StateAndRestored Test Output was not correct.", expectedOutput, restoredTaskHarness.getOutput());
}

Example 3 with MockInputSplitProvider

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

the class SourceFunctionUtil method runSourceFunction.

public static <T extends Serializable> List<T> runSourceFunction(SourceFunction<T> sourceFunction) throws Exception {
    final List<T> outputs = new ArrayList<T>();
    if (sourceFunction instanceof RichFunction) {
        AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class);
        when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig());
        RuntimeContext runtimeContext = new StreamingRuntimeContext(operator, new MockEnvironment("MockTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024), new HashMap<String, Accumulator<?, ?>>());
        ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext);
        ((RichFunction) sourceFunction).open(new Configuration());
    }
    try {
        SourceFunction.SourceContext<T> ctx = new CollectingSourceContext<T>(new Object(), outputs);
        sourceFunction.run(ctx);
    } catch (Exception e) {
        throw new RuntimeException("Cannot invoke source.", e);
    }
    return outputs;
}

Example 4 with MockInputSplitProvider

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

the class RocksDBAsyncSnapshotTest method testFullyAsyncSnapshot.

/**
	 * This ensures that asynchronous state handles are actually materialized asynchronously.
	 *
	 * <p>We use latches to block at various stages and see if the code still continues through
	 * the parts that are not asynchronous. If the checkpoint is not done asynchronously the
	 * test will simply lock forever.
	 */
@Test
public void testFullyAsyncSnapshot() throws Exception {
    LocalFileSystem localFS = new LocalFileSystem();
    localFS.initialize(new URI("file:///"), new Configuration());
    PowerMockito.stub(PowerMockito.method(FileSystem.class, "get", URI.class, Configuration.class)).toReturn(localFS);
    final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setupOutputForSingletonOperatorChain();
    testHarness.configureForKeyedStream(new KeySelector<String, String>() {

        @Override
        public String getKey(String value) throws Exception {
            return value;
        }
    }, BasicTypeInfo.STRING_TYPE_INFO);
    StreamConfig streamConfig = testHarness.getStreamConfig();
    File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
    RocksDBStateBackend backend = new RocksDBStateBackend(new MemoryStateBackend());
    backend.setDbStoragePath(dbDir.getAbsolutePath());
    streamConfig.setStateBackend(backend);
    streamConfig.setStreamOperator(new AsyncCheckpointOperator());
    final OneShotLatch delayCheckpointLatch = new OneShotLatch();
    final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
    StreamMockEnvironment mockEnv = new StreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize) {

        @Override
        public void acknowledgeCheckpoint(long checkpointId, CheckpointMetrics checkpointMetrics, SubtaskState checkpointStateHandles) {
            super.acknowledgeCheckpoint(checkpointId, checkpointMetrics);
            // even though the async checkpoint would not finish
            try {
                delayCheckpointLatch.await();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            // should be one k/v state
            assertNotNull(checkpointStateHandles.getManagedKeyedState());
            // we now know that the checkpoint went through
            ensureCheckpointLatch.trigger();
        }
    };
    testHarness.invoke(mockEnv);
    // wait for the task to be running
    for (Field field : StreamTask.class.getDeclaredFields()) {
        if (field.getName().equals("isRunning")) {
            field.setAccessible(true);
            while (!field.getBoolean(task)) {
                Thread.sleep(10);
            }
        }
    }
    task.triggerCheckpoint(new CheckpointMetaData(42, 17), CheckpointOptions.forFullCheckpoint());
    testHarness.processElement(new StreamRecord<>("Wohoo", 0));
    // now we allow the checkpoint
    delayCheckpointLatch.trigger();
    // wait for the checkpoint to go through
    ensureCheckpointLatch.await();
    testHarness.endInput();
    ExecutorService threadPool = task.getAsyncOperationsThreadPool();
    threadPool.shutdown();
    Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
    testHarness.waitForTaskCompletion();
    if (mockEnv.wasFailedExternally()) {
        fail("Unexpected exception during execution.");
    }
}

Example 5 with MockInputSplitProvider

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

the class RocksDBAsyncSnapshotTest method testCancelFullyAsyncCheckpoints.

/**
	 * This tests ensures that canceling of asynchronous snapshots works as expected and does not block.
	 * @throws Exception
	 */
@Test
@Ignore
public void testCancelFullyAsyncCheckpoints() throws Exception {
    LocalFileSystem localFS = new LocalFileSystem();
    localFS.initialize(new URI("file:///"), new Configuration());
    PowerMockito.stub(PowerMockito.method(FileSystem.class, "get", URI.class, Configuration.class)).toReturn(localFS);
    final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setupOutputForSingletonOperatorChain();
    testHarness.configureForKeyedStream(new KeySelector<String, String>() {

        @Override
        public String getKey(String value) throws Exception {
            return value;
        }
    }, BasicTypeInfo.STRING_TYPE_INFO);
    StreamConfig streamConfig = testHarness.getStreamConfig();
    File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
    BlockingStreamMemoryStateBackend memoryStateBackend = new BlockingStreamMemoryStateBackend();
    RocksDBStateBackend backend = new RocksDBStateBackend(memoryStateBackend);
    backend.setDbStoragePath(dbDir.getAbsolutePath());
    streamConfig.setStateBackend(backend);
    streamConfig.setStreamOperator(new AsyncCheckpointOperator());
    StreamMockEnvironment mockEnv = new StreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize);
    BlockingStreamMemoryStateBackend.waitFirstWriteLatch = new OneShotLatch();
    BlockingStreamMemoryStateBackend.unblockCancelLatch = new OneShotLatch();
    testHarness.invoke(mockEnv);
    // wait for the task to be running
    for (Field field : StreamTask.class.getDeclaredFields()) {
        if (field.getName().equals("isRunning")) {
            field.setAccessible(true);
            while (!field.getBoolean(task)) {
                Thread.sleep(10);
            }
        }
    }
    task.triggerCheckpoint(new CheckpointMetaData(42, 17), CheckpointOptions.forFullCheckpoint());
    testHarness.processElement(new StreamRecord<>("Wohoo", 0));
    BlockingStreamMemoryStateBackend.waitFirstWriteLatch.await();
    task.cancel();
    BlockingStreamMemoryStateBackend.unblockCancelLatch.trigger();
    testHarness.endInput();
    try {
        ExecutorService threadPool = task.getAsyncOperationsThreadPool();
        threadPool.shutdown();
        Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
        testHarness.waitForTaskCompletion();
        if (mockEnv.wasFailedExternally()) {
            throw new AsynchronousException(new InterruptedException("Exception was thrown as expected."));
        }
        fail("Operation completed. Cancel failed.");
    } catch (Exception expected) {
        AsynchronousException asynchronousException = null;
        if (expected instanceof AsynchronousException) {
            asynchronousException = (AsynchronousException) expected;
        } else if (expected.getCause() instanceof AsynchronousException) {
            asynchronousException = (AsynchronousException) expected.getCause();
        } else {
            fail("Unexpected exception: " + expected);
        }
        // we expect the exception from canceling snapshots
        Throwable innerCause = asynchronousException.getCause();
        Assert.assertTrue("Unexpected inner cause: " + innerCause, //future canceled
        innerCause instanceof CancellationException || //thread interrupted
        innerCause instanceof InterruptedException);
    }
}