Search in sources :

Example 6 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class TestTaskStateManager method prioritizedOperatorState.

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
    TaskStateSnapshot jmTaskStateSnapshot = getLastJobManagerTaskStateSnapshot();
    TaskStateSnapshot tmTaskStateSnapshot = getLastTaskManagerTaskStateSnapshot();
    if (jmTaskStateSnapshot == null) {
        return PrioritizedOperatorSubtaskState.emptyNotRestored();
    } else {
        OperatorSubtaskState jmOpState = jmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
        if (jmOpState == null) {
            return PrioritizedOperatorSubtaskState.emptyNotRestored();
        } else {
            List<OperatorSubtaskState> tmStateCollection = Collections.emptyList();
            if (tmTaskStateSnapshot != null) {
                OperatorSubtaskState tmOpState = tmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
                if (tmOpState != null) {
                    tmStateCollection = Collections.singletonList(tmOpState);
                }
            }
            PrioritizedOperatorSubtaskState.Builder builder = new PrioritizedOperatorSubtaskState.Builder(jmOpState, tmStateCollection, reportedCheckpointId);
            return builder.build();
        }
    }
}
Also used : TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) PrioritizedOperatorSubtaskState(org.apache.flink.runtime.checkpoint.PrioritizedOperatorSubtaskState) PrioritizedOperatorSubtaskState(org.apache.flink.runtime.checkpoint.PrioritizedOperatorSubtaskState) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) Nonnull(javax.annotation.Nonnull)

Example 7 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class MultipleInputStreamTaskTest method testTriggeringStopWithSavepointWithDrain.

