org.apache.flink.runtime.io.network.partition.consumer.InputChannelID Java Examples

private BufferResponse(
		@Nullable Buffer buffer,
		Buffer.DataType dataType,
		boolean isCompressed,
		int sequenceNumber,
		InputChannelID receiverId,
		int backlog,
		int bufferSize) {
	this.buffer = buffer;
	this.dataType = dataType;
	this.isCompressed = isCompressed;
	this.sequenceNumber = sequenceNumber;
	this.receiverId = checkNotNull(receiverId);
	this.backlog = backlog;
	this.bufferSize = bufferSize;
}
 

static ErrorResponse readFrom(ByteBuf buffer) throws Exception {
	try (ObjectInputStream ois = new ObjectInputStream(new ByteBufInputStream(buffer))) {
		Object obj = ois.readObject();

		if (!(obj instanceof Throwable)) {
			throw new ClassCastException("Read object expected to be of type Throwable, " +
					"actual type is " + obj.getClass() + ".");
		} else {
			if (buffer.readBoolean()) {
				InputChannelID receiverId = InputChannelID.fromByteBuf(buffer);
				return new ErrorResponse((Throwable) obj, receiverId);
			} else {
				return new ErrorResponse((Throwable) obj);
			}
		}
	}
}
 

/**
 * Tests a fix for FLINK-1627.
 *
 * <p> FLINK-1627 discovered a race condition, which could lead to an infinite loop when a
 * receiver was cancelled during a certain time of decoding a message. The test reproduces the
 * input, which lead to the infinite loop: when the handler gets a reference to the buffer
 * provider of the receiving input channel, but the respective input channel is released (and
 * the corresponding buffer provider destroyed), the handler did not notice this.
 *
 * @see <a href="https://issues.apache.org/jira/browse/FLINK-1627">FLINK-1627</a>
 */
@Test(timeout = 60000)
@SuppressWarnings("unchecked")
public void testReleaseInputChannelDuringDecode() throws Exception {
	// Mocks an input channel in a state as it was released during a decode.
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(null);
	when(bufferProvider.isDestroyed()).thenReturn(true);
	when(bufferProvider.addBufferListener(any(BufferListener.class))).thenReturn(false);

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	final BufferResponse receivedBuffer = createBufferResponse(
			TestBufferFactory.createBuffer(TestBufferFactory.BUFFER_SIZE), 0, inputChannel.getInputChannelId(), 2);

	final PartitionRequestClientHandler client = new PartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	client.channelRead(mock(ChannelHandlerContext.class), receivedBuffer);
}
 

@Test
public void testProducerFailedException() throws Exception {
	PartitionRequestQueue queue = new PartitionRequestQueue();

	ResultSubpartitionView view = new ReleasedResultSubpartitionView();

	ResultPartitionProvider partitionProvider =
		(partitionId, index, availabilityListener) -> view;

	EmbeddedChannel ch = new EmbeddedChannel(queue);

	CreditBasedSequenceNumberingViewReader seqView = new CreditBasedSequenceNumberingViewReader(new InputChannelID(), 2, queue);
	seqView.requestSubpartitionView(partitionProvider, new ResultPartitionID(), 0);
	// Add available buffer to trigger enqueue the erroneous view
	seqView.notifyDataAvailable();

	ch.runPendingTasks();

	// Read the enqueued msg
	Object msg = ch.readOutbound();

	assertEquals(msg.getClass(), NettyMessage.ErrorResponse.class);

	NettyMessage.ErrorResponse err = (NettyMessage.ErrorResponse) msg;
	assertTrue(err.cause instanceof CancelTaskException);
}
 

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

@Override
public void userEventTriggered(ChannelHandlerContext ctx, Object msg) throws Exception {
	// The user event triggered event loop callback is used for thread-safe
	// hand over of reader queues and cancelled producers.

	if (msg instanceof NetworkSequenceViewReader) {
		enqueueAvailableReader((NetworkSequenceViewReader) msg);
	} else if (msg.getClass() == InputChannelID.class) {
		// Release partition view that get a cancel request.
		InputChannelID toCancel = (InputChannelID) msg;

		// remove reader from queue of available readers
		availableReaders.removeIf(reader -> reader.getReceiverId().equals(toCancel));

		// remove reader from queue of all readers and release its resource
		final NetworkSequenceViewReader toRelease = allReaders.remove(toCancel);
		if (toRelease != null) {
			releaseViewReader(toRelease);
		}
	} else {
		ctx.fireUserEventTriggered(msg);
	}
}
 

