Search in sources :

Example 1 with StreamMockEnvironment

use of org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment in project flink by apache.

the class SourceOperatorEventTimeTest method createTestOperator.

// ------------------------------------------------------------------------
// test setup helpers
// ------------------------------------------------------------------------
private static <T> SourceOperator<T, MockSourceSplit> createTestOperator(SourceReader<T, MockSourceSplit> reader, WatermarkStrategy<T> watermarkStrategy, ProcessingTimeService timeService, boolean emitProgressiveWatermarks) throws Exception {
    final OperatorStateStore operatorStateStore = new MemoryStateBackend().createOperatorStateBackend(new MockEnvironmentBuilder().build(), "test-operator", Collections.emptyList(), new CloseableRegistry());
    final StateInitializationContext stateContext = new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
    final SourceOperator<T, MockSourceSplit> sourceOperator = new TestingSourceOperator<>(reader, watermarkStrategy, timeService, emitProgressiveWatermarks);
    sourceOperator.setup(new SourceOperatorStreamTask<Integer>(new StreamMockEnvironment(new Configuration(), new Configuration(), new ExecutionConfig(), 1L, new MockInputSplitProvider(), 1, new TestTaskStateManager())), new MockStreamConfig(new Configuration(), 1), new MockOutput<>(new ArrayList<>()));
    sourceOperator.initializeState(stateContext);
    sourceOperator.open();
    return sourceOperator;
}
Also used : OperatorStateStore(org.apache.flink.api.common.state.OperatorStateStore) MockEnvironmentBuilder(org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder) Configuration(org.apache.flink.configuration.Configuration) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) ArrayList(java.util.ArrayList) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) CloseableRegistry(org.apache.flink.core.fs.CloseableRegistry) TestTaskStateManager(org.apache.flink.runtime.state.TestTaskStateManager) StateInitializationContext(org.apache.flink.runtime.state.StateInitializationContext) StateInitializationContextImpl(org.apache.flink.runtime.state.StateInitializationContextImpl) MockStreamConfig(org.apache.flink.streaming.util.MockStreamConfig) StreamMockEnvironment(org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment) MockSourceSplit(org.apache.flink.api.connector.source.mocks.MockSourceSplit) MockInputSplitProvider(org.apache.flink.runtime.operators.testutils.MockInputSplitProvider)

Example 2 with StreamMockEnvironment

use of org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment in project flink by apache.

the class StatefulOperatorChainedTaskTest method createRunAndCheckpointOperatorChain.

private JobManagerTaskRestore createRunAndCheckpointOperatorChain(OperatorID headId, OneInputStreamOperator<String, String> headOperator, OperatorID tailId, OneInputStreamOperator<String, String> tailOperator, Optional<JobManagerTaskRestore> restore) throws Exception {
    File localRootDir = temporaryFolder.newFolder();
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(OneInputStreamTask::new, 1, 1, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO, localRootDir);
    testHarness.setupOperatorChain(headId, headOperator).chain(tailId, tailOperator, StringSerializer.INSTANCE, true).finish();
    if (restore.isPresent()) {
        JobManagerTaskRestore taskRestore = restore.get();
        testHarness.setTaskStateSnapshot(taskRestore.getRestoreCheckpointId(), taskRestore.getTaskStateSnapshot());
    }
    StreamMockEnvironment environment = new StreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.getExecutionConfig(), testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize, testHarness.getTaskStateManager());
    Configuration configuration = new Configuration();
    configuration.setString(STATE_BACKEND.key(), "rocksdb");
    File file = temporaryFolder.newFolder();
    configuration.setString(CHECKPOINTS_DIRECTORY.key(), file.toURI().toString());
    configuration.setString(INCREMENTAL_CHECKPOINTS.key(), "true");
    environment.setTaskManagerInfo(new TestingTaskManagerRuntimeInfo(configuration, System.getProperty("java.io.tmpdir").split(",|" + File.pathSeparator)));
    testHarness.invoke(environment);
    testHarness.waitForTaskRunning();
    OneInputStreamTask<String, String> streamTask = testHarness.getTask();
    processRecords(testHarness);
    triggerCheckpoint(testHarness, streamTask);
    TestTaskStateManager taskStateManager = testHarness.getTaskStateManager();
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(taskStateManager.getReportedCheckpointId(), taskStateManager.getLastJobManagerTaskStateSnapshot());
    testHarness.endInput();
    testHarness.waitForTaskCompletion();
    return jobManagerTaskRestore;
}
Also used : TestTaskStateManager(org.apache.flink.runtime.state.TestTaskStateManager) TestingTaskManagerRuntimeInfo(org.apache.flink.runtime.util.TestingTaskManagerRuntimeInfo) Configuration(org.apache.flink.configuration.Configuration) OneInputStreamTask(org.apache.flink.streaming.runtime.tasks.OneInputStreamTask) OneInputStreamTaskTestHarness(org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness) JobManagerTaskRestore(org.apache.flink.runtime.checkpoint.JobManagerTaskRestore) StreamMockEnvironment(org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment) File(java.io.File) MockInputSplitProvider(org.apache.flink.runtime.operators.testutils.MockInputSplitProvider)