@Test
public void testTriggeringStopWithSavepointWithDrain() throws Exception {
    SourceOperatorFactory<Integer> sourceOperatorFactory = new SourceOperatorFactory<>(new MockSource(Boundedness.CONTINUOUS_UNBOUNDED, 2), WatermarkStrategy.noWatermarks());
    CompletableFuture<Boolean> checkpointCompleted = new CompletableFuture<>();
    CheckpointResponder checkpointResponder = new TestCheckpointResponder() {

        @Override
        public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
            super.acknowledgeCheckpoint(jobID, executionAttemptID, checkpointId, checkpointMetrics, subtaskState);
            checkpointCompleted.complete(null);
        }
    };
    try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(MultipleInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).setCollectNetworkEvents().modifyStreamConfig(config -> config.setCheckpointingEnabled(true)).modifyExecutionConfig(applyObjectReuse(objectReuse)).addInput(BasicTypeInfo.INT_TYPE_INFO).addInput(BasicTypeInfo.INT_TYPE_INFO).addInput(BasicTypeInfo.INT_TYPE_INFO).setTaskStateSnapshot(1, TaskStateSnapshot.FINISHED_ON_RESTORE).setupOperatorChain(new LifeCycleMonitorMultipleInputOperatorFactory()).finishForSingletonOperatorChain(StringSerializer.INSTANCE).setCheckpointResponder(checkpointResponder).build()) {
        CompletableFuture<Boolean> triggerResult = testHarness.streamTask.triggerCheckpointAsync(new CheckpointMetaData(2, 2), CheckpointOptions.alignedNoTimeout(SavepointType.terminate(SavepointFormatType.CANONICAL), CheckpointStorageLocationReference.getDefault()));
        checkpointCompleted.whenComplete((ignored, exception) -> testHarness.streamTask.notifyCheckpointCompleteAsync(2));
        testHarness.waitForTaskCompletion();
        testHarness.finishProcessing();
        assertTrue(triggerResult.isDone());
        assertTrue(triggerResult.get());
        assertTrue(checkpointCompleted.isDone());
    }
}
Also used : EndOfData(org.apache.flink.runtime.io.network.api.EndOfData) TaskIOMetricGroup(org.apache.flink.runtime.metrics.groups.TaskIOMetricGroup) Arrays(java.util.Arrays) TestCheckpointResponder(org.apache.flink.runtime.taskmanager.TestCheckpointResponder) SharedObjects(org.apache.flink.testutils.junit.SharedObjects) NoMoreSplitsEvent(org.apache.flink.runtime.source.event.NoMoreSplitsEvent) AbstractStreamOperatorFactory(org.apache.flink.streaming.api.operators.AbstractStreamOperatorFactory) Duration(java.time.Duration) Map(java.util.Map) WatermarkStatus(org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus) TypeSerializer(org.apache.flink.api.common.typeutils.TypeSerializer) BoundedOneInput(org.apache.flink.streaming.api.operators.BoundedOneInput) Serializable(java.io.Serializable) StopMode(org.apache.flink.runtime.io.network.api.StopMode) MetricNames(org.apache.flink.runtime.metrics.MetricNames) CheckpointBarrier(org.apache.flink.runtime.io.network.api.CheckpointBarrier) Matchers.contains(org.hamcrest.Matchers.contains) WatermarkMetricOperator(org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTest.WatermarkMetricOperator) OneInputStreamOperator(org.apache.flink.streaming.api.operators.OneInputStreamOperator) Boundedness(org.apache.flink.api.connector.source.Boundedness) Counter(org.apache.flink.metrics.Counter) RunWith(org.junit.runner.RunWith) ResultPartitionWriter(org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter) TimestampAssigner(org.apache.flink.api.common.eventtime.TimestampAssigner) DataOutputView(org.apache.flink.core.memory.DataOutputView) AbstractInput(org.apache.flink.streaming.api.operators.AbstractInput) ArrayList(java.util.ArrayList) CompletingCheckpointResponder(org.apache.flink.streaming.util.CompletingCheckpointResponder) InternalOperatorMetricGroup(org.apache.flink.runtime.metrics.groups.InternalOperatorMetricGroup) Gauge(org.apache.flink.metrics.Gauge) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) TestHarnessUtil(org.apache.flink.streaming.util.TestHarnessUtil) Before(org.junit.Before) CheckpointStorageLocationReference(org.apache.flink.runtime.state.CheckpointStorageLocationReference) SourceReader(org.apache.flink.api.connector.source.SourceReader) Parameter(org.junit.runners.Parameterized.Parameter) Assert.assertTrue(org.junit.Assert.assertTrue) Test(org.junit.Test) IOException(java.io.IOException) InterceptingTaskMetricGroup(org.apache.flink.runtime.metrics.util.InterceptingTaskMetricGroup) AddSplitEvent(org.apache.flink.runtime.source.event.AddSplitEvent) StreamMultipleInputProcessor(org.apache.flink.streaming.runtime.io.StreamMultipleInputProcessor) AbstractStreamOperator(org.apache.flink.streaming.api.operators.AbstractStreamOperator) StreamOperator(org.apache.flink.streaming.api.operators.StreamOperator) JobID(org.apache.flink.api.common.JobID) UnregisteredMetricGroups(org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups) StreamTaskFinalCheckpointsTest.processMailTillCheckpointSucceeds(org.apache.flink.streaming.runtime.tasks.StreamTaskFinalCheckpointsTest.processMailTillCheckpointSucceeds) Assert(org.junit.Assert) ArrayDeque(java.util.ArrayDeque) Assert.assertEquals(org.junit.Assert.assertEquals) Input(org.apache.flink.streaming.api.operators.Input) LifeCycleMonitorMultipleInputOperatorFactory(org.apache.flink.streaming.runtime.tasks.MultipleInputStreamTaskChainedSourcesCheckpointingTest.LifeCycleMonitorMultipleInputOperatorFactory) WatermarkGenerator(org.apache.flink.api.common.eventtime.WatermarkGenerator) SavepointType(org.apache.flink.runtime.checkpoint.SavepointType) StringSerializer(org.apache.flink.api.common.typeutils.base.StringSerializer) StreamTaskFinalCheckpointsTest.triggerCheckpoint(org.apache.flink.streaming.runtime.tasks.StreamTaskFinalCheckpointsTest.triggerCheckpoint) BasicTypeInfo(org.apache.flink.api.common.typeinfo.BasicTypeInfo) ResultPartition(org.apache.flink.runtime.io.network.partition.ResultPartition) IntSerializer(org.apache.flink.api.common.typeutils.base.IntSerializer) TypeSerializerSnapshot(org.apache.flink.api.common.typeutils.TypeSerializerSnapshot) Parameterized(org.junit.runners.Parameterized) SourceReaderContext(org.apache.flink.api.connector.source.SourceReaderContext) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) CheckpointType(org.apache.flink.runtime.checkpoint.CheckpointType) InterceptingOperatorMetricGroup(org.apache.flink.runtime.metrics.util.InterceptingOperatorMetricGroup) BoundedMultiInput(org.apache.flink.streaming.api.operators.BoundedMultiInput) MockSourceReader(org.apache.flink.api.connector.source.mocks.MockSourceReader) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) WatermarkStrategy(org.apache.flink.api.common.eventtime.WatermarkStrategy) MockSourceSplitSerializer(org.apache.flink.api.connector.source.mocks.MockSourceSplitSerializer) CheckpointOptions(org.apache.flink.runtime.checkpoint.CheckpointOptions) MultipleInputStreamOperator(org.apache.flink.streaming.api.operators.MultipleInputStreamOperator) List(java.util.List) SerializedValue(org.apache.flink.util.SerializedValue) Preconditions.checkArgument(org.apache.flink.util.Preconditions.checkArgument) Matchers.containsInAnyOrder(org.hamcrest.Matchers.containsInAnyOrder) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) CancelCheckpointMarker(org.apache.flink.runtime.io.network.api.CancelCheckpointMarker) OperatorID(org.apache.flink.runtime.jobgraph.OperatorID) CheckpointResponder(org.apache.flink.runtime.taskmanager.CheckpointResponder) CheckpointMetrics(org.apache.flink.runtime.checkpoint.CheckpointMetrics) StreamConfig(org.apache.flink.streaming.api.graph.StreamConfig) SavepointFormatType(org.apache.flink.core.execution.SavepointFormatType) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) Parameters(org.junit.runners.Parameterized.Parameters) CoreMatchers.not(org.hamcrest.CoreMatchers.not) EndOfPartitionEvent(org.apache.flink.runtime.io.network.api.EndOfPartitionEvent) AbstractStreamOperatorV2(org.apache.flink.streaming.api.operators.AbstractStreamOperatorV2) Watermark(org.apache.flink.streaming.api.watermark.Watermark) ResultPartitionType(org.apache.flink.runtime.io.network.partition.ResultPartitionType) HashMap(java.util.HashMap) CompletableFuture(java.util.concurrent.CompletableFuture) Metric(org.apache.flink.metrics.Metric) SourceOperatorFactory(org.apache.flink.streaming.api.operators.SourceOperatorFactory) MockSourceSplit(org.apache.flink.api.connector.source.mocks.MockSourceSplit) TaskMetricGroup(org.apache.flink.runtime.metrics.groups.TaskMetricGroup) StreamRecord(org.apache.flink.streaming.runtime.streamrecord.StreamRecord) MockSource(org.apache.flink.api.connector.source.mocks.MockSource) OperatorMetricGroup(org.apache.flink.metrics.groups.OperatorMetricGroup) DataInputView(org.apache.flink.core.memory.DataInputView) SharedReference(org.apache.flink.testutils.junit.SharedReference) StreamOperatorParameters(org.apache.flink.streaming.api.operators.StreamOperatorParameters) WatermarkOutput(org.apache.flink.api.common.eventtime.WatermarkOutput) IsMapContaining(org.hamcrest.collection.IsMapContaining) Consumer(java.util.function.Consumer) ExecutionAttemptID(org.apache.flink.runtime.executiongraph.ExecutionAttemptID) Rule(org.junit.Rule) PartitionTestUtils(org.apache.flink.runtime.io.network.partition.PartitionTestUtils) LatencyMarker(org.apache.flink.streaming.runtime.streamrecord.LatencyMarker) Collections(java.util.Collections) MockSource(org.apache.flink.api.connector.source.mocks.MockSource) ExecutionAttemptID(org.apache.flink.runtime.executiongraph.ExecutionAttemptID) LifeCycleMonitorMultipleInputOperatorFactory(org.apache.flink.streaming.runtime.tasks.MultipleInputStreamTaskChainedSourcesCheckpointingTest.LifeCycleMonitorMultipleInputOperatorFactory) TestCheckpointResponder(org.apache.flink.runtime.taskmanager.TestCheckpointResponder) CompletingCheckpointResponder(org.apache.flink.streaming.util.CompletingCheckpointResponder) CheckpointResponder(org.apache.flink.runtime.taskmanager.CheckpointResponder) CheckpointMetrics(org.apache.flink.runtime.checkpoint.CheckpointMetrics) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) SourceOperatorFactory(org.apache.flink.streaming.api.operators.SourceOperatorFactory) CompletableFuture(java.util.concurrent.CompletableFuture) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) TestCheckpointResponder(org.apache.flink.runtime.taskmanager.TestCheckpointResponder) JobID(org.apache.flink.api.common.JobID) Test(org.junit.Test)

