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

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

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

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

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

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

@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 EmbeddedChannel createEmbeddedChannel() {
	NettyProtocol protocol = new NettyProtocol(
			mock(ResultPartitionProvider.class),
			mock(TaskEventDispatcher.class),
			true);

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

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

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

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

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

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

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

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

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

public RuntimeEnvironment(
		JobID jobId,
		JobVertexID jobVertexId,
		ExecutionAttemptID executionId,
		ExecutionConfig executionConfig,
		TaskInfo taskInfo,
		Configuration jobConfiguration,
		Configuration taskConfiguration,
		ClassLoader userCodeClassLoader,
		MemoryManager memManager,
		IOManager ioManager,
		BroadcastVariableManager bcVarManager,
		TaskStateManager taskStateManager,
		GlobalAggregateManager aggregateManager,
		AccumulatorRegistry accumulatorRegistry,
		TaskKvStateRegistry kvStateRegistry,
		InputSplitProvider splitProvider,
		Map<String, Future<Path>> distCacheEntries,
		ResultPartitionWriter[] writers,
		InputGate[] inputGates,
		TaskEventDispatcher taskEventDispatcher,
		CheckpointResponder checkpointResponder,
		TaskManagerRuntimeInfo taskManagerInfo,
		TaskMetricGroup metrics,
		Task containingTask) {

	this.jobId = checkNotNull(jobId);
	this.jobVertexId = checkNotNull(jobVertexId);
	this.executionId = checkNotNull(executionId);
	this.taskInfo = checkNotNull(taskInfo);
	this.executionConfig = checkNotNull(executionConfig);
	this.jobConfiguration = checkNotNull(jobConfiguration);
	this.taskConfiguration = checkNotNull(taskConfiguration);
	this.userCodeClassLoader = checkNotNull(userCodeClassLoader);
	this.memManager = checkNotNull(memManager);
	this.ioManager = checkNotNull(ioManager);
	this.bcVarManager = checkNotNull(bcVarManager);
	this.taskStateManager = checkNotNull(taskStateManager);
	this.aggregateManager = checkNotNull(aggregateManager);
	this.accumulatorRegistry = checkNotNull(accumulatorRegistry);
	this.kvStateRegistry = checkNotNull(kvStateRegistry);
	this.splitProvider = checkNotNull(splitProvider);
	this.distCacheEntries = checkNotNull(distCacheEntries);
	this.writers = checkNotNull(writers);
	this.inputGates = checkNotNull(inputGates);
	this.taskEventDispatcher = checkNotNull(taskEventDispatcher);
	this.checkpointResponder = checkNotNull(checkpointResponder);
	this.taskManagerInfo = checkNotNull(taskManagerInfo);
	this.containingTask = containingTask;
	this.metrics = metrics;
}
 

@Override
public TaskEventDispatcher getTaskEventDispatcher() {
	return taskEventDispatcher;
}
 

NettyProtocol(ResultPartitionProvider partitionProvider, TaskEventDispatcher taskEventDispatcher, boolean creditBasedEnabled) {
	this.partitionProvider = partitionProvider;
	this.taskEventDispatcher = taskEventDispatcher;
	this.creditBasedEnabled = creditBasedEnabled;
}
 

/**
 * Verifies remote closes trigger the release of all resources.
 */
@Test
public void testRemoteClose() throws Exception {
	final TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);

	final CountDownLatch sync = new CountDownLatch(1);

	final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	when(partitionManager
		.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
		.thenAnswer(new Answer<ResultSubpartitionView>() {
			@Override
			public ResultSubpartitionView answer(InvocationOnMock invocationOnMock) throws Throwable {
				BufferAvailabilityListener listener = (BufferAvailabilityListener) invocationOnMock.getArguments()[2];
				listener.notifyDataAvailable();
				return new CancelPartitionRequestTest.InfiniteSubpartitionView(outboundBuffers, sync);
			}
		});

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

		@Override
		public ChannelHandler[] getClientChannelHandlers() {
			return new ChannelHandler[]{
				new NettyMessage.NettyMessageEncoder(),
				// Close on read
				new ChannelInboundHandlerAdapter() {
					@Override
					public void channelRead(ChannelHandlerContext ctx, Object msg)
						throws Exception {

						ctx.channel().close();
					}
				}
			};
		}
	};

	NettyTestUtil.NettyServerAndClient serverAndClient = null;

	try {
		serverAndClient = initServerAndClient(protocol, createConfig());

		Channel ch = connect(serverAndClient);

		// Write something to trigger close by server
		ch.writeAndFlush(new NettyMessage.PartitionRequest(new ResultPartitionID(), 0, new InputChannelID(), Integer.MAX_VALUE));

		// 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 released partition.");
		}
	} finally {
		shutdown(serverAndClient);
	}
}
 

/**
 * Verifies that requests for non-existing (failed/cancelled) input channels are properly
 * cancelled. The receiver receives data, but there is no input channel to receive the data.
 * This should cancel the request.
 */
@Test
public void testCancelPartitionRequest() throws Exception {

	NettyServerAndClient serverAndClient = null;

	try {
		TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);

		ResultPartitionManager partitions = mock(ResultPartitionManager.class);

		ResultPartitionID pid = new ResultPartitionID();

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

		serverAndClient = initServerAndClient(protocol);

		Channel ch = connect(serverAndClient);

		// Request for non-existing input channel => results in cancel request
		ch.writeAndFlush(new PartitionRequest(pid, 0, new 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.");
		}

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

@Override
public TaskEventDispatcher getTaskEventDispatcher() {
	return taskEventDispatcher;
}
 

@Override
public TaskEventDispatcher getTaskEventDispatcher() {
	throw new UnsupportedOperationException();
}