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

/**
 * Sends a task event backwards to an intermediate result partition producer.
 *
 * <p>Backwards task events flow between readers and writers and therefore
 * will only work when both are running at the same time, which is only
 * guaranteed to be the case when both the respective producer and
 * consumer task run pipelined.
 */
@Override
public void sendTaskEvent(ResultPartitionID partitionId, TaskEvent event, final RemoteInputChannel inputChannel) throws IOException {
	checkNotClosed();

	tcpChannel.writeAndFlush(new TaskEventRequest(event, partitionId, inputChannel.getInputChannelId()))
			.addListener(
					new ChannelFutureListener() {
						@Override
						public void operationComplete(ChannelFuture future) throws Exception {
							if (!future.isSuccess()) {
								SocketAddress remoteAddr = future.channel().remoteAddress();
								inputChannel.onError(new LocalTransportException(
									String.format("Sending the task event to '%s' failed.", remoteAddr),
									future.channel().localAddress(), future.cause()
								));
							}
						}
					});
}
 

/**
 * Iterates over all input channels and collects the minimum number of queued buffers in a
 * channel in a best-effort way.
 *
 * @return minimum number of queued buffers per channel (<tt>0</tt> if no channels exist)
 */
int refreshAndGetMin() {
	int min = Integer.MAX_VALUE;

	Collection<InputChannel> channels = inputGate.getInputChannels().values();

	for (InputChannel channel : channels) {
		if (channel instanceof RemoteInputChannel) {
			RemoteInputChannel rc = (RemoteInputChannel) channel;

			int size = rc.unsynchronizedGetNumberOfQueuedBuffers();
			min = Math.min(min, size);
		}
	}

	if (min == Integer.MAX_VALUE) { // in case all channels are local, or the channel collection was empty
		return 0;
	}
	return min;
}
 

/**
 * Verifies that {@link RemoteInputChannel#onError(Throwable)} is called when a
 * {@link BufferResponse} is received but no available buffer in input channel.
 */
@Test
public void testThrowExceptionForNoAvailableBuffer() throws Exception {
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = spy(InputChannelBuilder.newBuilder().buildRemoteChannel(inputGate));

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

	assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());

	final BufferResponse bufferResponse = createBufferResponse(
		TestBufferFactory.createBuffer(TestBufferFactory.BUFFER_SIZE),
		0,
		inputChannel.getInputChannelId(),
		2,
		new NetworkBufferAllocator(handler));
	assertNull(bufferResponse.getBuffer());

	handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);
	verify(inputChannel, times(1)).onError(any(IllegalStateException.class));
}
 

/**
 * Verifies that {@link RemoteInputChannel#onError(Throwable)} is called when a
 * {@link BufferResponse} is received but no available buffer in input channel.
 */
@Test
public void testThrowExceptionForNoAvailableBuffer() throws Exception {
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = spy(InputChannelBuilder.newBuilder().buildRemoteAndSetToGate(inputGate));

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

	assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());

	final BufferResponse bufferResponse = createBufferResponse(
		TestBufferFactory.createBuffer(TestBufferFactory.BUFFER_SIZE), 0, inputChannel.getInputChannelId(), 2);
	handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

	verify(inputChannel, times(1)).onError(any(IllegalStateException.class));
}
 

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

private void notifyAllChannelsOfErrorAndClose(Throwable cause) {
	if (channelError.compareAndSet(null, cause)) {
		try {
			for (RemoteInputChannel inputChannel : inputChannels.values()) {
				inputChannel.onError(cause);
			}
		} catch (Throwable t) {
			// We can only swallow the Exception at this point. :(
			LOG.warn("An Exception was thrown during error notification of a remote input channel.", t);
		} finally {
			inputChannels.clear();
			clientOutboundMessages.clear();

			if (ctx != null) {
				ctx.close();
			}
		}
	}
}
 

/**
 * Iterates over all input channels and collects the average number of queued buffers in a
 * channel in a best-effort way.
 *
 * @return average number of queued buffers per channel
 */
float refreshAndGetAvg() {
	long total = 0;
	int count = 0;

	for (InputChannel channel : inputGate.getInputChannels().values()) {
		if (channel instanceof RemoteInputChannel) {
			RemoteInputChannel rc = (RemoteInputChannel) channel;

			int size = rc.unsynchronizedGetNumberOfQueuedBuffers();
			total += size;
			++count;
		}
	}

	return count == 0 ? 0 : total / (float) count;
}
 

/**
 * Sends a task event backwards to an intermediate result partition producer.
 *
 * <p>Backwards task events flow between readers and writers and therefore
 * will only work when both are running at the same time, which is only
 * guaranteed to be the case when both the respective producer and
 * consumer task run pipelined.
 */
