Java Code Examples for org.apache.flink.shaded.netty4.io.netty.channel.embedded.EmbeddedChannel#runPendingTasks()

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

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

@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 List<NettyMessage> decodeMessages(EmbeddedChannel channel, List<ByteBuf> inputBuffers) {
	for (ByteBuf buffer : inputBuffers) {
		channel.writeInbound(buffer);
	}

	channel.runPendingTasks();

	List<NettyMessage> decodedMessages = new ArrayList<>();
	Object input;
	while ((input = channel.readInbound()) != null) {
		assertTrue(input instanceof NettyMessage);
		decodedMessages.add((NettyMessage) input);
	}

	return decodedMessages;
}
 

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

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

@Test
public void testResumeConsumption() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = createPartitionRequestClient(channel, handler);

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, 2);
	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client);

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		inputChannel.resumeConsumption();
		channel.runPendingTasks();
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(ResumeConsumption.class));
		assertEquals(inputChannel.getInputChannelId(), ((ResumeConsumption) readFromOutbound).receiverId);

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.close();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

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

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

/**
 * 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 {@link PartitionRequestQueue#enqueueAvailableReader(NetworkSequenceViewReader)},
 * verifying the reader would be enqueued in the pipeline if the next sending buffer is an event
 * even though it has no available credits.
 */
@Test
public void testEnqueueReaderByNotifyingEventBuffer() throws Exception {
	// setup
	final ResultSubpartitionView view = new NextIsEventResultSubpartitionView();

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

	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(partitionProvider, new ResultPartitionID(), 0);

	// block the channel so that we see an intermediate state in the test
	ByteBuf channelBlockingBuffer = blockChannel(channel);
	assertNull(channel.readOutbound());

	// Notify an available event buffer to trigger enqueue the reader
	reader.notifyDataAvailable();

	channel.runPendingTasks();

	// The reader is enqueued in the pipeline because the next buffer is an event, even though no credits are available
	assertThat(queue.getAvailableReaders(), contains(reader)); // contains only (this) one!
	assertEquals(0, reader.getNumCreditsAvailable());

	// Flush the buffer to make the channel writable again and see the final results
	channel.flush();
	assertSame(channelBlockingBuffer, channel.readOutbound());

	assertEquals(0, queue.getAvailableReaders().size());
	assertEquals(0, reader.getNumCreditsAvailable());
	assertNull(channel.readOutbound());
}
 

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

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

/**
 * Tests {@link PartitionRequestQueue#enqueueAvailableReader(NetworkSequenceViewReader)},
 * verifying the reader would be enqueued in the pipeline if the next sending buffer is an event
 * even though it has no available credits.
 */
