org.apache.flink.runtime.io.network.partition.BufferAvailabilityListener Java Examples

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {
	ResultSubpartitionView view = mock(ResultSubpartitionView.class);
	when(view.isReleased()).thenReturn(true);
	when(view.getFailureCause()).thenReturn(new Exception("Expected test exception"));

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager
			.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
			.thenReturn(view);

	SingleInputGate inputGate = mock(SingleInputGate.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(inputGate.getBufferProvider()).thenReturn(bufferProvider);

	LocalInputChannel ch = createLocalInputChannel(inputGate, partitionManager);

	ch.requestSubpartition(0);

	// Should throw an instance of CancelTaskException.
	ch.getNextBuffer();
}
 

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

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {
	ResultSubpartitionView view = mock(ResultSubpartitionView.class);
	when(view.isReleased()).thenReturn(true);
	when(view.getFailureCause()).thenReturn(new Exception("Expected test exception"));

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager
			.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
			.thenReturn(view);

	SingleInputGate inputGate = mock(SingleInputGate.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(inputGate.getBufferProvider()).thenReturn(bufferProvider);

	LocalInputChannel ch = createLocalInputChannel(
			inputGate, partitionManager, new Tuple2<>(0, 0));

	ch.requestSubpartition(0);

	// Should throw an instance of CancelTaskException.
	ch.getNextBuffer();
}
 

@Test(expected = CancelTaskException.class)
public void testProducerFailedException() throws Exception {
	ResultSubpartitionView view = mock(ResultSubpartitionView.class);
	when(view.isReleased()).thenReturn(true);
	when(view.getFailureCause()).thenReturn(new Exception("Expected test exception"));

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager
			.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class)))
			.thenReturn(view);

	SingleInputGate inputGate = mock(SingleInputGate.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(inputGate.getBufferProvider()).thenReturn(bufferProvider);

	LocalInputChannel ch = createLocalInputChannel(inputGate, partitionManager);

	ch.requestSubpartition(0);

	// Should throw an instance of CancelTaskException.
	ch.getNextBuffer();
}
 

/**
 * 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 = InputChannelBuilder.newBuilder()
		.setPartitionManager(partitionManager)
		.buildUnknownAndSetToGate(inputGate);

	// Update to a local channel and verify that no request is triggered
	ResultPartitionID resultPartitionID = unknown.getPartitionId();
	ResourceID location = ResourceID.generate();
	inputGate.updateInputChannel(location, createRemoteWithIdAndLocation(resultPartitionID.getPartitionId(), location));

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

@Override
public ResultSubpartitionView createSubpartitionView(
		ResultPartitionID partitionId,
		int subpartitionIndex,
		BufferAvailabilityListener availabilityListener) throws IOException {
	++counter;
	return subpartitionView;
}
 

/**
 * Tests {@link PartitionRequestQueue#enqueueAvailableReader(NetworkSequenceViewReader)},
 * verifying the reader would be enqueued in the pipeline after resuming data consumption if there
 * are credit and data available.
 */
@Test
public void testEnqueueReaderByResumingConsumption() throws Exception {
	PipelinedSubpartition subpartition = PipelinedSubpartitionTest.createPipelinedSubpartition();
	Buffer.DataType dataType1 = Buffer.DataType.ALIGNED_EXACTLY_ONCE_CHECKPOINT_BARRIER;
	Buffer.DataType dataType2 = Buffer.DataType.DATA_BUFFER;
	subpartition.add(createEventBufferConsumer(4096, dataType1));
	subpartition.add(createEventBufferConsumer(4096, dataType2));

	BufferAvailabilityListener bufferAvailabilityListener = new NoOpBufferAvailablityListener();
	PipelinedSubpartitionView view = subpartition.createReadView(bufferAvailabilityListener);
	ResultPartitionProvider partitionProvider = (partitionId, index, availabilityListener) -> view;

	InputChannelID receiverId = new InputChannelID();
	PartitionRequestQueue queue = new PartitionRequestQueue();
	CreditBasedSequenceNumberingViewReader reader = new CreditBasedSequenceNumberingViewReader(receiverId, 0, queue);
	EmbeddedChannel channel = new EmbeddedChannel(queue);

	reader.requestSubpartitionView(partitionProvider, new ResultPartitionID(), 0);
	queue.notifyReaderCreated(reader);
	// we have adequate credits
	reader.addCredit(Integer.MAX_VALUE);
	assertTrue(reader.isAvailable());

	reader.notifyDataAvailable();
	channel.runPendingTasks();
	assertFalse(reader.isAvailable());
	assertEquals(1, subpartition.unsynchronizedGetNumberOfQueuedBuffers());

	queue.addCreditOrResumeConsumption(receiverId, NetworkSequenceViewReader::resumeConsumption);
	assertFalse(reader.isAvailable());
	assertEquals(0, subpartition.unsynchronizedGetNumberOfQueuedBuffers());

	Object data1 = channel.readOutbound();
	assertEquals(dataType1, ((NettyMessage.BufferResponse) data1).buffer.getDataType());
	Object data2 = channel.readOutbound();
	assertEquals(dataType2, ((NettyMessage.BufferResponse) data2).buffer.getDataType());
}
 

