org.apache.flink.runtime.io.network.TaskEventDispatcher Java Examples

@Test
public void testResumeConsumption() {
	final InputChannelID inputChannelID = new InputChannelID();
	final PartitionRequestQueue partitionRequestQueue = new PartitionRequestQueue();
	final TestViewReader testViewReader = new TestViewReader(inputChannelID, 2, partitionRequestQueue);
	final PartitionRequestServerHandler serverHandler = new PartitionRequestServerHandler(
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		partitionRequestQueue);
	final EmbeddedChannel channel = new EmbeddedChannel(serverHandler);
	partitionRequestQueue.notifyReaderCreated(testViewReader);

	// Write the message of resume consumption to server
	channel.writeInbound(new ResumeConsumption(inputChannelID));
	channel.runPendingTasks();

	assertTrue(testViewReader.consumptionResumed);
}
 

/**
 * Tests that {@link PartitionRequestServerHandler} responds {@link ErrorResponse} with wrapped
 * {@link PartitionNotFoundException} after receiving invalid {@link PartitionRequest}.
 */
@Test
public void testResponsePartitionNotFoundException() {
	final PartitionRequestServerHandler serverHandler = new PartitionRequestServerHandler(
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new PartitionRequestQueue(),
		true);
	final EmbeddedChannel channel = new EmbeddedChannel(serverHandler);
	final ResultPartitionID partitionId = new ResultPartitionID();

	// Write the message of partition request to server
	channel.writeInbound(new PartitionRequest(partitionId, 0, new InputChannelID(), 2));
	channel.runPendingTasks();

	// Read the response message after handling partition request
	final Object msg = channel.readOutbound();
	assertThat(msg, instanceOf(ErrorResponse.class));

	final ErrorResponse err = (ErrorResponse) msg;
	assertThat(err.cause, instanceOf(PartitionNotFoundException.class));

	final ResultPartitionID actualPartitionId = ((PartitionNotFoundException) err.cause).getPartitionId();
	assertThat(partitionId, is(actualPartitionId));
}
 

private LocalInputChannel createLocalInputChannel(
		SingleInputGate inputGate,
		ResultPartitionManager partitionManager,
		Tuple2<Integer, Integer> initialAndMaxRequestBackoff)
		throws IOException, InterruptedException {

	return new LocalInputChannel(
			inputGate,
			0,
			new ResultPartitionID(),
			partitionManager,
			mock(TaskEventDispatcher.class),
			initialAndMaxRequestBackoff._1(),
			initialAndMaxRequestBackoff._2(),
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());
}
 

public UnknownInputChannel(
		SingleInputGate gate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventDispatcher taskEventDispatcher,
		ConnectionManager connectionManager,
		int initialBackoff,
		int maxBackoff,
		TaskIOMetricGroup metrics) {

	super(gate, channelIndex, partitionId, initialBackoff, maxBackoff, null, null);

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
	this.connectionManager = checkNotNull(connectionManager);
	this.metrics = checkNotNull(metrics);
	this.initialBackoff = initialBackoff;
	this.maxBackoff = maxBackoff;
}
 

/**
 * Tests that an update channel does not trigger a partition request before the UDF has
 * requested any partitions. Otherwise, this can lead to races when registering a listener at
 * the gate (e.g. in UnionInputGate), which can result in missed buffer notifications at the
 * listener.
 */
@Test
public void testUpdateChannelBeforeRequest() throws Exception {
	SingleInputGate inputGate = createInputGate(1);

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	InputChannel unknown = new UnknownInputChannel(
		inputGate,
		0,
		new ResultPartitionID(),
		partitionManager,
		new TaskEventDispatcher(),
		new LocalConnectionManager(),
		0, 0, UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	inputGate.setInputChannel(unknown.partitionId.getPartitionId(), unknown);

	// Update to a local channel and verify that no request is triggered
	inputGate.updateInputChannel(new InputChannelDeploymentDescriptor(
		unknown.partitionId,
		ResultPartitionLocation.createLocal()));

	verify(partitionManager, never()).createSubpartitionView(
		any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));
}
 