@Test
public void testEnqueueReaderByNotifyingEventBuffer() throws Exception {
	// setup
	final ResultSubpartitionView view = new NextIsEventResultSubpartitionView();

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

	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(partitionProvider, new ResultPartitionID(), 0);

	// block the channel so that we see an intermediate state in the test
	ByteBuf channelBlockingBuffer = blockChannel(channel);
	assertNull(channel.readOutbound());

	// Notify an available event buffer to trigger enqueue the reader
	reader.notifyDataAvailable();

	channel.runPendingTasks();

	// The reader is enqueued in the pipeline because the next buffer is an event, even though no credits are available
	assertThat(queue.getAvailableReaders(), contains(reader)); // contains only (this) one!
	assertEquals(0, reader.getNumCreditsAvailable());

	// Flush the buffer to make the channel writable again and see the final results
	channel.flush();
	assertSame(channelBlockingBuffer, channel.readOutbound());

	assertEquals(0, queue.getAvailableReaders().size());
	assertEquals(0, reader.getNumCreditsAvailable());
	assertNull(channel.readOutbound());
}
 

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.addCreditOrResumeConsumption(receiverId, viewReader -> viewReader.addCredit(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();
}
 

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

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline, but {@link AddCredit}
 * message is not sent actually when this input channel is released.
 */
@Test
public void testNotifyCreditAvailableAfterReleased() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new PartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		inputChannel.requestSubpartition(0);

		// This should send the partition request
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// Trigger request floating buffers via buffer response to notify credits available
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(2, inputChannel.getUnannouncedCredit());

		// Release the input channel
		inputGate.releaseAllResources();

		// it should send a close request after releasing the input channel,
		// but will not notify credits for a released input channel.
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(CloseRequest.class));

		channel.runPendingTasks();

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Tests {@link PartitionRequestQueue#enqueueAvailableReader(NetworkSequenceViewReader)},
 * verifying the reader would be enqueued in the pipeline iff it has both available credits and buffers.
 */
@Test
public void testEnqueueReaderByNotifyingBufferAndCredit() throws Exception {
	// setup
	final ResultSubpartitionView view = new DefaultBufferResultSubpartitionView(10);

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

	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(partitionProvider, new ResultPartitionID(), 0);
	queue.notifyReaderCreated(reader);

	// block the channel so that we see an intermediate state in the test
	ByteBuf channelBlockingBuffer = blockChannel(channel);
	assertNull(channel.readOutbound());

	// Notify available buffers to trigger enqueue the reader
	final int notifyNumBuffers = 5;
	for (int i = 0; i < notifyNumBuffers; i++) {
		reader.notifyDataAvailable();
	}

	channel.runPendingTasks();

	// the reader is not enqueued in the pipeline because no credits are available
	// -> it should still have the same number of pending buffers
	assertEquals(0, queue.getAvailableReaders().size());
	assertTrue(reader.hasBuffersAvailable());
	assertFalse(reader.isRegisteredAsAvailable());
	assertEquals(0, reader.getNumCreditsAvailable());

	// Notify available credits to trigger enqueue the reader again
	final int notifyNumCredits = 3;
	for (int i = 1; i <= notifyNumCredits; i++) {
		queue.addCreditOrResumeConsumption(receiverId, viewReader -> viewReader.addCredit(1));

		// the reader is enqueued in the pipeline because it has both available buffers and credits
		// since the channel is blocked though, we will not process anything and only enqueue the
		// reader once
		assertTrue(reader.isRegisteredAsAvailable());
		assertThat(queue.getAvailableReaders(), contains(reader)); // contains only (this) one!
		assertEquals(i, reader.getNumCreditsAvailable());
		assertTrue(reader.hasBuffersAvailable());
	}

	// Flush the buffer to make the channel writable again and see the final results
	channel.flush();
	assertSame(channelBlockingBuffer, channel.readOutbound());

	assertEquals(0, queue.getAvailableReaders().size());
	assertEquals(0, reader.getNumCreditsAvailable());
	assertTrue(reader.hasBuffersAvailable());
	assertFalse(reader.isRegisteredAsAvailable());
	for (int i = 1; i <= notifyNumCredits; i++) {
		assertThat(channel.readOutbound(), instanceOf(NettyMessage.BufferResponse.class));
	}
	assertNull(channel.readOutbound());
}
 

/**
 * Tests {@link PartitionRequestQueue#enqueueAvailableReader(NetworkSequenceViewReader)},
 * verifying the reader would be enqueued in the pipeline iff it has both available credits and buffers.
 */
@Test
public void testEnqueueReaderByNotifyingBufferAndCredit() throws Exception {
	// setup
	final ResultSubpartitionView view = new DefaultBufferResultSubpartitionView(10);

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

	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(partitionProvider, new ResultPartitionID(), 0);
	queue.notifyReaderCreated(reader);

	// block the channel so that we see an intermediate state in the test
	ByteBuf channelBlockingBuffer = blockChannel(channel);
	assertNull(channel.readOutbound());

	// Notify available buffers to trigger enqueue the reader
	final int notifyNumBuffers = 5;
	for (int i = 0; i < notifyNumBuffers; i++) {
		reader.notifyDataAvailable();
	}

	channel.runPendingTasks();

	// the reader is not enqueued in the pipeline because no credits are available
	// -> it should still have the same number of pending buffers
	assertEquals(0, queue.getAvailableReaders().size());
	assertTrue(reader.hasBuffersAvailable());
	assertFalse(reader.isRegisteredAsAvailable());
	assertEquals(0, reader.getNumCreditsAvailable());

	// Notify available credits to trigger enqueue the reader again
	final int notifyNumCredits = 3;
	for (int i = 1; i <= notifyNumCredits; i++) {
		queue.addCredit(receiverId, 1);

		// the reader is enqueued in the pipeline because it has both available buffers and credits
		// since the channel is blocked though, we will not process anything and only enqueue the
		// reader once
		assertTrue(reader.isRegisteredAsAvailable());
		assertThat(queue.getAvailableReaders(), contains(reader)); // contains only (this) one!
		assertEquals(i, reader.getNumCreditsAvailable());
		assertTrue(reader.hasBuffersAvailable());
	}

	// Flush the buffer to make the channel writable again and see the final results
	channel.flush();
	assertSame(channelBlockingBuffer, channel.readOutbound());

	assertEquals(0, queue.getAvailableReaders().size());
	assertEquals(0, reader.getNumCreditsAvailable());
	assertTrue(reader.hasBuffersAvailable());
	assertFalse(reader.isRegisteredAsAvailable());
	for (int i = 1; i <= notifyNumCredits; i++) {
		assertThat(channel.readOutbound(), instanceOf(NettyMessage.BufferResponse.class));
	}
	assertNull(channel.readOutbound());
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline for notifying credits,
 * and verifies the behaviour of credit notification by triggering channel's writability changed.
 */
@Test
public void testNotifyCreditAvailable() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new NettyPartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, 2);
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel1 = createRemoteInputChannel(inputGate, client, networkBufferPool);
	final RemoteInputChannel inputChannel2 = createRemoteInputChannel(inputGate, client, networkBufferPool);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		inputChannel1.requestSubpartition(0);
		inputChannel2.requestSubpartition(0);

		// The two input channels should send partition requests
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel1.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel2.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// The buffer response will take one available buffer from input channel, and it will trigger
		// requesting (backlog + numExclusiveBuffers - numAvailableBuffers) floating buffers
		final BufferResponse bufferResponse1 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel1.getInputChannelId(), 1);
		final BufferResponse bufferResponse2 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel2.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse2);

		assertEquals(2, inputChannel1.getUnannouncedCredit());
		assertEquals(2, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The two input channels should notify credits availability via the writable channel
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel1.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel2.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);
		assertNull(channel.readOutbound());

		ByteBuf channelBlockingBuffer = blockChannel(channel);

		// Trigger notify credits availability via buffer response on the condition of an un-writable channel
		final BufferResponse bufferResponse3 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 1, inputChannel1.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse3);

		assertEquals(1, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The input channel will not notify credits via un-writable channel
		assertFalse(channel.isWritable());
		assertNull(channel.readOutbound());

		// Flush the buffer to make the channel writable again
		channel.flush();
		assertSame(channelBlockingBuffer, channel.readOutbound());

		// The input channel should notify credits via channel's writability changed event
		assertTrue(channel.isWritable());
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(1, ((AddCredit) readFromOutbound).credit);
		assertEquals(0, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		// no more messages
		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.close();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}