Example 3 with StreamMockEnvironment

use of org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment 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 {
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(OneInputStreamTask::new, 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 = temporaryFolder.newFolder();
    RocksDBStateBackend backend = new RocksDBStateBackend(new MemoryStateBackend());
    backend.setDbStoragePath(dbDir.getAbsolutePath());
    streamConfig.setStateBackend(backend);
    streamConfig.setStreamOperator(new AsyncCheckpointOperator());
    streamConfig.setOperatorID(new OperatorID());
    final OneShotLatch delayCheckpointLatch = new OneShotLatch();
    final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
    CheckpointResponder checkpointResponderMock = new CheckpointResponder() {

        @Override
        public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
            // even though the async checkpoint would not finish
            try {
                delayCheckpointLatch.await();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            boolean hasManagedKeyedState = false;
            for (Map.Entry<OperatorID, OperatorSubtaskState> entry : subtaskState.getSubtaskStateMappings()) {
                OperatorSubtaskState state = entry.getValue();
                if (state != null) {
                    hasManagedKeyedState |= state.getManagedKeyedState() != null;
                }
            }
            // should be one k/v state
            assertTrue(hasManagedKeyedState);
            // we now know that the checkpoint went through
            ensureCheckpointLatch.trigger();
        }

        @Override
        public void reportCheckpointMetrics(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics) {
        }

        @Override
        public void declineCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointException checkpointException) {
        }
    };
    JobID jobID = new JobID();
    ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
    TestTaskStateManager taskStateManagerTestMock = new TestTaskStateManager(jobID, executionAttemptID, checkpointResponderMock, TestLocalRecoveryConfig.disabled(), new InMemoryStateChangelogStorage(), new HashMap<>(), -1L, new OneShotLatch());
    StreamMockEnvironment mockEnv = new StreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize, taskStateManagerTestMock);
    AtomicReference<Throwable> errorRef = new AtomicReference<>();
    mockEnv.setExternalExceptionHandler(errorRef::set);
    testHarness.invoke(mockEnv);
    testHarness.waitForTaskRunning();
    final OneInputStreamTask<String, String> task = testHarness.getTask();
    task.triggerCheckpointAsync(new CheckpointMetaData(42, 17), CheckpointOptions.forCheckpointWithDefaultLocation()).get();
    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 (errorRef.get() != null) {
        fail("Unexpected exception during execution.");
    }
}
Also used : OneInputStreamTask(org.apache.flink.streaming.runtime.tasks.OneInputStreamTask) CheckpointException(org.apache.flink.runtime.checkpoint.CheckpointException) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) CheckpointMetrics(org.apache.flink.runtime.checkpoint.CheckpointMetrics) OperatorID(org.apache.flink.runtime.jobgraph.OperatorID) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) InMemoryStateChangelogStorage(org.apache.flink.runtime.state.changelog.inmemory.InMemoryStateChangelogStorage) OneShotLatch(org.apache.flink.core.testutils.OneShotLatch) StreamMockEnvironment(org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment) MockInputSplitProvider(org.apache.flink.runtime.operators.testutils.MockInputSplitProvider) ExecutionAttemptID(org.apache.flink.runtime.executiongraph.ExecutionAttemptID) CheckpointResponder(org.apache.flink.runtime.taskmanager.CheckpointResponder) StreamConfig(org.apache.flink.streaming.api.graph.StreamConfig) AtomicReference(java.util.concurrent.atomic.AtomicReference) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) CheckpointException(org.apache.flink.runtime.checkpoint.CheckpointException) CancelTaskException(org.apache.flink.runtime.execution.CancelTaskException) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException) TestTaskStateManager(org.apache.flink.runtime.state.TestTaskStateManager) OneInputStreamTaskTestHarness(org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness) ExecutorService(java.util.concurrent.ExecutorService) File(java.io.File) Map(java.util.Map) HashMap(java.util.HashMap) JobID(org.apache.flink.api.common.JobID) Test(org.junit.Test)

