org.apache.flink.runtime.state.TaskLocalStateStore Java Examples

The following examples show how to use org.apache.flink.runtime.state.TaskLocalStateStore. 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: StreamTaskStateInitializerImplTest.java    From flink with Apache License 2.0 5 votes vote down vote up
private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				ProcessingTimeService processingTimeService,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 
Example #2
Source File: JvmExitOnFatalErrorTest.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up
public static void main(String[] args) throws Exception {

			System.err.println("creating task");

			// we suppress process exits via errors here to not
			// have a test that exits accidentally due to a programming error
			try {
				final Configuration taskManagerConfig = new Configuration();
				taskManagerConfig.setBoolean(TaskManagerOptions.KILL_ON_OUT_OF_MEMORY, true);

				final JobID jid = new JobID();
				final AllocationID allocationID = new AllocationID();
				final JobVertexID jobVertexId = new JobVertexID();
				final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
				final AllocationID slotAllocationId = new AllocationID();

				final SerializedValue<ExecutionConfig> execConfig = new SerializedValue<>(new ExecutionConfig());

				final JobInformation jobInformation = new JobInformation(
						jid, "Test Job", execConfig, new Configuration(),
						Collections.emptyList(), Collections.emptyList());

				final TaskInformation taskInformation = new TaskInformation(
						jobVertexId, "Test Task", 1, 1, OomInvokable.class.getName(), new Configuration());

				final MemoryManager memoryManager = new MemoryManager(1024 * 1024, 1);
				final IOManager ioManager = new IOManagerAsync();

				final NetworkEnvironment networkEnvironment = mock(NetworkEnvironment.class);
				when(networkEnvironment.createKvStateTaskRegistry(jid, jobVertexId)).thenReturn(mock(TaskKvStateRegistry.class));
				TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);
				when(networkEnvironment.getTaskEventDispatcher()).thenReturn(taskEventDispatcher);

				final TaskManagerRuntimeInfo tmInfo = TaskManagerConfiguration.fromConfiguration(taskManagerConfig);

				final Executor executor = Executors.newCachedThreadPool();

				BlobCacheService blobService =
					new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

				final TaskLocalStateStore localStateStore =
					new TaskLocalStateStoreImpl(
						jid,
						allocationID,
						jobVertexId,
						0,
						TestLocalRecoveryConfig.disabled(),
						executor);

				final TaskStateManager slotStateManager =
					new TaskStateManagerImpl(
						jid,
						executionAttemptID,
						localStateStore,
						null,
						mock(CheckpointResponder.class));

				Task task = new Task(
						jobInformation,
						taskInformation,
						executionAttemptID,
						slotAllocationId,
						0,       // subtaskIndex
						0,       // attemptNumber
						Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
						Collections.<InputGateDeploymentDescriptor>emptyList(),
						0,       // targetSlotNumber
						memoryManager,
						ioManager,
						networkEnvironment,
						new BroadcastVariableManager(),
						slotStateManager,
						new NoOpTaskManagerActions(),
						new NoOpInputSplitProvider(),
						new NoOpCheckpointResponder(),
						new TestGlobalAggregateManager(),
						blobService,
						new BlobLibraryCacheManager(
							blobService.getPermanentBlobService(),
							FlinkUserCodeClassLoaders.ResolveOrder.CHILD_FIRST,
							new String[0]),
						new FileCache(tmInfo.getTmpDirectories(), blobService.getPermanentBlobService()),
						tmInfo,
						UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
						new NoOpResultPartitionConsumableNotifier(),
						new NoOpPartitionProducerStateChecker(),
						executor);

				System.err.println("starting task thread");

				task.startTaskThread();
			}
			catch (Throwable t) {
				System.err.println("ERROR STARTING TASK");
				t.printStackTrace();
			}

			System.err.println("parking the main thread");
			CommonTestUtils.blockForeverNonInterruptibly();
		}
 
Example #3
Source File: StreamTaskStateInitializerImplTest.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up
private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	ProcessingTimeService processingTimeService = new TestProcessingTimeService();

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend,
			processingTimeService);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend,
			processingTimeService) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 
