Task implements Runnable { public void onPartitionStateUpdate( IntermediateDataSetID resultId, IntermediateResultPartitionID partitionId, ExecutionState partitionState) throws IOException, InterruptedException { if (executionState == ExecutionState.RUNNING) { final SingleInputGate inputGate = inputGatesById.get(resultId); if (inputGate != null) { if (partitionState == ExecutionState.RUNNING) { inputGate.retriggerPartitionRequest(partitionId); } else if (partitionState == ExecutionState.CANCELED || partitionState == ExecutionState.CANCELING || partitionState == ExecutionState.FAILED) { cancelExecution(); } else { failExternally(new IllegalStateException("Received unexpected partition state " + partitionState + " for partition request. This is a bug.")); } } else { failExternally(new IllegalStateException("Received partition state for " + "unknown input gate " + resultId + ". This is a bug.")); } } else { LOG.debug("Ignoring partition state notification for not running task."); } } Task(TaskDeploymentDescriptor tdd,
				MemoryManager memManager,
				IOManager ioManager,
				NetworkEnvironment networkEnvironment,
				BroadcastVariableManager bcVarManager,
				ActorGateway taskManagerActor,
				ActorGateway jobManagerActor,
				FiniteDuration actorAskTimeout,
				LibraryCacheManager libraryCache,
				FileCache fileCache,
				TaskManagerRuntimeInfo taskManagerConfig); JobID getJobID(); JobVertexID getJobVertexId(); ExecutionAttemptID getExecutionId(); TaskInfo getTaskInfo(); Configuration getJobConfiguration(); Configuration getTaskConfiguration(); ResultPartitionWriter[] getAllWriters(); SingleInputGate[] getAllInputGates(); ResultPartition[] getProducedPartitions(); SingleInputGate getInputGateById(IntermediateDataSetID id); AccumulatorRegistry getAccumulatorRegistry(); Thread getExecutingThread(); ExecutionState getExecutionState(); boolean isCanceledOrFailed(); Throwable getFailureCause(); void startTaskThread(); @Override void run(); void stopExecution(); void cancelExecution(); void failExternally(Throwable cause); void registerExecutionListener(ActorGateway listener); void triggerCheckpointBarrier(final long checkpointID, final long checkpointTimestamp); void notifyCheckpointComplete(final long checkpointID); void onPartitionStateUpdate(
			IntermediateDataSetID resultId,
			IntermediateResultPartitionID partitionId,
			ExecutionState partitionState); @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 = Maps .newHashMapWithExpectedSize(ExecutionState.values().length); for (ExecutionState state : ExecutionState.values()) { expected.put(state, ExecutionState.FAILED); } expected.put(ExecutionState.RUNNING, 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.getPartitionId(), state); ExecutionState newTaskState = task.getExecutionState(); assertEquals(expected.get(state), newTaskState); } verify(inputGate, times(1)).retriggerPartitionRequest(eq(partitionId.getPartitionId())); }