Task implements Runnable, TaskActions { @VisibleForTesting void onPartitionStateUpdate( IntermediateDataSetID intermediateDataSetId, ResultPartitionID resultPartitionId, ExecutionState producerState) throws IOException, InterruptedException { if (executionState == ExecutionState.RUNNING) { final SingleInputGate inputGate = inputGatesById.get(intermediateDataSetId); if (inputGate != null) { if (producerState == ExecutionState.SCHEDULED || producerState == ExecutionState.DEPLOYING || producerState == ExecutionState.RUNNING || producerState == ExecutionState.FINISHED) { inputGate.retriggerPartitionRequest(resultPartitionId.getPartitionId()); } else if (producerState == ExecutionState.CANCELING || producerState == ExecutionState.CANCELED || producerState == ExecutionState.FAILED) { if (LOG.isDebugEnabled()) { LOG.debug("Cancelling task {} after the producer of partition {} with attempt ID {} has entered state {}.", taskNameWithSubtask, resultPartitionId.getPartitionId(), resultPartitionId.getProducerId(), producerState); } cancelExecution(); } else { String msg = String.format("Producer with attempt ID %s of partition %s in unexpected state %s.", resultPartitionId.getProducerId(), resultPartitionId.getPartitionId(), producerState); failExternally(new IllegalStateException(msg)); } } else { failExternally(new IllegalStateException("Received partition producer state for " + "unknown input gate " + intermediateDataSetId + ".")); } } else { LOG.debug("Task {} ignored a partition producer state notification, because it's not running.", taskNameWithSubtask); } } Task(
		JobInformation jobInformation,
		TaskInformation taskInformation,
		ExecutionAttemptID executionAttemptID,
		AllocationID slotAllocationId,
		int subtaskIndex,
		int attemptNumber,
		Collection<ResultPartitionDeploymentDescriptor> resultPartitionDeploymentDescriptors,
		Collection<InputGateDeploymentDescriptor> inputGateDeploymentDescriptors,
		int targetSlotNumber,
		TaskStateSnapshot taskStateHandles,
		MemoryManager memManager,
		IOManager ioManager,
		NetworkEnvironment networkEnvironment,
		BroadcastVariableManager bcVarManager,
		TaskManagerActions taskManagerActions,
		InputSplitProvider inputSplitProvider,
		CheckpointResponder checkpointResponder,
		BlobCacheService blobService,
		LibraryCacheManager libraryCache,
		FileCache fileCache,
		TaskManagerRuntimeInfo taskManagerConfig,
		@Nonnull TaskMetricGroup metricGroup,
		ResultPartitionConsumableNotifier resultPartitionConsumableNotifier,
		PartitionProducerStateChecker partitionProducerStateChecker,
		Executor executor); JobID getJobID(); JobVertexID getJobVertexId(); ExecutionAttemptID getExecutionId(); AllocationID getAllocationId(); TaskInfo getTaskInfo(); Configuration getJobConfiguration(); Configuration getTaskConfiguration(); ResultPartitionWriter[] getAllWriters(); SingleInputGate[] getAllInputGates(); ResultPartition[] getProducedPartitions(); SingleInputGate getInputGateById(IntermediateDataSetID id); AccumulatorRegistry getAccumulatorRegistry(); TaskMetricGroup getMetricGroup(); Thread getExecutingThread(); ExecutionState getExecutionState(); boolean isCanceledOrFailed(); Throwable getFailureCause(); void startTaskThread(); @Override void run(); void stopExecution(); void cancelExecution(); @Override void failExternally(Throwable cause); void registerExecutionListener(TaskExecutionStateListener listener); @Override void triggerPartitionProducerStateCheck(
		JobID jobId,
		final IntermediateDataSetID intermediateDataSetId,
		final ResultPartitionID resultPartitionId); void triggerCheckpointBarrier(
			final long checkpointID,
			long checkpointTimestamp,
			final CheckpointOptions checkpointOptions); void notifyCheckpointComplete(final long checkpointID); @Override String toString();  }
@Test public void testOnPartitionStateUpdate() throws Exception { IntermediateDataSetID resultId = new IntermediateDataSetID(); ResultPartitionID partitionId = new ResultPartitionID(); SingleInputGate inputGate = mock(SingleInputGate.class); when(inputGate.getConsumedResultId()).thenReturn(resultId); final Task task = createTask(InvokableBlockingInInvoke.class); setInputGate(task, inputGate); final Map<ExecutionState, ExecutionState> expected = new HashMap<>(ExecutionState.values().length); for (ExecutionState state : ExecutionState.values()) { expected.put(state, ExecutionState.FAILED); } expected.put(ExecutionState.RUNNING, ExecutionState.RUNNING); expected.put(ExecutionState.SCHEDULED, ExecutionState.RUNNING); expected.put(ExecutionState.DEPLOYING, ExecutionState.RUNNING); expected.put(ExecutionState.FINISHED, ExecutionState.RUNNING); expected.put(ExecutionState.CANCELED, ExecutionState.CANCELING); expected.put(ExecutionState.CANCELING, ExecutionState.CANCELING); expected.put(ExecutionState.FAILED, ExecutionState.CANCELING); for (ExecutionState state : ExecutionState.values()) { setState(task, ExecutionState.RUNNING); task.onPartitionStateUpdate(resultId, partitionId, state); ExecutionState newTaskState = task.getExecutionState(); assertEquals(expected.get(state), newTaskState); } verify(inputGate, times(4)).retriggerPartitionRequest(eq(partitionId.getPartitionId())); }