Examples with TestTaskStateManager org.apache.flink.runtime.state.TestTaskStateManager used on opensource projects

Example 16 with TestTaskStateManager

use of org.apache.flink.runtime.state.TestTaskStateManager in project flink by apache.

the class SubtaskCheckpointCoordinatorTest method testNotifyCheckpointComplete.

@Test
public void testNotifyCheckpointComplete() throws Exception {
    TestTaskStateManager stateManager = new TestTaskStateManager();
    MockEnvironment mockEnvironment = MockEnvironment.builder().setTaskStateManager(stateManager).build();
    try (SubtaskCheckpointCoordinator subtaskCheckpointCoordinator = new MockSubtaskCheckpointCoordinatorBuilder().setEnvironment(mockEnvironment).build()) {
        final OperatorChain<?, ?> operatorChain = getOperatorChain(mockEnvironment);
        long checkpointId = 42L;
        {
            subtaskCheckpointCoordinator.notifyCheckpointComplete(checkpointId, operatorChain, () -> true);
            assertEquals(checkpointId, stateManager.getNotifiedCompletedCheckpointId());
        }
        long newCheckpointId = checkpointId + 1;
        {
            subtaskCheckpointCoordinator.notifyCheckpointComplete(newCheckpointId, operatorChain, () -> false);
            // even task is not running, state manager could still receive the notification.
            assertEquals(newCheckpointId, stateManager.getNotifiedCompletedCheckpointId());
        }
    }
}

Example 17 with TestTaskStateManager

use of org.apache.flink.runtime.state.TestTaskStateManager in project flink by apache.

the class SynchronousCheckpointITCase method createTask.

// --------------------------		Boilerplate tools copied from the TaskAsyncCallTest
// --------------------------
private Task createTask(Class<? extends TaskInvokable> invokableClass) throws Exception {
    ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
    PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
    Executor executor = mock(Executor.class);
    ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();
    TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup();
    JobInformation jobInformation = new JobInformation(new JobID(), "Job Name", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.emptyList(), Collections.emptyList());
    TaskInformation taskInformation = new TaskInformation(new JobVertexID(), "Test Task", 1, 1, invokableClass.getName(), new Configuration());
    return new Task(jobInformation, taskInformation, new ExecutionAttemptID(), new AllocationID(), 0, 0, Collections.<ResultPartitionDeploymentDescriptor>emptyList(), Collections.<InputGateDeploymentDescriptor>emptyList(), mock(MemoryManager.class), mock(IOManager.class), shuffleEnvironment, new KvStateService(new KvStateRegistry(), null, null), mock(BroadcastVariableManager.class), new TaskEventDispatcher(), ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES, new TestTaskStateManager(), mock(TaskManagerActions.class), mock(InputSplitProvider.class), mock(CheckpointResponder.class), new NoOpTaskOperatorEventGateway(), new TestGlobalAggregateManager(), TestingClassLoaderLease.newBuilder().build(), mock(FileCache.class), new TestingTaskManagerRuntimeInfo(), taskMetricGroup, consumableNotifier, partitionProducerStateChecker, executor);
}

Example 18 with TestTaskStateManager

use of org.apache.flink.runtime.state.TestTaskStateManager in project flink by apache.

the class StreamTaskMailboxTestHarnessBuilder method buildUnrestored.

public StreamTaskMailboxTestHarness<OUT> buildUnrestored() throws Exception {
    TestTaskStateManagerBuilder taskStateManagerBuilder = TestTaskStateManager.builder().setLocalRecoveryConfig(localRecoveryConfig).setCheckpointResponder(checkpointResponder);
    if (taskStateSnapshots != null) {
        taskStateManagerBuilder.setReportedCheckpointId(taskStateSnapshots.keySet().iterator().next()).setJobManagerTaskStateSnapshotsByCheckpointId(taskStateSnapshots);
    }
    TestTaskStateManager taskStateManager = taskStateManagerBuilder.build();
    StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(new JobID(), new ExecutionAttemptID(), jobConfig, streamConfig.getConfiguration(), executionConfig, memorySize, new MockInputSplitProvider(), bufferSize, taskStateManager, collectNetworkEvents);
    streamMockEnvironment.setCheckpointResponder(taskStateManager.getCheckpointResponder());
    streamMockEnvironment.setTaskManagerInfo(taskManagerRuntimeInfo);
    initializeInputs(streamMockEnvironment);
    checkState(inputGates != null, "InputGates hasn't been initialised");
    StreamElementSerializer<OUT> outputStreamRecordSerializer = new StreamElementSerializer<>(outputSerializer);
    Queue<Object> outputList = new ArrayDeque<>();
    streamMockEnvironment.addOutput(outputList, outputStreamRecordSerializer);
    streamMockEnvironment.setTaskMetricGroup(taskMetricGroup);
    for (ResultPartitionWriter writer : additionalOutputs) {
        streamMockEnvironment.addOutput(writer);
    }
    StreamTask<OUT, ?> task = taskFactory.apply(streamMockEnvironment);
    return new StreamTaskMailboxTestHarness<>(task, outputList, inputGates, streamMockEnvironment);
}