/**
 * Returns a deserialized buffer message as it would be received during runtime.
 */
static BufferResponse createBufferResponse(
		Buffer buffer,
		int sequenceNumber,
		InputChannelID receivingChannelId,
		int backlog) throws IOException {

	// Mock buffer to serialize
	BufferResponse resp = new BufferResponse(buffer, sequenceNumber, receivingChannelId, backlog);

	ByteBuf serialized = resp.write(UnpooledByteBufAllocator.DEFAULT);

	// Skip general header bytes
	serialized.readBytes(NettyMessage.FRAME_HEADER_LENGTH);

	// Deserialize the bytes again. We have to go this way, because we only partly deserialize
	// the header of the response and wait for a buffer from the buffer pool to copy the payload
	// data into.
	BufferResponse deserialized = BufferResponse.readFrom(serialized);

	return deserialized;
}
 

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

static ErrorResponse readFrom(ByteBuf buffer) throws Exception {
	try (ObjectInputStream ois = new ObjectInputStream(new ByteBufInputStream(buffer))) {
		Object obj = ois.readObject();

		if (!(obj instanceof Throwable)) {
			throw new ClassCastException("Read object expected to be of type Throwable, " +
					"actual type is " + obj.getClass() + ".");
		} else {
			if (buffer.readBoolean()) {
				InputChannelID receiverId = InputChannelID.fromByteBuf(buffer);
				return new ErrorResponse((Throwable) obj, receiverId);
			} else {
				return new ErrorResponse((Throwable) obj);
			}
		}
	}
}
 

private void testBufferWriting(ResultSubpartitionView view) throws IOException {
	// setup
	ResultPartitionProvider partitionProvider =
		(partitionId, index, availabilityListener) -> view;

	final InputChannelID receiverId = new InputChannelID();
	final PartitionRequestQueue queue = new PartitionRequestQueue();
	final SequenceNumberingViewReader reader = new SequenceNumberingViewReader(receiverId, queue);
	final EmbeddedChannel channel = new EmbeddedChannel(queue);

	reader.requestSubpartitionView(partitionProvider, new ResultPartitionID(), 0);

	// notify about buffer availability and encode one buffer
	reader.notifyDataAvailable();

	channel.runPendingTasks();

	Object read = channel.readOutbound();
	assertNotNull(read);
	if (read instanceof NettyMessage.ErrorResponse) {
		((NettyMessage.ErrorResponse) read).cause.printStackTrace();
	}
	assertThat(read, instanceOf(NettyMessage.BufferResponse.class));
	read = channel.readOutbound();
	assertNull(read);
}
 

/**
 * Tests a fix for FLINK-1627.
 *
 * <p> FLINK-1627 discovered a race condition, which could lead to an infinite loop when a
 * receiver was cancelled during a certain time of decoding a message. The test reproduces the
 * input, which lead to the infinite loop: when the handler gets a reference to the buffer
 * provider of the receiving input channel, but the respective input channel is released (and
 * the corresponding buffer provider destroyed), the handler did not notice this.
 *
 * @see <a href="https://issues.apache.org/jira/browse/FLINK-1627">FLINK-1627</a>
 */
@Test(timeout = 60000)
@SuppressWarnings("unchecked")
public void testReleaseInputChannelDuringDecode() throws Exception {
	// Mocks an input channel in a state as it was released during a decode.
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(null);
	when(bufferProvider.isDestroyed()).thenReturn(true);
	when(bufferProvider.addBufferListener(any(BufferListener.class))).thenReturn(false);

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	final BufferResponse receivedBuffer = createBufferResponse(
			TestBufferFactory.createBuffer(TestBufferFactory.BUFFER_SIZE), 0, inputChannel.getInputChannelId(), 2);

	final CreditBasedPartitionRequestClientHandler client = new CreditBasedPartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	client.channelRead(mock(ChannelHandlerContext.class), receivedBuffer);
}
 

@Override
public void userEventTriggered(ChannelHandlerContext ctx, Object msg) throws Exception {
	// The user event triggered event loop callback is used for thread-safe
	// hand over of reader queues and cancelled producers.

	if (msg instanceof NetworkSequenceViewReader) {
		enqueueAvailableReader((NetworkSequenceViewReader) msg);
	} else if (msg.getClass() == InputChannelID.class) {
		// Release partition view that get a cancel request.
		InputChannelID toCancel = (InputChannelID) msg;

		// remove reader from queue of available readers
		availableReaders.removeIf(reader -> reader.getReceiverId().equals(toCancel));

		// remove reader from queue of all readers and release its resource
		final NetworkSequenceViewReader toRelease = allReaders.remove(toCancel);
		if (toRelease != null) {
			releaseViewReader(toRelease);
		}
	} else {
		ctx.fireUserEventTriggered(msg);
	}
}
 

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