Example #4
Source File: LocalStateForwardingTest.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up
/**
 * 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(true, 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);
			}
		};

	TaskStateManagerImpl taskStateManager =
		new TaskStateManagerImpl(
			jobID,
			executionAttemptID,
			taskLocalStateStore,
			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());
}
 
Example #5
Source File: JvmExitOnFatalErrorTest.java    From flink with Apache License 2.0 4 votes vote down vote up
public static void main(String[] args) throws Exception {

			System.err.println("creating task");

			// we suppress process exits via errors here to not
			// have a test that exits accidentally due to a programming error
			try {
				final Configuration taskManagerConfig = new Configuration();
				taskManagerConfig.setBoolean(TaskManagerOptions.KILL_ON_OUT_OF_MEMORY, true);

				final JobID jid = new JobID();
				final AllocationID allocationID = new AllocationID();
				final JobVertexID jobVertexId = new JobVertexID();
				final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
				final AllocationID slotAllocationId = new AllocationID();

				final SerializedValue<ExecutionConfig> execConfig = new SerializedValue<>(new ExecutionConfig());

				final JobInformation jobInformation = new JobInformation(
						jid, "Test Job", execConfig, new Configuration(),
						Collections.emptyList(), Collections.emptyList());

				final TaskInformation taskInformation = new TaskInformation(
						jobVertexId, "Test Task", 1, 1, OomInvokable.class.getName(), new Configuration());

				final MemoryManager memoryManager = new MemoryManager(1024 * 1024, 1);
				final IOManager ioManager = new IOManagerAsync();

				final ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();

				final TaskManagerRuntimeInfo tmInfo = TaskManagerConfiguration.fromConfiguration(taskManagerConfig);

				final Executor executor = Executors.newCachedThreadPool();

				BlobCacheService blobService =
					new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class));

				final TaskLocalStateStore localStateStore =
					new TaskLocalStateStoreImpl(
						jid,
						allocationID,
						jobVertexId,
						0,
						TestLocalRecoveryConfig.disabled(),
						executor);

				final TaskStateManager slotStateManager =
					new TaskStateManagerImpl(
						jid,
						executionAttemptID,
						localStateStore,
						null,
						mock(CheckpointResponder.class));

				Task task = new Task(
						jobInformation,
						taskInformation,
						executionAttemptID,
						slotAllocationId,
						0,       // subtaskIndex
						0,       // attemptNumber
						Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
						Collections.<InputGateDeploymentDescriptor>emptyList(),
						0,       // targetSlotNumber
						memoryManager,
						ioManager,
						shuffleEnvironment,
						new KvStateService(new KvStateRegistry(), null, null),
						new BroadcastVariableManager(),
						new TaskEventDispatcher(),
						slotStateManager,
						new NoOpTaskManagerActions(),
						new NoOpInputSplitProvider(),
						new NoOpCheckpointResponder(),
						new TestGlobalAggregateManager(),
						blobService,
						new BlobLibraryCacheManager(
							blobService.getPermanentBlobService(),
							FlinkUserCodeClassLoaders.ResolveOrder.CHILD_FIRST,
							new String[0]),
						new FileCache(tmInfo.getTmpDirectories(), blobService.getPermanentBlobService()),
						tmInfo,
						UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
						new NoOpResultPartitionConsumableNotifier(),
						new NoOpPartitionProducerStateChecker(),
						executor);

				System.err.println("starting task thread");

				task.startTaskThread();
			}
			catch (Throwable t) {
				System.err.println("ERROR STARTING TASK");
				t.printStackTrace();
			}

			System.err.println("parking the main thread");
			CommonTestUtils.blockForeverNonInterruptibly();
		}
 
Example #6
Source File: StreamTaskStateInitializerImplTest.java    From flink with Apache License 2.0 4 votes vote down vote up
private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	ProcessingTimeService processingTimeService = new TestProcessingTimeService();

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend,
			processingTimeService);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend,
			processingTimeService) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 
Example #7
Source File: LocalStateForwardingTest.java    From flink with Apache License 2.0 4 votes vote down vote up
/**
 * 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(true, 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);
			}
		};

	TaskStateManagerImpl taskStateManager =
		new TaskStateManagerImpl(
			jobID,
			executionAttemptID,
			taskLocalStateStore,
			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());
}
 
Example #8
Source File: JvmExitOnFatalErrorTest.java    From flink with Apache License 2.0 4 votes vote down vote up
public static void main(String[] args) throws Exception {

			System.err.println("creating task");

			// we suppress process exits via errors here to not
			// have a test that exits accidentally due to a programming error
			try {
				final Configuration taskManagerConfig = new Configuration();
				taskManagerConfig.setBoolean(TaskManagerOptions.KILL_ON_OUT_OF_MEMORY, true);

				final JobID jid = new JobID();
				final AllocationID allocationID = new AllocationID();
				final JobVertexID jobVertexId = new JobVertexID();
				final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
				final AllocationID slotAllocationId = new AllocationID();

				final SerializedValue<ExecutionConfig> execConfig = new SerializedValue<>(new ExecutionConfig());

				final JobInformation jobInformation = new JobInformation(
						jid, "Test Job", execConfig, new Configuration(),
						Collections.emptyList(), Collections.emptyList());

				final TaskInformation taskInformation = new TaskInformation(
						jobVertexId, "Test Task", 1, 1, OomInvokable.class.getName(), new Configuration());

				final MemoryManager memoryManager = MemoryManagerBuilder.newBuilder().setMemorySize(1024 * 1024).build();
				final IOManager ioManager = new IOManagerAsync();

				final ShuffleEnvironment<?, ?> shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build();

				final Configuration copiedConf = new Configuration(taskManagerConfig);
				final TaskManagerRuntimeInfo tmInfo = TaskManagerConfiguration
					.fromConfiguration(
						taskManagerConfig,
						TaskExecutorResourceUtils.resourceSpecFromConfigForLocalExecution(copiedConf),
						InetAddress.getLoopbackAddress().getHostAddress());

				final Executor executor = Executors.newCachedThreadPool();

				final TaskLocalStateStore localStateStore =
					new TaskLocalStateStoreImpl(
						jid,
						allocationID,
						jobVertexId,
						0,
						TestLocalRecoveryConfig.disabled(),
						executor);

				final TaskStateManager slotStateManager =
					new TaskStateManagerImpl(
						jid,
						executionAttemptID,
						localStateStore,
						null,
						mock(CheckpointResponder.class));

				Task task = new Task(
						jobInformation,
						taskInformation,
						executionAttemptID,
						slotAllocationId,
						0,       // subtaskIndex
						0,       // attemptNumber
						Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
						Collections.<InputGateDeploymentDescriptor>emptyList(),
						0,       // targetSlotNumber
						memoryManager,
						ioManager,
						shuffleEnvironment,
						new KvStateService(new KvStateRegistry(), null, null),
						new BroadcastVariableManager(),
						new TaskEventDispatcher(),
						ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES,
						slotStateManager,
						new NoOpTaskManagerActions(),
						new NoOpInputSplitProvider(),
						NoOpCheckpointResponder.INSTANCE,
						new NoOpTaskOperatorEventGateway(),
						new TestGlobalAggregateManager(),
						TestingClassLoaderLease.newBuilder().build(),
						new FileCache(tmInfo.getTmpDirectories(), VoidPermanentBlobService.INSTANCE),
						tmInfo,
						UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
						new NoOpResultPartitionConsumableNotifier(),
						new NoOpPartitionProducerStateChecker(),
						executor);

				System.err.println("starting task thread");

				task.startTaskThread();
			}
			catch (Throwable t) {
				System.err.println("ERROR STARTING TASK");
				t.printStackTrace();
			}

			System.err.println("parking the main thread");
			CommonTestUtils.blockForeverNonInterruptibly();
		}
 
Example #9
Source File: LocalStateForwardingTest.java    From flink with Apache License 2.0 4 votes vote down vote up
/**
 * 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(true, 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);
			}
		};

	TaskStateManagerImpl taskStateManager =
		new TaskStateManagerImpl(
			jobID,
			executionAttemptID,
			taskLocalStateStore,
			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());
}