Example 8 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class LocalStateForwardingTest method testReportingFromSnapshotToTaskStateManager.

/**
 * This tests the forwarding of jm and tm-local state from the futures reported by the backends,
 * through the async checkpointing thread to the {@link
 * org.apache.flink.runtime.state.TaskStateManager}.
 */
@Test
public void testReportingFromSnapshotToTaskStateManager() throws Exception {
    TestTaskStateManager taskStateManager = new TestTaskStateManager();
    StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(new Configuration(), new Configuration(), new ExecutionConfig(), 1024 * 1024, new MockInputSplitProvider(), 0, taskStateManager);
    StreamTask testStreamTask = new StreamTaskTest.NoOpStreamTask(streamMockEnvironment);
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(0L, 0L);
    CheckpointMetricsBuilder checkpointMetrics = new CheckpointMetricsBuilder();
    Map<OperatorID, OperatorSnapshotFutures> snapshots = new HashMap<>(1);
    OperatorSnapshotFutures osFuture = new OperatorSnapshotFutures();
    osFuture.setKeyedStateManagedFuture(createSnapshotResult(KeyedStateHandle.class));
    osFuture.setKeyedStateRawFuture(createSnapshotResult(KeyedStateHandle.class));
    osFuture.setOperatorStateManagedFuture(createSnapshotResult(OperatorStateHandle.class));
    osFuture.setOperatorStateRawFuture(createSnapshotResult(OperatorStateHandle.class));
    osFuture.setInputChannelStateFuture(createSnapshotCollectionResult(InputChannelStateHandle.class));
    osFuture.setResultSubpartitionStateFuture(createSnapshotCollectionResult(ResultSubpartitionStateHandle.class));
    OperatorID operatorID = new OperatorID();
    snapshots.put(operatorID, osFuture);
    AsyncCheckpointRunnable checkpointRunnable = new AsyncCheckpointRunnable(snapshots, checkpointMetaData, checkpointMetrics, 0L, testStreamTask.getName(), asyncCheckpointRunnable -> {
    }, testStreamTask.getEnvironment(), testStreamTask, false, false, () -> true);
    checkpointMetrics.setAlignmentDurationNanos(0L);
    checkpointMetrics.setBytesProcessedDuringAlignment(0L);
    checkpointRunnable.run();
    TaskStateSnapshot lastJobManagerTaskStateSnapshot = taskStateManager.getLastJobManagerTaskStateSnapshot();
    TaskStateSnapshot lastTaskManagerTaskStateSnapshot = taskStateManager.getLastTaskManagerTaskStateSnapshot();
    OperatorSubtaskState jmState = lastJobManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    OperatorSubtaskState tmState = lastTaskManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    performCheck(osFuture.getKeyedStateManagedFuture(), jmState.getManagedKeyedState(), tmState.getManagedKeyedState());
    performCheck(osFuture.getKeyedStateRawFuture(), jmState.getRawKeyedState(), tmState.getRawKeyedState());
    performCheck(osFuture.getOperatorStateManagedFuture(), jmState.getManagedOperatorState(), tmState.getManagedOperatorState());
    performCheck(osFuture.getOperatorStateRawFuture(), jmState.getRawOperatorState(), tmState.getRawOperatorState());
    performCollectionCheck(osFuture.getInputChannelStateFuture(), jmState.getInputChannelState(), tmState.getInputChannelState());
    performCollectionCheck(osFuture.getResultSubpartitionStateFuture(), jmState.getResultSubpartitionState(), tmState.getResultSubpartitionState());
}
Also used : OperatorSnapshotFutures(org.apache.flink.streaming.api.operators.OperatorSnapshotFutures) Configuration(org.apache.flink.configuration.Configuration) CheckpointMetricsBuilder(org.apache.flink.runtime.checkpoint.CheckpointMetricsBuilder) HashMap(java.util.HashMap) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) OperatorID(org.apache.flink.runtime.jobgraph.OperatorID) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) KeyedStateHandle(org.apache.flink.runtime.state.KeyedStateHandle) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) TestTaskStateManager(org.apache.flink.runtime.state.TestTaskStateManager) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) ResultSubpartitionStateHandle(org.apache.flink.runtime.state.ResultSubpartitionStateHandle) OperatorStateHandle(org.apache.flink.runtime.state.OperatorStateHandle) MockInputSplitProvider(org.apache.flink.runtime.operators.testutils.MockInputSplitProvider) InputChannelStateHandle(org.apache.flink.runtime.state.InputChannelStateHandle) Test(org.junit.Test)