/**
 * Tests that reading from a channel when after the partition has been
 * released are handled and don't lead to NPEs.
 */
@Test
public void testGetNextAfterPartitionReleased() throws Exception {
	ResultSubpartitionView reader = mock(ResultSubpartitionView.class);
	SingleInputGate gate = mock(SingleInputGate.class);
	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	when(partitionManager.createSubpartitionView(
		any(ResultPartitionID.class),
		anyInt(),
		any(BufferAvailabilityListener.class))).thenReturn(reader);

	LocalInputChannel channel = createLocalInputChannel(gate, partitionManager);

	channel.requestSubpartition(0);

	// Null buffer but not released
	when(reader.getNextBuffer()).thenReturn(null);
	when(reader.isReleased()).thenReturn(false);

	assertFalse(channel.getNextBuffer().isPresent());

	// Null buffer and released
	when(reader.getNextBuffer()).thenReturn(null);
	when(reader.isReleased()).thenReturn(true);

	try {
		channel.getNextBuffer();
		fail("Did not throw expected CancelTaskException");
	} catch (CancelTaskException ignored) {
	}

	channel.releaseAllResources();
	assertFalse(channel.getNextBuffer().isPresent());
}
 

@Test
public void testBackwardsEventWithUninitializedChannel() throws Exception {
	// Setup environment
	final TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);
	when(taskEventDispatcher.publish(any(ResultPartitionID.class), any(TaskEvent.class))).thenReturn(true);

	final ResultSubpartitionView iterator = mock(ResultSubpartitionView.class);
	when(iterator.getNextBuffer()).thenReturn(
		new BufferAndBacklog(new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(1024), FreeingBufferRecycler.INSTANCE), false, 0, false));

	final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager.createSubpartitionView(
		any(ResultPartitionID.class),
		anyInt(),
		any(BufferAvailabilityListener.class))).thenReturn(iterator);

	// Setup reader with one local and one unknown input channel

	NettyShuffleEnvironment environment = createNettyShuffleEnvironment();
	final SingleInputGate inputGate = createInputGate(environment, 2, ResultPartitionType.PIPELINED);
	try {
		// Local
		ResultPartitionID localPartitionId = new ResultPartitionID();

		InputChannelBuilder.newBuilder()
			.setPartitionId(localPartitionId)
			.setPartitionManager(partitionManager)
			.setTaskEventPublisher(taskEventDispatcher)
			.buildLocalAndSetToGate(inputGate);

		// Unknown
		ResultPartitionID unknownPartitionId = new ResultPartitionID();

		InputChannelBuilder.newBuilder()
			.setChannelIndex(1)
			.setPartitionId(unknownPartitionId)
			.setPartitionManager(partitionManager)
			.setTaskEventPublisher(taskEventDispatcher)
			.buildUnknownAndSetToGate(inputGate);

		inputGate.setup();

		// Only the local channel can request
		verify(partitionManager, times(1)).createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));

		// Send event backwards and initialize unknown channel afterwards
		final TaskEvent event = new TestTaskEvent();
		inputGate.sendTaskEvent(event);

		// Only the local channel can send out the event
		verify(taskEventDispatcher, times(1)).publish(any(ResultPartitionID.class), any(TaskEvent.class));

		// After the update, the pending event should be send to local channel

		ResourceID location = ResourceID.generate();
		inputGate.updateInputChannel(location, createRemoteWithIdAndLocation(unknownPartitionId.getPartitionId(), location));

		verify(partitionManager, times(2)).createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));
		verify(taskEventDispatcher, times(2)).publish(any(ResultPartitionID.class), any(TaskEvent.class));
	}
	finally {
		inputGate.close();
		environment.close();
	}
}
 

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

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

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