@Override
public void cancelRequestFor(InputChannelID inputChannelId) {
	if (inputChannelId == null || ctx == null) {
		return;
	}

	if (cancelled.putIfAbsent(inputChannelId, inputChannelId) == null) {
		ctx.writeAndFlush(new NettyMessage.CancelPartitionRequest(inputChannelId));
	}
}
 

AddCredit(ResultPartitionID partitionId, int credit, InputChannelID receiverId) {
	checkArgument(credit > 0, "The announced credit should be greater than 0");

	this.partitionId = partitionId;
	this.credit = credit;
	this.receiverId = receiverId;
}
 

private void testBufferWriting(ResultSubpartitionView view) throws IOException {
	// setup
	ResultPartitionProvider partitionProvider =
		(partitionId, index, availabilityListener) -> view;

	final InputChannelID receiverId = new InputChannelID();
	final PartitionRequestQueue queue = new PartitionRequestQueue();
	final CreditBasedSequenceNumberingViewReader reader = new CreditBasedSequenceNumberingViewReader(
		receiverId,
		Integer.MAX_VALUE,
		queue);
	final EmbeddedChannel channel = new EmbeddedChannel(queue);

	reader.requestSubpartitionView(partitionProvider, new ResultPartitionID(), 0);

	// notify about buffer availability and encode one buffer
	reader.notifyDataAvailable();

	channel.runPendingTasks();

	Object read = channel.readOutbound();
	assertNotNull(read);
	if (read instanceof NettyMessage.ErrorResponse) {
		((NettyMessage.ErrorResponse) read).cause.printStackTrace();
	}
	assertThat(read, instanceOf(NettyMessage.BufferResponse.class));
	read = channel.readOutbound();
	assertNull(read);
}
 

@Override
public void cancelRequestFor(InputChannelID inputChannelId) {
	if (inputChannelId == null || ctx == null) {
		return;
	}

	if (cancelled.putIfAbsent(inputChannelId, inputChannelId) == null) {
		ctx.writeAndFlush(new NettyMessage.CancelPartitionRequest(inputChannelId));
	}
}
 

static AddCredit readFrom(ByteBuf buffer) {
	ResultPartitionID partitionId =
		new ResultPartitionID(
			IntermediateResultPartitionID.fromByteBuf(buffer),
			ExecutionAttemptID.fromByteBuf(buffer));
	int credit = buffer.readInt();
	InputChannelID receiverId = InputChannelID.fromByteBuf(buffer);

	return new AddCredit(partitionId, credit, receiverId);
}
 

CreditBasedSequenceNumberingViewReader(
		InputChannelID receiverId,
		int initialCredit,
		PartitionRequestQueue requestQueue) {

	this.receiverId = receiverId;
	this.numCreditsAvailable = initialCredit;
	this.requestQueue = requestQueue;
}
 

/**
 * Allocates a pooled network buffer for the specific input channel.
 *
 * @param receiverId The id of the requested input channel.
 * @return The pooled network buffer.
 */
@Nullable
Buffer allocatePooledNetworkBuffer(InputChannelID receiverId) {
	Buffer buffer = null;

	RemoteInputChannel inputChannel = networkClientHandler.getInputChannel(receiverId);

	// If the input channel has been released, we cannot allocate buffer and the received message
	// will be discarded.
	if (inputChannel != null) {
		buffer = inputChannel.requestBuffer();
	}

	return buffer;
}
 

/**
 * Verifies that {@link RemoteInputChannel#onFailedPartitionRequest()} is called when a
 * {@link PartitionNotFoundException} is received.
 */
@Test
public void testReceivePartitionNotFoundException() throws Exception {
	// Minimal mock of a remote input channel
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(TestBufferFactory.createBuffer(0));

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	final ErrorResponse partitionNotFound = new ErrorResponse(
		new PartitionNotFoundException(new ResultPartitionID()),
		inputChannel.getInputChannelId());

	final CreditBasedPartitionRequestClientHandler client = new CreditBasedPartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	// Mock channel context
	ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
	when(ctx.channel()).thenReturn(mock(Channel.class));

	client.channelActive(ctx);

	client.channelRead(ctx, partitionNotFound);

	verify(inputChannel, times(1)).onFailedPartitionRequest();
}
 