Example 4 with StreamMockEnvironment

use of org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment 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.
 */
@Test
public void testCancelFullyAsyncCheckpoints() throws Exception {
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(OneInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setupOutputForSingletonOperatorChain();
    testHarness.configureForKeyedStream(value -> value, BasicTypeInfo.STRING_TYPE_INFO);
    StreamConfig streamConfig = testHarness.getStreamConfig();
    File dbDir = temporaryFolder.newFolder();
    final EmbeddedRocksDBStateBackend.PriorityQueueStateType timerServicePriorityQueueType = RocksDBOptions.TIMER_SERVICE_FACTORY.defaultValue();
    final int skipStreams;
    if (timerServicePriorityQueueType == EmbeddedRocksDBStateBackend.PriorityQueueStateType.HEAP) {
        // we skip the first created stream, because it is used to checkpoint the timer service,
        // which is
        // currently not asynchronous.
        skipStreams = 1;
    } else if (timerServicePriorityQueueType == EmbeddedRocksDBStateBackend.PriorityQueueStateType.ROCKSDB) {
        skipStreams = 0;
    } else {
        throw new AssertionError(String.format("Unknown timer service priority queue type %s.", timerServicePriorityQueueType));
    }
    // this is the proper instance that we need to call.
    BlockerCheckpointStreamFactory blockerCheckpointStreamFactory = new BlockerCheckpointStreamFactory(4 * 1024 * 1024) {

        int count = skipStreams;

        @Override
        public CheckpointStateOutputStream createCheckpointStateOutputStream(CheckpointedStateScope scope) throws IOException {
            if (count > 0) {
                --count;
                return new BlockingCheckpointOutputStream(new MemCheckpointStreamFactory.MemoryCheckpointOutputStream(maxSize), null, null, Integer.MAX_VALUE);
            } else {
                return super.createCheckpointStateOutputStream(scope);
            }
        }
    };
    // to avoid serialization of the above factory instance, we need to pass it in
    // through a static variable
    StateBackend stateBackend = new BackendForTestStream(new StaticForwardFactory(blockerCheckpointStreamFactory));
    RocksDBStateBackend backend = new RocksDBStateBackend(stateBackend);
    backend.setDbStoragePath(dbDir.getAbsolutePath());
    streamConfig.setStateBackend(backend);
    streamConfig.setStreamOperator(new AsyncCheckpointOperator());
    streamConfig.setOperatorID(new OperatorID());
    TestTaskStateManager taskStateManagerTestMock = new TestTaskStateManager();
    StreamMockEnvironment mockEnv = new StreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize, taskStateManagerTestMock);
    blockerCheckpointStreamFactory.setBlockerLatch(new OneShotLatch());
    blockerCheckpointStreamFactory.setWaiterLatch(new OneShotLatch());
    testHarness.invoke(mockEnv);
    testHarness.waitForTaskRunning();
    final OneInputStreamTask<String, String> task = testHarness.getTask();
    task.triggerCheckpointAsync(new CheckpointMetaData(42, 17), CheckpointOptions.forCheckpointWithDefaultLocation()).get();
    testHarness.processElement(new StreamRecord<>("Wohoo", 0));
    blockerCheckpointStreamFactory.getWaiterLatch().await();
    task.cancel();
    blockerCheckpointStreamFactory.getBlockerLatch().trigger();
    testHarness.endInput();
    ExecutorService threadPool = task.getAsyncOperationsThreadPool();
    threadPool.shutdown();
    Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
    Set<BlockingCheckpointOutputStream> createdStreams = blockerCheckpointStreamFactory.getAllCreatedStreams();
    for (BlockingCheckpointOutputStream stream : createdStreams) {
        Assert.assertTrue("Not all of the " + createdStreams.size() + " created streams have been closed.", stream.isClosed());
    }
    try {
        testHarness.waitForTaskCompletion();
        fail("Operation completed. Cancel failed.");
    } catch (Exception expected) {
        Throwable cause = expected.getCause();
        if (!(cause instanceof CancelTaskException)) {
            fail("Unexpected exception: " + expected);
        }
    }
}
Also used : OneInputStreamTask(org.apache.flink.streaming.runtime.tasks.OneInputStreamTask) OperatorID(org.apache.flink.runtime.jobgraph.OperatorID) StateBackend(org.apache.flink.runtime.state.StateBackend) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) AbstractKeyedStateBackend(org.apache.flink.runtime.state.AbstractKeyedStateBackend) MemCheckpointStreamFactory(org.apache.flink.runtime.state.memory.MemCheckpointStreamFactory) BlockerCheckpointStreamFactory(org.apache.flink.runtime.util.BlockerCheckpointStreamFactory) OneShotLatch(org.apache.flink.core.testutils.OneShotLatch) StreamMockEnvironment(org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment) MockInputSplitProvider(org.apache.flink.runtime.operators.testutils.MockInputSplitProvider) StreamConfig(org.apache.flink.streaming.api.graph.StreamConfig) BlockingCheckpointOutputStream(org.apache.flink.runtime.util.BlockingCheckpointOutputStream) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) CheckpointException(org.apache.flink.runtime.checkpoint.CheckpointException) CancelTaskException(org.apache.flink.runtime.execution.CancelTaskException) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException) BackendForTestStream(org.apache.flink.runtime.state.testutils.BackendForTestStream) TestTaskStateManager(org.apache.flink.runtime.state.TestTaskStateManager) CancelTaskException(org.apache.flink.runtime.execution.CancelTaskException) OneInputStreamTaskTestHarness(org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness) ExecutorService(java.util.concurrent.ExecutorService) CheckpointedStateScope(org.apache.flink.runtime.state.CheckpointedStateScope) File(java.io.File) Test(org.junit.Test)