private static ShuffleEnvironment<?, ?> createShuffleEnvironment(
		TaskManagerServicesConfiguration taskManagerServicesConfiguration,
		TaskEventDispatcher taskEventDispatcher,
		MetricGroup taskManagerMetricGroup,
		Executor ioExecutor) throws FlinkException {

	final ShuffleEnvironmentContext shuffleEnvironmentContext = new ShuffleEnvironmentContext(
		taskManagerServicesConfiguration.getConfiguration(),
		taskManagerServicesConfiguration.getResourceID(),
		taskManagerServicesConfiguration.getNetworkMemorySize(),
		taskManagerServicesConfiguration.isLocalCommunicationOnly(),
		taskManagerServicesConfiguration.getBindAddress(),
		taskEventDispatcher,
		taskManagerMetricGroup,
		ioExecutor);

	return ShuffleServiceLoader
		.loadShuffleServiceFactory(taskManagerServicesConfiguration.getConfiguration())
		.createShuffleEnvironment(shuffleEnvironmentContext);
}
 

private static ShuffleEnvironment<?, ?> createShuffleEnvironment(
		TaskManagerServicesConfiguration taskManagerServicesConfiguration,
		TaskEventDispatcher taskEventDispatcher,
		MetricGroup taskManagerMetricGroup) throws FlinkException {

	final ShuffleEnvironmentContext shuffleEnvironmentContext = new ShuffleEnvironmentContext(
		taskManagerServicesConfiguration.getConfiguration(),
		taskManagerServicesConfiguration.getResourceID(),
		taskManagerServicesConfiguration.getMaxJvmHeapMemory(),
		taskManagerServicesConfiguration.isLocalCommunicationOnly(),
		taskManagerServicesConfiguration.getTaskManagerAddress(),
		taskEventDispatcher,
		taskManagerMetricGroup);

	return ShuffleServiceLoader
		.loadShuffleServiceFactory(taskManagerServicesConfiguration.getConfiguration())
		.createShuffleEnvironment(shuffleEnvironmentContext);
}
 

public TestLocalInputChannelConsumer(
		int subpartitionIndex,
		int numberOfInputChannels,
		int numberOfExpectedBuffersPerChannel,
		BufferPool bufferPool,
		ResultPartitionManager partitionManager,
		TaskEventDispatcher taskEventDispatcher,
		ResultPartitionID[] consumedPartitionIds) throws IOException, InterruptedException {

	checkArgument(numberOfInputChannels >= 1);
	checkArgument(numberOfExpectedBuffersPerChannel >= 1);

	this.inputGate = new SingleInputGateBuilder()
		.setConsumedSubpartitionIndex(subpartitionIndex)
		.setNumberOfChannels(numberOfInputChannels)
		.setBufferPoolFactory(bufferPool)
		.build();

	// Setup input channels
	for (int i = 0; i < numberOfInputChannels; i++) {
		InputChannelBuilder.newBuilder()
			.setChannelIndex(i)
			.setPartitionManager(partitionManager)
			.setPartitionId(consumedPartitionIds[i])
			.setTaskEventPublisher(taskEventDispatcher)
			.buildLocalAndSetToGate(inputGate);
	}

	inputGate.setup();

	this.numberOfInputChannels = numberOfInputChannels;
	this.numberOfExpectedBuffersPerChannel = numberOfExpectedBuffersPerChannel;
}
 

@Override
public void start(ResultPartitionProvider partitionProvider, TaskEventDispatcher taskEventDispatcher) throws IOException {
	NettyProtocol partitionRequestProtocol = new NettyProtocol(
		partitionProvider,
		taskEventDispatcher,
		client.getConfig().isCreditBasedEnabled());

	client.init(partitionRequestProtocol, bufferPool);
	server.init(partitionRequestProtocol, bufferPool);
}
 

PartitionRequestServerHandler(
	ResultPartitionProvider partitionProvider,
	TaskEventDispatcher taskEventDispatcher,
	PartitionRequestQueue outboundQueue,
	boolean creditBasedEnabled) {

	this.partitionProvider = partitionProvider;
	this.taskEventDispatcher = taskEventDispatcher;
	this.outboundQueue = outboundQueue;
	this.creditBasedEnabled = creditBasedEnabled;
}
 

public LocalInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ResultPartitionManager partitionManager,
	TaskEventDispatcher taskEventDispatcher,
	TaskIOMetricGroup metrics) {

	this(inputGate, channelIndex, partitionId, partitionManager, taskEventDispatcher,
		0, 0, metrics);
}
 