/**
 * Tests a fix for FLINK-1761.
 *
 * <p>FLINK-1761 discovered an IndexOutOfBoundsException, when receiving buffers of size 0.
 */
@Test
public void testReceiveEmptyBuffer() throws Exception {
	// Minimal mock of a remote input channel
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(TestBufferFactory.createBuffer(0));

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	// An empty buffer of size 0
	final Buffer emptyBuffer = TestBufferFactory.createBuffer(0);

	final int backlog = 2;
	final BufferResponse receivedBuffer = createBufferResponse(
		emptyBuffer, 0, inputChannel.getInputChannelId(), backlog);

	final CreditBasedPartitionRequestClientHandler client = new CreditBasedPartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	// Read the empty buffer
	client.channelRead(mock(ChannelHandlerContext.class), receivedBuffer);

	// This should not throw an exception
	verify(inputChannel, never()).onError(any(Throwable.class));
	verify(inputChannel, times(1)).onEmptyBuffer(0, backlog);
}
 

private void testCancelPartitionRequest(boolean isAvailableView) throws Exception {
	// setup
	final ResultPartitionManager partitionManager = new ResultPartitionManager();
	final ResultPartition partition = createFinishedPartitionWithFilledData(partitionManager);
	final InputChannelID receiverId = new InputChannelID();
	final PartitionRequestQueue queue = new PartitionRequestQueue();
	final CreditBasedSequenceNumberingViewReader reader = new CreditBasedSequenceNumberingViewReader(receiverId, 0, queue);
	final EmbeddedChannel channel = new EmbeddedChannel(queue);

	reader.requestSubpartitionView(partitionManager, partition.getPartitionId(), 0);
	// add this reader into allReaders queue
	queue.notifyReaderCreated(reader);

	// block the channel so that we see an intermediate state in the test
	blockChannel(channel);

	// add credit to make this reader available for adding into availableReaders queue
	if (isAvailableView) {
		queue.addCredit(receiverId, 1);
		assertTrue(queue.getAvailableReaders().contains(reader));
	}

	// cancel this subpartition view
	queue.cancel(receiverId);
	channel.runPendingTasks();

	assertFalse(queue.getAvailableReaders().contains(reader));
	// the partition and its reader view should all be released
	assertTrue(reader.isReleased());
	assertTrue(partition.isReleased());
	for (ResultSubpartition subpartition : partition.getAllPartitions()) {
		assertTrue(subpartition.isReleased());
	}

	// cleanup
	channel.close();
}
 

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

@Test
public void testResumeConsumption() {
	NettyMessage.ResumeConsumption expected = new NettyMessage.ResumeConsumption(new InputChannelID());
	NettyMessage.ResumeConsumption actual = encodeAndDecode(expected, channel);

	assertEquals(expected.receiverId, actual.receiverId);
}
 

private void testEncodeDecodeBuffer(boolean testReadOnlyBuffer) {
	NetworkBuffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(1024), FreeingBufferRecycler.INSTANCE);

	for (int i = 0; i < 1024; i += 4) {
		buffer.writeInt(i);
	}

	Buffer testBuffer = testReadOnlyBuffer ? buffer.readOnlySlice() : buffer;

	NettyMessage.BufferResponse expected = new NettyMessage.BufferResponse(
		testBuffer, random.nextInt(), new InputChannelID(), random.nextInt());
	NettyMessage.BufferResponse actual = encodeAndDecode(expected);

	// Netty 4.1 is not copying the messages, but retaining slices of them. BufferResponse actual is in this case
	// holding a reference to the buffer. Buffer will be recycled only once "actual" will be released.
	assertFalse(buffer.isRecycled());
	assertFalse(testBuffer.isRecycled());

	final ByteBuf retainedSlice = actual.getNettyBuffer();

	// Ensure not recycled and same size as original buffer
	assertEquals(1, retainedSlice.refCnt());
	assertEquals(1024, retainedSlice.readableBytes());

	for (int i = 0; i < 1024; i += 4) {
		assertEquals(i, retainedSlice.readInt());
	}

	// Release the retained slice
	actual.releaseBuffer();
	assertEquals(0, retainedSlice.refCnt());
	assertTrue(buffer.isRecycled());
	assertTrue(testBuffer.isRecycled());

	assertEquals(expected.sequenceNumber, actual.sequenceNumber);
	assertEquals(expected.receiverId, actual.receiverId);
	assertEquals(expected.backlog, actual.backlog);
}
 

/**
 * Verifies that {@link RemoteInputChannel#onFailedPartitionRequest()} is called when a
 * {@link PartitionNotFoundException} is received.
 */
