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

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

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

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	Environment mockEnvironment = spy(new MockEnvironmentBuilder().build());

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	when(mockEnvironment.getTaskStateManager()).thenReturn(taskStateManager);

	StreamTask<?, ?> streamTask = new EmptyStreamTask(mockEnvironment);
	CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);

	// mock the operators
	StreamOperator<?> statelessOperator =
			mock(StreamOperator.class);

	final OperatorID operatorID = new OperatorID();
	when(statelessOperator.getOperatorID()).thenReturn(operatorID);

	// mock the returned empty snapshot result (all state handles are null)
	OperatorSnapshotFutures statelessOperatorSnapshotResult = new OperatorSnapshotFutures();
	when(statelessOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class)))
			.thenReturn(statelessOperatorSnapshotResult);

	// set up the task
	StreamOperator<?>[] streamOperators = {statelessOperator};
	OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
	when(operatorChain.getAllOperators()).thenReturn(streamOperators);

	Whitebox.setInternalState(streamTask, "isRunning", true);
	Whitebox.setInternalState(streamTask, "lock", new Object());
	Whitebox.setInternalState(streamTask, "operatorChain", operatorChain);
	Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry());
	Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
	Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", Executors.newCachedThreadPool());
	Whitebox.setInternalState(streamTask, "checkpointStorage", new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE));

	streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forCheckpointWithDefaultLocation());
	checkpointCompletedLatch.await(30, TimeUnit.SECONDS);
	streamTask.cancel();

	// ensure that 'null' was acknowledged as subtask state
	Assert.assertNull(checkpointResult.get(0));
}
 

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

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

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	// mock the operator with empty snapshot result (all state handles are null)
	StreamOperator<?> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());

	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder()
		.setTaskStateManager(taskStateManager)
		.build()) {

		RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
			mockEnvironment,
			operatorChain(statelessOperator)));

		waitTaskIsRunning(task.streamTask, task.invocationFuture);

		task.streamTask.triggerCheckpoint(
			new CheckpointMetaData(42L, 1L),
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			false);

		checkpointCompletedLatch.await(30, TimeUnit.SECONDS);

		// ensure that 'null' was acknowledged as subtask state
		Assert.assertNull(checkpointResult.get(0));

		task.streamTask.cancel();
		task.waitForTaskCompletion(true);
	}
}
 

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

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

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as stateless tasks. This
 * happens by translating an empty {@link SubtaskState} into reporting 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {

	// latch blocks until the async checkpoint thread acknowledges
	final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
	final List<SubtaskState> checkpointResult = new ArrayList<>(1);

	CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
	doAnswer(new Answer() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			SubtaskState subtaskState = invocation.getArgument(4);
			checkpointResult.add(subtaskState);
			checkpointCompletedLatch.trigger();
			return null;
		}
	}).when(checkpointResponder).acknowledgeCheckpoint(
		any(JobID.class),
		any(ExecutionAttemptID.class),
		anyLong(),
		any(CheckpointMetrics.class),
		nullable(TaskStateSnapshot.class));

	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(1L, 2L),
		new ExecutionAttemptID(1L, 2L),
		mock(TaskLocalStateStoreImpl.class),
		null,
		checkpointResponder);

	// mock the operator with empty snapshot result (all state handles are null)
	OneInputStreamOperator<String, String> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());

	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder()
		.setTaskStateManager(taskStateManager)
		.build()) {

		RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(
			mockEnvironment,
			operatorChain(statelessOperator)));

		waitTaskIsRunning(task.streamTask, task.invocationFuture);

		task.streamTask.triggerCheckpointAsync(
			new CheckpointMetaData(42L, 1L),
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			false);

		checkpointCompletedLatch.await(30, TimeUnit.SECONDS);

		// ensure that 'null' was acknowledged as subtask state
		Assert.assertNull(checkpointResult.get(0));

		task.streamTask.cancel();
		task.waitForTaskCompletion(true);
	}
}
 

@Test
public void testBeforeInvokeWithChannelStates() throws Exception {
	int numWriters = 2;
	int numGates = 2;
	RecoveryResultPartition[] partitions = new RecoveryResultPartition[numWriters];
	for (int i = 0; i < numWriters; i++) {
		partitions[i] = new RecoveryResultPartition();
	}
	RecoveryInputGate[] gates = new RecoveryInputGate[numGates];
	for (int i = 0; i < numGates; i++) {
		gates[i] = new RecoveryInputGate(partitions);
	}

	ChannelStateReader reader = new ResultPartitionTest.FiniteChannelStateReader(1, new int[] {0});
	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(),
		new ExecutionAttemptID(),
		new TestTaskLocalStateStore(),
		null,
		NoOpCheckpointResponder.INSTANCE,
		reader);
	MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setTaskStateManager(taskStateManager).build();
	mockEnvironment.addOutputs(asList(partitions));
	mockEnvironment.addInputs(asList(gates));
	StreamTask task = new MockStreamTaskBuilder(mockEnvironment).build();
	try {
		verifyResults(gates, partitions, false, false);

		task.beforeInvoke();

		verifyResults(gates, partitions, true, false);

		// execute the partition request mail inserted after input recovery completes
		task.mailboxProcessor.drain();

		for (RecoveryInputGate inputGate : gates) {
			assertTrue(inputGate.isPartitionRequested());
		}
	} finally {
		task.cleanUpInvoke();
	}
}