private EmbeddedChannel createEmbeddedChannel() {
	NettyProtocol protocol = new NettyProtocol(
			mock(ResultPartitionProvider.class),
			mock(TaskEventDispatcher.class));

	return new EmbeddedChannel(protocol.getClientChannelHandlers());
}
 

@Test
public void testExecutionFailsInNetworkRegistration() throws Exception {
	// mock a network manager that rejects registration
	final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	final ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
	final PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
	final TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);

	final NetworkEnvironment network = mock(NetworkEnvironment.class);
	when(network.getResultPartitionManager()).thenReturn(partitionManager);
	when(network.getDefaultIOMode()).thenReturn(IOManager.IOMode.SYNC);
	when(network.getTaskEventDispatcher()).thenReturn(taskEventDispatcher);
	doThrow(new RuntimeException("buffers")).when(network).registerTask(any(Task.class));

	final QueuedNoOpTaskManagerActions taskManagerActions = new QueuedNoOpTaskManagerActions();
	final Task task = new TaskBuilder()
		.setTaskManagerActions(taskManagerActions)
		.setConsumableNotifier(consumableNotifier)
		.setPartitionProducerStateChecker(partitionProducerStateChecker)
		.setNetworkEnvironment(network)
		.build();

	// should fail
	task.run();

	// verify final state
	assertEquals(ExecutionState.FAILED, task.getExecutionState());
	assertTrue(task.isCanceledOrFailed());
	assertTrue(task.getFailureCause().getMessage().contains("buffers"));

	taskManagerActions.validateListenerMessage(
		ExecutionState.FAILED, task, new RuntimeException("buffers"));
}
 

private SuperstepBarrier initSuperstepBarrier() {
	SuperstepBarrier barrier = new SuperstepBarrier(getUserCodeClassLoader());
	TaskEventDispatcher taskEventDispatcher = getEnvironment().getTaskEventDispatcher();
	ResultPartitionID partitionId = toSyncPartitionId;
	taskEventDispatcher.subscribeToEvent(partitionId, barrier, AllWorkersDoneEvent.class);
	taskEventDispatcher.subscribeToEvent(partitionId, barrier, TerminationEvent.class);
	return barrier;
}
 

private SuperstepBarrier initSuperstepBarrier() {
	SuperstepBarrier barrier = new SuperstepBarrier(getUserCodeClassLoader());
	TaskEventDispatcher taskEventDispatcher = getEnvironment().getTaskEventDispatcher();
	ResultPartitionID partitionId = toSyncPartitionId;
	taskEventDispatcher.subscribeToEvent(partitionId, barrier, AllWorkersDoneEvent.class);
	taskEventDispatcher.subscribeToEvent(partitionId, barrier, TerminationEvent.class);
	return barrier;
}
 

private EmbeddedChannel createEmbeddedChannel() {
	NettyProtocol protocol = new NettyProtocol(
			mock(ResultPartitionProvider.class),
			mock(TaskEventDispatcher.class),
			true);

	return new EmbeddedChannel(protocol.getClientChannelHandlers());
}
 

public TestLocalInputChannelConsumer(
		int subpartitionIndex,
		int numberOfInputChannels,
		int numberOfExpectedBuffersPerChannel,
		BufferPool bufferPool,
		ResultPartitionManager partitionManager,
		TaskEventDispatcher taskEventDispatcher,
		ResultPartitionID[] consumedPartitionIds) throws IOException, InterruptedException {

	checkArgument(numberOfInputChannels >= 1);
	checkArgument(numberOfExpectedBuffersPerChannel >= 1);

	this.inputGate = new SingleInputGateBuilder()
		.setConsumedSubpartitionIndex(subpartitionIndex)
		.setNumberOfChannels(numberOfInputChannels)
		.setBufferPoolFactory(bufferPool)
		.build();
	InputChannel[] inputChannels = new InputChannel[numberOfInputChannels];

	// Setup input channels
	for (int i = 0; i < numberOfInputChannels; i++) {
		inputChannels[i] = InputChannelBuilder.newBuilder()
			.setChannelIndex(i)
			.setPartitionManager(partitionManager)
			.setPartitionId(consumedPartitionIds[i])
			.setTaskEventPublisher(taskEventDispatcher)
			.buildLocalChannel(inputGate);
	}

	setupInputGate(inputGate, inputChannels);

	this.numberOfInputChannels = numberOfInputChannels;
	this.numberOfExpectedBuffersPerChannel = numberOfExpectedBuffersPerChannel;
}
 