@Test
public void testReceivePartitionNotFoundException() throws Exception {
	// Minimal mock of a remote input channel
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(TestBufferFactory.createBuffer(0));

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	final ErrorResponse partitionNotFound = new ErrorResponse(
		new PartitionNotFoundException(new ResultPartitionID()),
		inputChannel.getInputChannelId());

	final PartitionRequestClientHandler client = new PartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	// Mock channel context
	ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
	when(ctx.channel()).thenReturn(mock(Channel.class));

	client.channelActive(ctx);

	client.channelRead(ctx, partitionNotFound);

	verify(inputChannel, times(1)).onFailedPartitionRequest();
}
 

/**
 * Tests a fix for FLINK-1761.
 *
 * <p>FLINK-1761 discovered an IndexOutOfBoundsException, when receiving buffers of size 0.
 */
@Test
public void testReceiveEmptyBuffer() throws Exception {
	// Minimal mock of a remote input channel
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(TestBufferFactory.createBuffer(0));

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	// An empty buffer of size 0
	final Buffer emptyBuffer = TestBufferFactory.createBuffer(0);

	final int backlog = -1;
	final BufferResponse receivedBuffer = createBufferResponse(
		emptyBuffer, 0, inputChannel.getInputChannelId(), backlog);

	final PartitionRequestClientHandler client = new PartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	// Read the empty buffer
	client.channelRead(mock(ChannelHandlerContext.class), receivedBuffer);

	// This should not throw an exception
	verify(inputChannel, never()).onError(any(Throwable.class));
	verify(inputChannel, times(1)).onEmptyBuffer(0, backlog);
}
 

/**
 * In case of enqueuing an empty reader and a reader that actually has some buffers when channel is not writable,
 * on channelWritability change event should result in reading all of the messages.
 */
@Test
public void testNotifyReaderNonEmptyOnEmptyReaders() throws Exception {
	final int buffersToWrite = 5;
	PartitionRequestQueue queue = new PartitionRequestQueue();
	EmbeddedChannel channel = new EmbeddedChannel(queue);

	CreditBasedSequenceNumberingViewReader reader1 = new CreditBasedSequenceNumberingViewReader(new InputChannelID(0, 0), 10, queue);
	CreditBasedSequenceNumberingViewReader reader2 = new CreditBasedSequenceNumberingViewReader(new InputChannelID(1, 1), 10, queue);

	reader1.requestSubpartitionView((partitionId, index, availabilityListener) -> new EmptyAlwaysAvailableResultSubpartitionView(), new ResultPartitionID(), 0);
	reader1.notifyDataAvailable();
	assertTrue(reader1.isAvailable());
	assertFalse(reader1.isRegisteredAsAvailable());

	channel.unsafe().outboundBuffer().setUserDefinedWritability(1, false);
	assertFalse(channel.isWritable());

	reader1.notifyDataAvailable();
	channel.runPendingTasks();

	reader2.notifyDataAvailable();
	reader2.requestSubpartitionView((partitionId, index, availabilityListener) -> new DefaultBufferResultSubpartitionView(buffersToWrite), new ResultPartitionID(), 0);
	assertTrue(reader2.isAvailable());
	assertFalse(reader2.isRegisteredAsAvailable());

	reader2.notifyDataAvailable();

	// changing a channel writability should result in draining both reader1 and reader2
	channel.unsafe().outboundBuffer().setUserDefinedWritability(1, true);
	channel.runPendingTasks();
	assertEquals(buffersToWrite, channel.outboundMessages().size());
}
 

/**
 * Tests a fix for FLINK-1761.
 *
 * <p>FLINK-1761 discovered an IndexOutOfBoundsException, when receiving buffers of size 0.
 */
@Test
public void testReceiveEmptyBuffer() throws Exception {
	// Minimal mock of a remote input channel
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(TestBufferFactory.createBuffer(0));

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	// An empty buffer of size 0
	final Buffer emptyBuffer = TestBufferFactory.createBuffer(0);

	final int backlog = 2;
	final BufferResponse receivedBuffer = createBufferResponse(
		emptyBuffer, 0, inputChannel.getInputChannelId(), backlog);

	final CreditBasedPartitionRequestClientHandler client = new CreditBasedPartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	// Read the empty buffer
	client.channelRead(mock(ChannelHandlerContext.class), receivedBuffer);

	// This should not throw an exception
	verify(inputChannel, never()).onError(any(Throwable.class));
	verify(inputChannel, times(1)).onEmptyBuffer(0, backlog);
}