Example 9 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot in project flink by apache.

the class LocalStateForwardingTest method testReportingFromTaskStateManagerToResponderAndTaskLocalStateStore.

/**
 * This tests that state that was reported to the {@link
 * org.apache.flink.runtime.state.TaskStateManager} is also reported to {@link
 * org.apache.flink.runtime.taskmanager.CheckpointResponder} and {@link
 * TaskLocalStateStoreImpl}.
 */
@Test
public void testReportingFromTaskStateManagerToResponderAndTaskLocalStateStore() throws Exception {
    final JobID jobID = new JobID();
    final AllocationID allocationID = new AllocationID();
    final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
    final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(42L, 4711L);
    final CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
    final int subtaskIdx = 42;
    JobVertexID jobVertexID = new JobVertexID();
    TaskStateSnapshot jmSnapshot = new TaskStateSnapshot();
    TaskStateSnapshot tmSnapshot = new TaskStateSnapshot();
    final AtomicBoolean jmReported = new AtomicBoolean(false);
    final AtomicBoolean tmReported = new AtomicBoolean(false);
    TestCheckpointResponder checkpointResponder = new TestCheckpointResponder() {

        @Override
        public void acknowledgeCheckpoint(JobID lJobID, ExecutionAttemptID lExecutionAttemptID, long lCheckpointId, CheckpointMetrics lCheckpointMetrics, TaskStateSnapshot lSubtaskState) {
            Assert.assertEquals(jobID, lJobID);
            Assert.assertEquals(executionAttemptID, lExecutionAttemptID);
            Assert.assertEquals(checkpointMetaData.getCheckpointId(), lCheckpointId);
            Assert.assertEquals(checkpointMetrics, lCheckpointMetrics);
            jmReported.set(true);
        }
    };
    Executor executor = Executors.directExecutor();
    LocalRecoveryDirectoryProviderImpl directoryProvider = new LocalRecoveryDirectoryProviderImpl(temporaryFolder.newFolder(), jobID, jobVertexID, subtaskIdx);
    LocalRecoveryConfig localRecoveryConfig = new LocalRecoveryConfig(directoryProvider);
    TaskLocalStateStore taskLocalStateStore = new TaskLocalStateStoreImpl(jobID, allocationID, jobVertexID, subtaskIdx, localRecoveryConfig, executor) {

        @Override
        public void storeLocalState(@Nonnegative long checkpointId, @Nullable TaskStateSnapshot localState) {
            Assert.assertEquals(tmSnapshot, localState);
            tmReported.set(true);
        }
    };
    StateChangelogStorage<?> stateChangelogStorage = new InMemoryStateChangelogStorage();
    TaskStateManagerImpl taskStateManager = new TaskStateManagerImpl(jobID, executionAttemptID, taskLocalStateStore, stateChangelogStorage, null, checkpointResponder);
    taskStateManager.reportTaskStateSnapshots(checkpointMetaData, checkpointMetrics, jmSnapshot, tmSnapshot);
    Assert.assertTrue("Reporting for JM state was not called.", jmReported.get());
    Assert.assertTrue("Reporting for TM state was not called.", tmReported.get());
}
Also used : TaskStateManagerImpl(org.apache.flink.runtime.state.TaskStateManagerImpl) ExecutionAttemptID(org.apache.flink.runtime.executiongraph.ExecutionAttemptID) TaskLocalStateStore(org.apache.flink.runtime.state.TaskLocalStateStore) AllocationID(org.apache.flink.runtime.clusterframework.types.AllocationID) JobVertexID(org.apache.flink.runtime.jobgraph.JobVertexID) CheckpointMetrics(org.apache.flink.runtime.checkpoint.CheckpointMetrics) LocalRecoveryDirectoryProviderImpl(org.apache.flink.runtime.state.LocalRecoveryDirectoryProviderImpl) LocalRecoveryConfig(org.apache.flink.runtime.state.LocalRecoveryConfig) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) Executor(java.util.concurrent.Executor) InMemoryStateChangelogStorage(org.apache.flink.runtime.state.changelog.inmemory.InMemoryStateChangelogStorage) TaskLocalStateStoreImpl(org.apache.flink.runtime.state.TaskLocalStateStoreImpl) Nonnegative(javax.annotation.Nonnegative) TestCheckpointResponder(org.apache.flink.runtime.taskmanager.TestCheckpointResponder) JobID(org.apache.flink.api.common.JobID) Nullable(javax.annotation.Nullable) Test(org.junit.Test)