Aggregations

MockInputSplitProvider (org.apache.flink.runtime.operators.testutils.MockInputSplitProvider)4 TestTaskStateManager (org.apache.flink.runtime.state.TestTaskStateManager)4 StreamMockEnvironment (org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment)4 File (java.io.File)3 MemoryStateBackend (org.apache.flink.runtime.state.memory.MemoryStateBackend)3 OneInputStreamTask (org.apache.flink.streaming.runtime.tasks.OneInputStreamTask)3 OneInputStreamTaskTestHarness (org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness)3 IOException (java.io.IOException)2 ExecutionException (java.util.concurrent.ExecutionException)2 ExecutorService (java.util.concurrent.ExecutorService)2 Configuration (org.apache.flink.configuration.Configuration)2 OneShotLatch (org.apache.flink.core.testutils.OneShotLatch)2 CheckpointException (org.apache.flink.runtime.checkpoint.CheckpointException)2 CheckpointMetaData (org.apache.flink.runtime.checkpoint.CheckpointMetaData)2 CancelTaskException (org.apache.flink.runtime.execution.CancelTaskException)2 OperatorID (org.apache.flink.runtime.jobgraph.OperatorID)2 StreamConfig (org.apache.flink.streaming.api.graph.StreamConfig)2 Test (org.junit.Test)2 ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1