private SuperstepBarrier initSuperstepBarrier() {
	SuperstepBarrier barrier = new SuperstepBarrier(getUserCodeClassLoader());
	TaskEventDispatcher taskEventDispatcher = getEnvironment().getTaskEventDispatcher();
	ResultPartitionID partitionId = toSyncPartitionId;
	taskEventDispatcher.subscribeToEvent(partitionId, barrier, AllWorkersDoneEvent.class);
	taskEventDispatcher.subscribeToEvent(partitionId, barrier, TerminationEvent.class);
	return barrier;
}
 

/**
 * Sets up the environment including buffer pools and netty threads.
 *
 * @param writers
 * 		number of writers
 * @param channels
 * 		outgoing channels per writer
 * @param localMode
 * 		only local channels?
 * @param senderBufferPoolSize
 * 		buffer pool size for the sender (set to <tt>-1</tt> for default)
 * @param receiverBufferPoolSize
 * 		buffer pool size for the receiver (set to <tt>-1</tt> for default)
 */
public void setUp(
		int writers,
		int channels,
		boolean broadcastMode,
		boolean localMode,
		int senderBufferPoolSize,
		int receiverBufferPoolSize,
		Configuration config) throws Exception {
	this.broadcastMode = broadcastMode;
	this.localMode = localMode;
	this.channels = channels;
	this.partitionIds = new ResultPartitionID[writers];
	if (senderBufferPoolSize == -1) {
		senderBufferPoolSize = Math.max(2048, writers * channels * 4);
	}
	if (receiverBufferPoolSize == -1) {
		receiverBufferPoolSize = Math.max(2048, writers * channels * 4);
	}

	senderEnv = createShuffleEnvironment(senderBufferPoolSize, config);
	this.dataPort = senderEnv.start();
	if (localMode && senderBufferPoolSize == receiverBufferPoolSize) {
		receiverEnv = senderEnv;
	}
	else {
		receiverEnv = createShuffleEnvironment(receiverBufferPoolSize, config);
		receiverEnv.start();
	}

	gateFactory = new SingleInputGateFactory(
		location,
		receiverEnv.getConfiguration(),
		receiverEnv.getConnectionManager(),
		receiverEnv.getResultPartitionManager(),
		new TaskEventDispatcher(),
		receiverEnv.getNetworkBufferPool());

	generatePartitionIds();
}
 

private EmbeddedChannel createEmbeddedChannel() {
	NettyProtocol protocol = new NettyProtocol(
			mock(ResultPartitionProvider.class),
			mock(TaskEventDispatcher.class),
			true);

	return new EmbeddedChannel(protocol.getClientChannelHandlers());
}
 

/**
 * Sets up the environment including buffer pools and netty threads.
 *
 * @param writers
 * 		number of writers
 * @param channels
 * 		outgoing channels per writer
 * @param localMode
 * 		only local channels?
 * @param senderBufferPoolSize
 * 		buffer pool size for the sender (set to <tt>-1</tt> for default)
 * @param receiverBufferPoolSize
 * 		buffer pool size for the receiver (set to <tt>-1</tt> for default)
 */
public void setUp(
		int writers,
		int channels,
		boolean localMode,
		int senderBufferPoolSize,
		int receiverBufferPoolSize,
		Configuration config) throws Exception {
	this.localMode = localMode;
	this.channels = channels;
	this.partitionIds = new ResultPartitionID[writers];
	if (senderBufferPoolSize == -1) {
		senderBufferPoolSize = Math.max(2048, writers * channels * 4);
	}
	if (receiverBufferPoolSize == -1) {
		receiverBufferPoolSize = Math.max(2048, writers * channels * 4);
	}

	senderEnv = createShuffleEnvironment(senderBufferPoolSize, config);
	this.dataPort = senderEnv.start();
	if (localMode && senderBufferPoolSize == receiverBufferPoolSize) {
		receiverEnv = senderEnv;
	}
	else {
		receiverEnv = createShuffleEnvironment(receiverBufferPoolSize, config);
		receiverEnv.start();
	}

	gateFactory = new SingleInputGateBenchmarkFactory(
		location,
		receiverEnv.getConfiguration(),
		receiverEnv.getConnectionManager(),
		receiverEnv.getResultPartitionManager(),
		new TaskEventDispatcher(),
		receiverEnv.getNetworkBufferPool());

	generatePartitionIds();
}
 

