Java Code Examples for org.apache.flink.runtime.state.TestTaskStateManager#getLastJobManagerTaskStateSnapshot()

The following examples show how to use org.apache.flink.runtime.state.TestTaskStateManager#getLastJobManagerTaskStateSnapshot() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.

Example 1

Source File: StatefulOperatorChainedTaskTest.java From Flink-CEPplus with Apache License 2.0

4 votes

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;
}

Example 2

Source File: LocalStateForwardingTest.java From Flink-CEPplus with Apache License 2.0

4 votes

/**
 * 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() {

	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);
	CheckpointMetrics checkpointMetrics = new CheckpointMetrics();

	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));

	OperatorID operatorID = new OperatorID();
	snapshots.put(operatorID, osFuture);

	StreamTask.AsyncCheckpointRunnable checkpointRunnable =
		new StreamTask.AsyncCheckpointRunnable(
			testStreamTask,
			snapshots,
			checkpointMetaData,
			checkpointMetrics,
			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());
}

Example 3

Source File: StatefulOperatorChainedTaskTest.java From flink with Apache License 2.0

4 votes

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;
}

Example 4

Source File: LocalStateForwardingTest.java From flink with Apache License 2.0

4 votes

/**
 * 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() {

	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);
	CheckpointMetrics checkpointMetrics = new CheckpointMetrics();

	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));

	OperatorID operatorID = new OperatorID();
	snapshots.put(operatorID, osFuture);

	StreamTask.AsyncCheckpointRunnable checkpointRunnable =
		new StreamTask.AsyncCheckpointRunnable(
			testStreamTask,
			snapshots,
			checkpointMetaData,
			checkpointMetrics,
			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());
}

Example 5

Source File: StatefulOperatorChainedTaskTest.java From flink with Apache License 2.0

4 votes

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;
}

Example 6

Source File: LocalStateForwardingTest.java From flink with Apache License 2.0

4 votes

/**
 * 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);
	CheckpointMetrics checkpointMetrics = new CheckpointMetrics();

	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 -> {},
		asyncCheckpointRunnable -> {},
		testStreamTask.getEnvironment(),
		testStreamTask);

	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());
}

Example 7

Source File: OneInputStreamTaskTest.java From Flink-CEPplus with Apache License 2.0

2 votes

/**
 * 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 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;
	OneShotLatch waitForAcknowledgeLatch = new OneShotLatch();

	taskStateManager.setWaitForReportLatch(waitForAcknowledgeLatch);

	// reset number of restore calls
	TestingStreamOperator.numberRestoreCalls = 0;

	testHarness.invoke();
	testHarness.waitForTaskRunning(deadline.timeLeft().toMillis());

	final OneInputStreamTask<String, String> streamTask = testHarness.getTask();

	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, checkpointTimestamp);

	while (!streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation())) {}

	// since no state was set, there shouldn't be restore calls
	assertEquals(0, TestingStreamOperator.numberRestoreCalls);

	waitForAcknowledgeLatch.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;
}

Example 8

Source File: OneInputStreamTaskTest.java From flink with Apache License 2.0

2 votes

/**
 * 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 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;
	OneShotLatch waitForAcknowledgeLatch = new OneShotLatch();

	taskStateManager.setWaitForReportLatch(waitForAcknowledgeLatch);

	// reset number of restore calls
	TestingStreamOperator.numberRestoreCalls = 0;

	testHarness.invoke();
	testHarness.waitForTaskRunning(deadline.timeLeft().toMillis());

	final OneInputStreamTask<String, String> streamTask = testHarness.getTask();

	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, checkpointTimestamp);

	while (!streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation(), false)) {}

	// since no state was set, there shouldn't be restore calls
	assertEquals(0, TestingStreamOperator.numberRestoreCalls);

	waitForAcknowledgeLatch.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;
}

Example 9

Source File: OneInputStreamTaskTest.java From flink with Apache License 2.0

2 votes

/**
 * 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 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;
	OneShotLatch waitForAcknowledgeLatch = new OneShotLatch();

	taskStateManager.setWaitForReportLatch(waitForAcknowledgeLatch);

	// 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(), false).get();

	// since no state was set, there shouldn't be restore calls
	assertEquals(0, TestingStreamOperator.numberRestoreCalls);

	waitForAcknowledgeLatch.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;
}