Example 19 with TestTaskStateManager

use of org.apache.flink.runtime.state.TestTaskStateManager in project flink by apache.

the class TaskCheckpointingBehaviourTest method createTask.

// ------------------------------------------------------------------------
// Utilities
// ------------------------------------------------------------------------
private static Task createTask(StreamOperator<?> op, StateBackend backend, CheckpointResponder checkpointResponder) throws IOException {
    Configuration taskConfig = new Configuration();
    StreamConfig cfg = new StreamConfig(taskConfig);
    cfg.setStreamOperator(op);
    cfg.setOperatorID(new OperatorID());
    cfg.setStateBackend(backend);
    ExecutionConfig executionConfig = new ExecutionConfig();
    JobInformation jobInformation = new JobInformation(new JobID(), "test job name", new SerializedValue<>(executionConfig), new Configuration(), Collections.emptyList(), Collections.emptyList());
    TaskInformation taskInformation = new TaskInformation(new JobVertexID(), "test task name", 1, 11, TestStreamTask.class.getName(), taskConfig);
    ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();
    return new Task(jobInformation, taskInformation, new ExecutionAttemptID(), new AllocationID(), 0, 0, Collections.<ResultPartitionDeploymentDescriptor>emptyList(), Collections.<InputGateDeploymentDescriptor>emptyList(), mock(MemoryManager.class), mock(IOManager.class), shuffleEnvironment, new KvStateService(new KvStateRegistry(), null, null), mock(BroadcastVariableManager.class), new TaskEventDispatcher(), ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES, new TestTaskStateManager(), mock(TaskManagerActions.class), mock(InputSplitProvider.class), checkpointResponder, new NoOpTaskOperatorEventGateway(), new TestGlobalAggregateManager(), TestingClassLoaderLease.newBuilder().build(), new FileCache(new String[] { EnvironmentInformation.getTemporaryFileDirectory() }, VoidPermanentBlobService.INSTANCE), new TestingTaskManagerRuntimeInfo(), UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(), new NoOpResultPartitionConsumableNotifier(), mock(PartitionProducerStateChecker.class), Executors.directExecutor());
}

Example 20 with TestTaskStateManager

use of org.apache.flink.runtime.state.TestTaskStateManager in project flink by apache.

the class TwoInputStreamTaskTest method testOperatorMetricReuse.

@Test
public void testOperatorMetricReuse() throws Exception {
    final TwoInputStreamTaskTestHarness<String, String, String> testHarness = new TwoInputStreamTaskTestHarness<>(TwoInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setupOperatorChain(new OperatorID(), new DuplicatingOperator()).chain(new OperatorID(), new OneInputStreamTaskTest.DuplicatingOperator(), BasicTypeInfo.STRING_TYPE_INFO.createSerializer(new ExecutionConfig())).chain(new OperatorID(), new OneInputStreamTaskTest.DuplicatingOperator(), BasicTypeInfo.STRING_TYPE_INFO.createSerializer(new ExecutionConfig())).finish();
    final TaskMetricGroup taskMetricGroup = TaskManagerMetricGroup.createTaskManagerMetricGroup(NoOpMetricRegistry.INSTANCE, "host", ResourceID.generate()).addJob(new JobID(), "jobname").addTask(new JobVertexID(), new ExecutionAttemptID(), "task", 0, 0);
    final StreamMockEnvironment env = new StreamMockEnvironment(testHarness.jobConfig, testHarness.taskConfig, testHarness.memorySize, new MockInputSplitProvider(), testHarness.bufferSize, new TestTaskStateManager()) {

        @Override
        public TaskMetricGroup getMetricGroup() {
            return taskMetricGroup;
        }
    };
    final Counter numRecordsInCounter = taskMetricGroup.getIOMetricGroup().getNumRecordsInCounter();
    final Counter numRecordsOutCounter = taskMetricGroup.getIOMetricGroup().getNumRecordsOutCounter();
    testHarness.invoke(env);
    testHarness.waitForTaskRunning();
    final int numRecords1 = 5;
    final int numRecords2 = 3;
    for (int x = 0; x < numRecords1; x++) {
        testHarness.processElement(new StreamRecord<>("hello"), 0, 0);
    }
    for (int x = 0; x < numRecords2; x++) {
        testHarness.processElement(new StreamRecord<>("hello"), 1, 0);
    }
    testHarness.waitForInputProcessing();
    assertEquals(numRecords1 + numRecords2, numRecordsInCounter.getCount());
    assertEquals((numRecords1 + numRecords2) * 2 * 2 * 2, numRecordsOutCounter.getCount());
    testHarness.endInput();
    testHarness.waitForTaskCompletion();
}