Example 10 with TaskStateSnapshot

use of org.apache.flink.runtime.checkpoint.TaskStateSnapshot 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 = Deadline.fromNow(Duration.ofMinutes(2));
    final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(OneInputStreamTask::new, 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;
    int numberChainedTasks = 11;
    StreamConfig streamConfig = testHarness.getStreamConfig();
    configureChainedTestingStreamOperator(streamConfig, numberChainedTasks);
    TestTaskStateManager taskStateManager = testHarness.taskStateManager;
    // reset number of restore calls
    TestingStreamOperator.numberRestoreCalls = 0;
    testHarness.invoke();
    testHarness.waitForTaskRunning();
    final OneInputStreamTask<String, String> streamTask = testHarness.getTask();
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, checkpointTimestamp);
    streamTask.triggerCheckpointAsync(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation()).get();
    // since no state was set, there shouldn't be restore calls
    assertEquals(0, TestingStreamOperator.numberRestoreCalls);
    taskStateManager.getWaitForReportLatch().await();
    assertEquals(checkpointId, taskStateManager.getReportedCheckpointId());
    testHarness.endInput();
    testHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
    final OneInputStreamTaskTestHarness<String, String> restoredTaskHarness = new OneInputStreamTaskTestHarness<>(OneInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
    restoredTaskHarness.configureForKeyedStream(keySelector, BasicTypeInfo.STRING_TYPE_INFO);
    restoredTaskHarness.setTaskStateSnapshot(checkpointId, taskStateManager.getLastJobManagerTaskStateSnapshot());
    StreamConfig restoredTaskStreamConfig = restoredTaskHarness.getStreamConfig();
    configureChainedTestingStreamOperator(restoredTaskStreamConfig, numberChainedTasks);
    TaskStateSnapshot stateHandles = taskStateManager.getLastJobManagerTaskStateSnapshot();
    Assert.assertEquals(numberChainedTasks, stateHandles.getSubtaskStateMappings().size());
    TestingStreamOperator.numberRestoreCalls = 0;
    // transfer state to new harness
    restoredTaskHarness.taskStateManager.restoreLatestCheckpointState(taskStateManager.getJobManagerTaskStateSnapshotsByCheckpointId());
    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;
    TestingStreamOperator.numberSnapshotCalls = 0;
}
Also used : Deadline(org.apache.flink.api.common.time.Deadline) StreamConfig(org.apache.flink.streaming.api.graph.StreamConfig) CheckpointMetaData(org.apache.flink.runtime.checkpoint.CheckpointMetaData) TestTaskStateManager(org.apache.flink.runtime.state.TestTaskStateManager) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) Test(org.junit.Test)