@Override
public void sendTaskEvent(ResultPartitionID partitionId, TaskEvent event, final RemoteInputChannel inputChannel) throws IOException {
	checkNotClosed();

	tcpChannel.writeAndFlush(new TaskEventRequest(event, partitionId, inputChannel.getInputChannelId()))
			.addListener(
					new ChannelFutureListener() {
						@Override
						public void operationComplete(ChannelFuture future) throws Exception {
							if (!future.isSuccess()) {
								SocketAddress remoteAddr = future.channel().remoteAddress();
								inputChannel.onError(new LocalTransportException(
									String.format("Sending the task event to '%s' failed.", remoteAddr),
									future.channel().localAddress(), future.cause()
								));
							}
						}
					});
}
 

private void notifyAllChannelsOfErrorAndClose(Throwable cause) {
	if (channelError.compareAndSet(null, cause)) {
		try {
			for (RemoteInputChannel inputChannel : inputChannels.values()) {
				inputChannel.onError(cause);
			}
		} catch (Throwable t) {
			// We can only swallow the Exception at this point. :(
			LOG.warn("An Exception was thrown during error notification of a remote input channel.", t);
		} finally {
			inputChannels.clear();
			inputChannelsWithCredit.clear();

			if (ctx != null) {
				ctx.close();
			}
		}
	}
}
 

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

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

@Test
public void testDoublePartitionRequest() throws Exception {
	final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
	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);

		// The input channel should only send one partition request
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(numExclusiveBuffers, ((PartitionRequest) readFromOutbound).credit);

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

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

public static RemoteInputChannel createRemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ConnectionManager connectionManager) {

	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setConnectionManager(connectionManager)
		.buildRemoteChannel(inputGate);
}
 

@Test
public void testCancelBeforeActive() throws Exception {

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

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

	// Don't throw NPE
	client.cancelRequestFor(null);

	// Don't throw NPE, because channel is not active yet
	client.cancelRequestFor(inputChannel.getInputChannelId());
}
 

@Test
public void testDoublePartitionRequest() throws Exception {
	final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new NettyPartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

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

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

		// The input channel should only send one partition request
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(numExclusiveBuffers, ((PartitionRequest) readFromOutbound).credit);

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

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

/**
 * Tries to write&flush unannounced credits for the next input channel in queue.
 *
 * <p>This method may be called by the first input channel enqueuing, or the complete
 * future's callback in previous input channel, or the channel writability changed event.
 */
private void writeAndFlushNextMessageIfPossible(Channel channel) {
	if (channelError.get() != null || !channel.isWritable()) {
		return;
	}

	while (true) {
		RemoteInputChannel inputChannel = inputChannelsWithCredit.poll();

		// The input channel may be null because of the write callbacks
		// that are executed after each write.
		if (inputChannel == null) {
			return;
		}

		//It is no need to notify credit for the released channel.
		if (!inputChannel.isReleased()) {
			AddCredit msg = new AddCredit(
				inputChannel.getPartitionId(),
				inputChannel.getAndResetUnannouncedCredit(),
				inputChannel.getInputChannelId());

			// Write and flush and wait until this is done before
			// trying to continue with the next input channel.
			channel.writeAndFlush(msg).addListener(writeListener);

			return;
		}
	}
}
 

@Test
public void testDoNotFailHandlerOnSingleChannelFailure() throws Exception {
	// Setup
	final int bufferSize = 1024;
	final String expectedMessage = "test exception on buffer";
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, bufferSize, 2);
	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = new TestRemoteInputChannelForError(inputGate, expectedMessage);
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();

	try {
		inputGate.setInputChannels(inputChannel);
		inputGate.assignExclusiveSegments();
		inputGate.requestPartitions();
		handler.addInputChannel(inputChannel);

		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(bufferSize),
			0,
			inputChannel.getInputChannelId(),
			1,
			new NetworkBufferAllocator(handler));

		// It will trigger an expected exception from TestRemoteInputChannelForError#onBuffer
		handler.channelRead(null, bufferResponse);

		// The handler should not be tagged as error for above excepted exception
		handler.checkError();

		try {
			// The input channel should be tagged as error and the respective exception is thrown via #getNext
			inputGate.getNext();
		} catch (IOException ignored) {
			assertEquals(expectedMessage, ignored.getMessage());
		}
	} finally {
		// Cleanup
		releaseResource(inputGate, networkBufferPool);
	}
}
 

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

/**
 * Verifies that {@link RemoteInputChannel#onBuffer(Buffer, int, int)} is called when a
 * {@link BufferResponse} is received.
 */
@Test
public void testReceiveBuffer() throws Exception {
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

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

		final int backlog = 2;
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), backlog);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertEquals(2, inputChannel.getSenderBacklog());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

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