/**
 * 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), true);

		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();
		verify(view, times(0)).notifySubpartitionConsumed();
	}
	finally {
		shutdown(serverAndClient);
	}
}
 

/**
 * 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), true) {

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

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

/**
 * 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), true);

		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();
		verify(view, times(1)).notifySubpartitionConsumed();
	}
	finally {
		shutdown(serverAndClient);
	}
}
 

/**
 * Tests that reading from a channel when after the partition has been
 * released are handled and don't lead to NPEs.
 */
@Test
public void testGetNextAfterPartitionReleased() throws Exception {
	ResultSubpartitionView reader = mock(ResultSubpartitionView.class);
	SingleInputGate gate = mock(SingleInputGate.class);
	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	when(partitionManager.createSubpartitionView(
		any(ResultPartitionID.class),
		anyInt(),
		any(BufferAvailabilityListener.class))).thenReturn(reader);

	LocalInputChannel channel = new LocalInputChannel(
		gate,
		0,
		new ResultPartitionID(),
		partitionManager,
		new TaskEventDispatcher(),
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	channel.requestSubpartition(0);

	// Null buffer but not released
	when(reader.getNextBuffer()).thenReturn(null);
	when(reader.isReleased()).thenReturn(false);

	assertFalse(channel.getNextBuffer().isPresent());

	// Null buffer and released
	when(reader.getNextBuffer()).thenReturn(null);
	when(reader.isReleased()).thenReturn(true);

	try {
		channel.getNextBuffer();
		fail("Did not throw expected CancelTaskException");
	} catch (CancelTaskException ignored) {
	}

	channel.releaseAllResources();
	assertFalse(channel.getNextBuffer().isPresent());
}
 

@Test
public void testBackwardsEventWithUninitializedChannel() throws Exception {
	// Setup environment
	final TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);
	when(taskEventDispatcher.publish(any(ResultPartitionID.class), any(TaskEvent.class))).thenReturn(true);

	final ResultSubpartitionView iterator = mock(ResultSubpartitionView.class);
	when(iterator.getNextBuffer()).thenReturn(
		new BufferAndBacklog(new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(1024), FreeingBufferRecycler.INSTANCE), false, 0, false));

	final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
	when(partitionManager.createSubpartitionView(
		any(ResultPartitionID.class),
		anyInt(),
		any(BufferAvailabilityListener.class))).thenReturn(iterator);

	// Setup reader with one local and one unknown input channel

	final SingleInputGate inputGate = createInputGate();
	final BufferPool bufferPool = mock(BufferPool.class);
	when(bufferPool.getNumberOfRequiredMemorySegments()).thenReturn(2);

	inputGate.setBufferPool(bufferPool);

	// Local
	ResultPartitionID localPartitionId = new ResultPartitionID(new IntermediateResultPartitionID(), new ExecutionAttemptID());

	InputChannel local = new LocalInputChannel(inputGate, 0, localPartitionId, partitionManager, taskEventDispatcher, UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	// Unknown
	ResultPartitionID unknownPartitionId = new ResultPartitionID(new IntermediateResultPartitionID(), new ExecutionAttemptID());

	InputChannel unknown = new UnknownInputChannel(inputGate, 1, unknownPartitionId, partitionManager, taskEventDispatcher, mock(ConnectionManager.class), 0, 0, UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	// Set channels
	inputGate.setInputChannel(localPartitionId.getPartitionId(), local);
	inputGate.setInputChannel(unknownPartitionId.getPartitionId(), unknown);

	// Request partitions
	inputGate.requestPartitions();

	// Only the local channel can request
	verify(partitionManager, times(1)).createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));

	// Send event backwards and initialize unknown channel afterwards
	final TaskEvent event = new TestTaskEvent();
	inputGate.sendTaskEvent(event);

	// Only the local channel can send out the event
	verify(taskEventDispatcher, times(1)).publish(any(ResultPartitionID.class), any(TaskEvent.class));

	// After the update, the pending event should be send to local channel
	inputGate.updateInputChannel(new InputChannelDeploymentDescriptor(new ResultPartitionID(unknownPartitionId.getPartitionId(), unknownPartitionId.getProducerId()), ResultPartitionLocation.createLocal()));

	verify(partitionManager, times(2)).createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));
	verify(taskEventDispatcher, times(2)).publish(any(ResultPartitionID.class), any(TaskEvent.class));
}
 

/**
 * 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), true) {

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

@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(0)).notifySubpartitionConsumed();
	}
	finally {
		shutdown(serverAndClient);
	}
}