Aggregations

TaskStateSnapshot (org.apache.flink.runtime.checkpoint.TaskStateSnapshot)42 Test (org.junit.Test)28 OperatorID (org.apache.flink.runtime.jobgraph.OperatorID)19 JobID (org.apache.flink.api.common.JobID)17 ExecutionAttemptID (org.apache.flink.runtime.executiongraph.ExecutionAttemptID)16 CheckpointMetaData (org.apache.flink.runtime.checkpoint.CheckpointMetaData)13 JobManagerTaskRestore (org.apache.flink.runtime.checkpoint.JobManagerTaskRestore)13 OperatorSubtaskState (org.apache.flink.runtime.checkpoint.OperatorSubtaskState)13 CheckpointMetrics (org.apache.flink.runtime.checkpoint.CheckpointMetrics)12 TestCheckpointResponder (org.apache.flink.runtime.taskmanager.TestCheckpointResponder)9 CheckpointResponder (org.apache.flink.runtime.taskmanager.CheckpointResponder)8 IOException (java.io.IOException)7 HashMap (java.util.HashMap)6 OperatorStateHandle (org.apache.flink.runtime.state.OperatorStateHandle)6 CompletableFuture (java.util.concurrent.CompletableFuture)5 KeyedStateHandle (org.apache.flink.runtime.state.KeyedStateHandle)5 TestTaskStateManager (org.apache.flink.runtime.state.TestTaskStateManager)5 InMemoryStateChangelogStorage (org.apache.flink.runtime.state.changelog.inmemory.InMemoryStateChangelogStorage)5 ArrayList (java.util.ArrayList)4 Map (java.util.Map)4