private NetworkEnvironment createNettyNetworkEnvironment(
		@SuppressWarnings("SameParameterValue") int bufferPoolSize, Configuration config) throws Exception {

	int segmentSize =
		checkedDownCast(
			MemorySize.parse(config.getString(TaskManagerOptions.MEMORY_SEGMENT_SIZE))
				.getBytes());

	// we need this because many configs have been written with a "-1" entry
	// similar to TaskManagerServicesConfiguration#fromConfiguration()
	// -> please note that this directly influences the number of netty threads!
	int slots = config.getInteger(TaskManagerOptions.NUM_TASK_SLOTS, 1);
	if (slots == -1) {
		slots = 1;
	}

	final NetworkBufferPool bufferPool = new NetworkBufferPool(bufferPoolSize, segmentSize);

	final NettyConnectionManager nettyConnectionManager = new NettyConnectionManager(
		new NettyConfig(LOCAL_ADDRESS, 0, segmentSize, slots, config));

	return new NetworkEnvironment(
		bufferPool,
		nettyConnectionManager,
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new KvStateRegistry(),
		null,
		null,
		IOMode.SYNC,
		TaskManagerOptions.NETWORK_REQUEST_BACKOFF_INITIAL.defaultValue(),
		TaskManagerOptions.NETWORK_REQUEST_BACKOFF_MAX.defaultValue(),
		TaskManagerOptions.NETWORK_BUFFERS_PER_CHANNEL.defaultValue(),
		TaskManagerOptions.NETWORK_EXTRA_BUFFERS_PER_GATE.defaultValue(),
		true);
}
 

TaskManagerServices(
	UnresolvedTaskManagerLocation unresolvedTaskManagerLocation,
	long managedMemorySize,
	IOManager ioManager,
	ShuffleEnvironment<?, ?> shuffleEnvironment,
	KvStateService kvStateService,
	BroadcastVariableManager broadcastVariableManager,
	TaskSlotTable<Task> taskSlotTable,
	JobTable jobTable,
	JobLeaderService jobLeaderService,
	TaskExecutorLocalStateStoresManager taskManagerStateStore,
	TaskEventDispatcher taskEventDispatcher,
	ExecutorService ioExecutor,
	LibraryCacheManager libraryCacheManager) {

	this.unresolvedTaskManagerLocation = Preconditions.checkNotNull(unresolvedTaskManagerLocation);
	this.managedMemorySize = managedMemorySize;
	this.ioManager = Preconditions.checkNotNull(ioManager);
	this.shuffleEnvironment = Preconditions.checkNotNull(shuffleEnvironment);
	this.kvStateService = Preconditions.checkNotNull(kvStateService);
	this.broadcastVariableManager = Preconditions.checkNotNull(broadcastVariableManager);
	this.taskSlotTable = Preconditions.checkNotNull(taskSlotTable);
	this.jobTable = Preconditions.checkNotNull(jobTable);
	this.jobLeaderService = Preconditions.checkNotNull(jobLeaderService);
	this.taskManagerStateStore = Preconditions.checkNotNull(taskManagerStateStore);
	this.taskEventDispatcher = Preconditions.checkNotNull(taskEventDispatcher);
	this.ioExecutor = Preconditions.checkNotNull(ioExecutor);
	this.libraryCacheManager = Preconditions.checkNotNull(libraryCacheManager);
}
 

/**
 * Creates and returns the task manager services.
 *
 * @param taskManagerServicesConfiguration task manager configuration
 * @param taskManagerMetricGroup metric group of the task manager
 * @param taskIOExecutor executor for async IO operations
 * @return task manager components
 * @throws Exception
 */