/**
 * Iterates over all input channels and collects the maximum number of queued buffers in a
 * channel in a best-effort way.
 *
 * @return maximum number of queued buffers per channel
 */
int refreshAndGetMax() {
	int max = 0;

	for (InputChannel channel : inputGate.getInputChannels().values()) {
		if (channel instanceof RemoteInputChannel) {
			RemoteInputChannel rc = (RemoteInputChannel) channel;

			int size = rc.unsynchronizedGetNumberOfQueuedBuffers();
			max = Math.max(max, size);
		}
	}

	return max;
}
 

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

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

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

		// The input channel should only send one partition request
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(numExclusiveBuffers, ((PartitionRequest) readFromOutbound).credit);

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

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

@Test
public void testCancelBeforeActive() throws Exception {

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

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

	// Don't throw NPE
	client.cancelRequestFor(null);

	// Don't throw NPE, because channel is not active yet
	client.cancelRequestFor(inputChannel.getInputChannelId());
}
 

@Test
public void testConsumptionWithRemoteChannels() throws Exception {
	final int numberOfChannels = 11;
	final int buffersPerChannel = 1000;

	final ConnectionManager connManager = createDummyConnectionManager();
	final Source[] sources = new Source[numberOfChannels];

	final SingleInputGate gate = new SingleInputGate(
			"Test Task Name",
			new JobID(),
			new IntermediateDataSetID(), ResultPartitionType.PIPELINED,
			0,
			numberOfChannels,
			mock(TaskActions.class),
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup(),
			true);

	for (int i = 0; i < numberOfChannels; i++) {
		RemoteInputChannel channel = new RemoteInputChannel(
				gate, i, new ResultPartitionID(), mock(ConnectionID.class),
				connManager, 0, 0, UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());
		gate.setInputChannel(new IntermediateResultPartitionID(), channel);

		sources[i] = new RemoteChannelSource(channel);
	}

	ProducerThread producer = new ProducerThread(sources, numberOfChannels * buffersPerChannel, 4, 10);
	ConsumerThread consumer = new ConsumerThread(gate, numberOfChannels * buffersPerChannel);
	producer.start();
	consumer.start();

	// the 'sync()' call checks for exceptions and failed assertions
	producer.sync();
	consumer.sync();
}
 

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

@Override
public int calculateUsedBuffers(SingleInputGate inputGate) {
	int availableFloatingBuffers = 0;
	BufferPool bufferPool = inputGate.getBufferPool();
	if (bufferPool != null) {
		int requestedFloatingBuffers = bufferPool.bestEffortGetNumOfUsedBuffers();
		for (InputChannel ic : inputGate.getInputChannels().values()) {
			if (ic instanceof RemoteInputChannel) {
				availableFloatingBuffers += ((RemoteInputChannel) ic).unsynchronizedGetFloatingBuffersAvailable();
			}
		}
		return Math.max(0, requestedFloatingBuffers - availableFloatingBuffers);
	}
	return 0;
}
 

@Override
public int calculateUsedBuffers(SingleInputGate inputGate) {
	int usedBuffers = 0;
	for (InputChannel ic : inputGate.getInputChannels().values()) {
		if (ic instanceof RemoteInputChannel) {
			usedBuffers += ((RemoteInputChannel) ic).unsynchronizedGetExclusiveBuffersUsed();
		}
	}
	return usedBuffers;
}
 

public static RemoteInputChannel createRemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ConnectionManager connectionManager) {

	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setConnectionManager(connectionManager)
		.setMemorySegmentProvider(new UnpooledMemorySegmentProvider(32 * 1024))
		.buildRemoteAndSetToGate(inputGate);
}
 

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

private static RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartition resultPartition,
		ConnectionManager connManager) {
	return InputChannelBuilder.newBuilder()
		.setChannelIndex(channelIndex)
		.setPartitionId(resultPartition.getPartitionId())
		.setConnectionManager(connManager)
		.buildRemoteChannel(inputGate);
}
 

/**
 * Triggered by notifying credit available in the client handler pipeline.
 *
 * <p>Enqueues the input channel and will trigger write&flush unannounced credits
 * for this input channel if it is the first one in the queue.
 */
@Override
public void userEventTriggered(ChannelHandlerContext ctx, Object msg) throws Exception {
	if (msg instanceof RemoteInputChannel) {
		boolean triggerWrite = inputChannelsWithCredit.isEmpty();

		inputChannelsWithCredit.add((RemoteInputChannel) msg);

		if (triggerWrite) {
			writeAndFlushNextMessageIfPossible(ctx.channel());
		}
	} else {
		ctx.fireUserEventTriggered(msg);
	}
}