public static TaskManagerServices fromConfiguration(
		TaskManagerServicesConfiguration taskManagerServicesConfiguration,
		MetricGroup taskManagerMetricGroup,
		Executor taskIOExecutor) throws Exception {

	// pre-start checks
	checkTempDirs(taskManagerServicesConfiguration.getTmpDirPaths());

	final TaskEventDispatcher taskEventDispatcher = new TaskEventDispatcher();

	// start the I/O manager, it will create some temp directories.
	final IOManager ioManager = new IOManagerAsync(taskManagerServicesConfiguration.getTmpDirPaths());

	final ShuffleEnvironment<?, ?> shuffleEnvironment = createShuffleEnvironment(
		taskManagerServicesConfiguration,
		taskEventDispatcher,
		taskManagerMetricGroup);
	final int dataPort = shuffleEnvironment.start();

	final KvStateService kvStateService = KvStateService.fromConfiguration(taskManagerServicesConfiguration);
	kvStateService.start();

	final TaskManagerLocation taskManagerLocation = new TaskManagerLocation(
		taskManagerServicesConfiguration.getResourceID(),
		taskManagerServicesConfiguration.getTaskManagerAddress(),
		dataPort);

	// this call has to happen strictly after the network stack has been initialized
	final MemoryManager memoryManager = createMemoryManager(taskManagerServicesConfiguration);
	final long managedMemorySize = memoryManager.getMemorySize();

	final BroadcastVariableManager broadcastVariableManager = new BroadcastVariableManager();

	final int numOfSlots = taskManagerServicesConfiguration.getNumberOfSlots();
	final List<ResourceProfile> resourceProfiles =
		Collections.nCopies(numOfSlots, computeSlotResourceProfile(numOfSlots, managedMemorySize));

	final TimerService<AllocationID> timerService = new TimerService<>(
		new ScheduledThreadPoolExecutor(1),
		taskManagerServicesConfiguration.getTimerServiceShutdownTimeout());

	final TaskSlotTable taskSlotTable = new TaskSlotTable(resourceProfiles, timerService);

	final JobManagerTable jobManagerTable = new JobManagerTable();

	final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation, taskManagerServicesConfiguration.getRetryingRegistrationConfiguration());

	final String[] stateRootDirectoryStrings = taskManagerServicesConfiguration.getLocalRecoveryStateRootDirectories();

	final File[] stateRootDirectoryFiles = new File[stateRootDirectoryStrings.length];

	for (int i = 0; i < stateRootDirectoryStrings.length; ++i) {
		stateRootDirectoryFiles[i] = new File(stateRootDirectoryStrings[i], LOCAL_STATE_SUB_DIRECTORY_ROOT);
	}

	final TaskExecutorLocalStateStoresManager taskStateManager = new TaskExecutorLocalStateStoresManager(
		taskManagerServicesConfiguration.isLocalRecoveryEnabled(),
		stateRootDirectoryFiles,
		taskIOExecutor);

	return new TaskManagerServices(
		taskManagerLocation,
		memoryManager,
		ioManager,
		shuffleEnvironment,
		kvStateService,
		broadcastVariableManager,
		taskSlotTable,
		jobManagerTable,
		jobLeaderService,
		taskStateManager,
		taskEventDispatcher);
}
 

private static Task createTask(
	StreamOperator<?> op,
	StateBackend backend,
	CheckpointResponder checkpointResponder) throws IOException {

	Configuration taskConfig = new Configuration();
	StreamConfig cfg = new StreamConfig(taskConfig);
	cfg.setStreamOperator(op);
	cfg.setOperatorID(new OperatorID());
	cfg.setStateBackend(backend);

	ExecutionConfig executionConfig = new ExecutionConfig();

	JobInformation jobInformation = new JobInformation(
			new JobID(),
			"test job name",
			new SerializedValue<>(executionConfig),
			new Configuration(),
			Collections.emptyList(),
			Collections.emptyList());

	TaskInformation taskInformation = new TaskInformation(
			new JobVertexID(),
			"test task name",
			1,
			11,
			TestStreamTask.class.getName(),
			taskConfig);

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

	return new Task(
			jobInformation,
			taskInformation,
			new ExecutionAttemptID(),
			new AllocationID(),
			0,
			0,
			Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
			Collections.<InputGateDeploymentDescriptor>emptyList(),
			0,
			mock(MemoryManager.class),
			mock(IOManager.class),
			shuffleEnvironment,
			new KvStateService(new KvStateRegistry(), null, null),
			mock(BroadcastVariableManager.class),
			new TaskEventDispatcher(),
			ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES,
			new TestTaskStateManager(),
			mock(TaskManagerActions.class),
			mock(InputSplitProvider.class),
			checkpointResponder,
			new NoOpTaskOperatorEventGateway(),
			new TestGlobalAggregateManager(),
			TestingClassLoaderLease.newBuilder().build(),
			new FileCache(new String[] { EnvironmentInformation.getTemporaryFileDirectory() },
				VoidPermanentBlobService.INSTANCE),
			new TestingTaskManagerRuntimeInfo(),
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
			new NoOpResultPartitionConsumableNotifier(),
			mock(PartitionProducerStateChecker.class),
			Executors.directExecutor());
}
 

/**
 * Tests request back off configuration is correctly forwarded to the channels.
 */
@Test
public void testRequestBackoffConfiguration() throws Exception {
	IntermediateResultPartitionID[] partitionIds = new IntermediateResultPartitionID[] {
		new IntermediateResultPartitionID(),
		new IntermediateResultPartitionID(),
		new IntermediateResultPartitionID()
	};

	ResourceID localLocation = ResourceID.generate();
	ShuffleDescriptor[] channelDescs = new ShuffleDescriptor[]{
		// Local
		createRemoteWithIdAndLocation(partitionIds[0], localLocation),
		// Remote
		createRemoteWithIdAndLocation(partitionIds[1], ResourceID.generate()),
		// Unknown
		new UnknownShuffleDescriptor(new ResultPartitionID(partitionIds[2], new ExecutionAttemptID()))};

	InputGateDeploymentDescriptor gateDesc = new InputGateDeploymentDescriptor(
		new IntermediateDataSetID(),
		ResultPartitionType.PIPELINED,
		0,
		channelDescs);

	int initialBackoff = 137;
	int maxBackoff = 1001;

	final NettyShuffleEnvironment netEnv = new NettyShuffleEnvironmentBuilder()
		.setPartitionRequestInitialBackoff(initialBackoff)
		.setPartitionRequestMaxBackoff(maxBackoff)
		.setIsCreditBased(enableCreditBasedFlowControl)
		.build();

	SingleInputGate gate = new SingleInputGateFactory(
		localLocation,
		netEnv.getConfiguration(),
		netEnv.getConnectionManager(),
		netEnv.getResultPartitionManager(),
		new TaskEventDispatcher(),
		netEnv.getNetworkBufferPool())
		.create(
			"TestTask",
			gateDesc,
			SingleInputGateBuilder.NO_OP_PRODUCER_CHECKER,
			InputChannelTestUtils.newUnregisteredInputChannelMetrics());

	try {
		assertEquals(gateDesc.getConsumedPartitionType(), gate.getConsumedPartitionType());

		Map<IntermediateResultPartitionID, InputChannel> channelMap = gate.getInputChannels();

		assertEquals(3, channelMap.size());
		InputChannel localChannel = channelMap.get(partitionIds[0]);
		assertEquals(LocalInputChannel.class, localChannel.getClass());

		InputChannel remoteChannel = channelMap.get(partitionIds[1]);
		assertEquals(RemoteInputChannel.class, remoteChannel.getClass());

		InputChannel unknownChannel = channelMap.get(partitionIds[2]);
		assertEquals(UnknownInputChannel.class, unknownChannel.getClass());

		InputChannel[] channels =
			new InputChannel[] {localChannel, remoteChannel, unknownChannel};
		for (InputChannel ch : channels) {
			assertEquals(0, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff * 2, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(initialBackoff * 2 * 2, ch.getCurrentBackoff());

			assertTrue(ch.increaseBackoff());
			assertEquals(maxBackoff, ch.getCurrentBackoff());

			assertFalse(ch.increaseBackoff());
		}
	} finally {
		gate.close();
		netEnv.close();
	}
}
 

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

@Test
public void testDuplicateCancel() throws Exception {

	NettyServerAndClient serverAndClient = null;

	try {
		final TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);

		ResultPartitionManager partitions = mock(ResultPartitionManager.class);

		ResultPartitionID pid = new ResultPartitionID();

		final CountDownLatch sync = new CountDownLatch(1);

		final ResultSubpartitionView view = spy(new InfiniteSubpartitionView(outboundBuffers, sync));

		// Return infinite subpartition
		when(partitions.createSubpartitionView(eq(pid), eq(0), any(BufferAvailabilityListener.class)))
				.thenAnswer(new Answer<ResultSubpartitionView>() {
					@Override
					public ResultSubpartitionView answer(InvocationOnMock invocationOnMock) throws Throwable {
						BufferAvailabilityListener listener = (BufferAvailabilityListener) invocationOnMock.getArguments()[2];
						listener.notifyDataAvailable();
						return view;
					}
				});

		NettyProtocol protocol = new NettyProtocol(
				partitions, mock(TaskEventDispatcher.class), true);

		serverAndClient = initServerAndClient(protocol);

		Channel ch = connect(serverAndClient);

		// Request for non-existing input channel => results in cancel request
		InputChannelID inputChannelId = new InputChannelID();

		ch.writeAndFlush(new PartitionRequest(pid, 0, inputChannelId, Integer.MAX_VALUE)).await();

		// Wait for the notification
		if (!sync.await(TestingUtils.TESTING_DURATION().toMillis(), TimeUnit.MILLISECONDS)) {
			fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION().toMillis() +
					" ms to be notified about cancelled partition.");
		}

		ch.writeAndFlush(new CancelPartitionRequest(inputChannelId)).await();

		ch.close();

		NettyTestUtil.awaitClose(ch);

		verify(view, times(1)).releaseAllResources();
		verify(view, times(1)).notifySubpartitionConsumed();
	}
	finally {
		shutdown(serverAndClient);
	}
}
 

private NettyConnectionManager createNettyConnectionManager(NettyConfig config) {
	return new NettyConnectionManager(new ResultPartitionManager(), new TaskEventDispatcher(), config, true);
}
 

/**
 * Verifies that failed client requests via {@link PartitionRequestClient} are correctly
 * attributed to the respective {@link RemoteInputChannel}.
 */
@Test
public void testExceptionOnWrite() throws Exception {

	NettyProtocol protocol = new NettyProtocol(
			mock(ResultPartitionProvider.class),
			mock(TaskEventDispatcher.class),
			true) {

		@Override
		public ChannelHandler[] getServerChannelHandlers() {
			return new ChannelHandler[0];
		}
	};

	// We need a real server and client in this test, because Netty's EmbeddedChannel is
	// not failing the ChannelPromise of failed writes.
	NettyServerAndClient serverAndClient = initServerAndClient(protocol, createConfig());

	Channel ch = connect(serverAndClient);

	NetworkClientHandler handler = getClientHandler(ch);

	// Last outbound handler throws Exception after 1st write
	ch.pipeline().addFirst(new ChannelOutboundHandlerAdapter() {
		int writeNum = 0;

		@Override
		public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
				throws Exception {

			if (writeNum >= 1) {
				throw new RuntimeException("Expected test exception.");
			}

			writeNum++;
			ctx.write(msg, promise);
		}
	});

	PartitionRequestClient requestClient = new NettyPartitionRequestClient(
			ch, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	// Create input channels
	RemoteInputChannel[] rich = new RemoteInputChannel[] {
			createRemoteInputChannel(), createRemoteInputChannel()};

	final CountDownLatch sync = new CountDownLatch(1);

	// Do this with explicit synchronization. Otherwise this is not robust against slow timings
	// of the callback (e.g. we cannot just verify that it was called once, because there is
	// a chance that we do this too early).
	doAnswer(new Answer<Void>() {
		@Override
		public Void answer(InvocationOnMock invocation) throws Throwable {
			sync.countDown();
			return null;
		}
	}).when(rich[1]).onError(isA(LocalTransportException.class));

	// First request is successful
	requestClient.requestSubpartition(new ResultPartitionID(), 0, rich[0], 0);

	// Second request is *not* successful
	requestClient.requestSubpartition(new ResultPartitionID(), 0, rich[1], 0);

	// Wait for the notification and it could confirm all the request operations are done
	if (!sync.await(TestingUtils.TESTING_DURATION().toMillis(), TimeUnit.MILLISECONDS)) {
		fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION().toMillis() +
				" ms to be notified about the channel error.");
	}

	// Only the second channel should be notified about the error
	verify(rich[0], times(0)).onError(any(LocalTransportException.class));

	shutdown(serverAndClient);
}