org.apache.flink.runtime.io.network.buffer.BufferProvider Java Examples

private boolean waitForBuffer(BufferProvider bufferProvider, NettyMessage.BufferResponse bufferResponse) {

			stagedBufferResponse = bufferResponse;

			if (bufferProvider.addBufferListener(this)) {
				if (ctx.channel().config().isAutoRead()) {
					ctx.channel().config().setAutoRead(false);
				}

				return true;
			}
			else {
				stagedBufferResponse = null;

				return false;
			}
		}
 

public TestPartitionProducerBufferSource(
		int parallelism,
		BufferProvider bufferProvider,
		int numberOfBuffersToProduce) {

	this.bufferProvider = bufferProvider;
	this.channelIndexes = Lists.newArrayListWithCapacity(
			parallelism * numberOfBuffersToProduce);

	// Array of channel indexes to produce buffers for
	for (byte i = 0; i < parallelism; i++) {
		for (int j = 0; j < numberOfBuffersToProduce; j++) {
			channelIndexes.add(i);
		}
	}

	// Random buffer to channel ordering
	Collections.shuffle(channelIndexes);
}
 

@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 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 CreditBasedPartitionRequestClientHandler client = new CreditBasedPartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

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

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

public TestPartitionProducerBufferSource(
		int parallelism,
		BufferProvider bufferProvider,
		int numberOfBuffersToProduce) {

	this.bufferProvider = bufferProvider;
	this.channelIndexes = Lists.newArrayListWithCapacity(
			parallelism * numberOfBuffersToProduce);

	// Array of channel indexes to produce buffers for
	for (byte i = 0; i < parallelism; i++) {
		for (int j = 0; j < numberOfBuffersToProduce; j++) {
			channelIndexes.add(i);
		}
	}

	// Random buffer to channel ordering
	Collections.shuffle(channelIndexes);
}
 

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

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

private boolean waitForBuffer(BufferProvider bufferProvider, NettyMessage.BufferResponse bufferResponse) {

			stagedBufferResponse = bufferResponse;

			if (bufferProvider.addBufferListener(this)) {
				if (ctx.channel().config().isAutoRead()) {
					ctx.channel().config().setAutoRead(false);
				}

				return true;
			}
			else {
				stagedBufferResponse = null;

				return false;
			}
		}
 

public TestPartitionProducerBufferSource(
		int parallelism,
		BufferProvider bufferProvider,
		int numberOfBuffersToProduce) {

	this.bufferProvider = bufferProvider;
	this.channelIndexes = Lists.newArrayListWithCapacity(
			parallelism * numberOfBuffersToProduce);

	// Array of channel indexes to produce buffers for
	for (byte i = 0; i < parallelism; i++) {
		for (int j = 0; j < numberOfBuffersToProduce; j++) {
			channelIndexes.add(i);
		}
	}

	// Random buffer to channel ordering
	Collections.shuffle(channelIndexes);
}
 

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

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

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

public RecordOrEventCollectingResultPartitionWriter(
		Collection<Object> output,
		BufferProvider bufferProvider,
		TypeSerializer<T> serializer) {
	super(bufferProvider);
	this.output = checkNotNull(output);
	this.delegate = new NonReusingDeserializationDelegate<>(checkNotNull(serializer));
}
 

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

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

public BufferProvider getBufferProvider() throws IOException {
	if (isReleased.get()) {
		return null;
	}

	return inputGate.getBufferProvider();
}
 

public RecordOrEventCollectingResultPartitionWriter(
		Collection<Object> output,
		BufferProvider bufferProvider,
		TypeSerializer<T> serializer) {
	super(bufferProvider);
	this.output = checkNotNull(output);
	this.delegate = new NonReusingDeserializationDelegate<>(checkNotNull(serializer));
}
 

public BufferProvider getBufferProvider() throws IOException {
	if (isReleased.get()) {
		return null;
	}

	return inputGate.getBufferProvider();
}
 

/**
 * 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 CreditBasedPartitionRequestClientHandler client = new CreditBasedPartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

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

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

public RecordOrEventCollectingResultPartitionWriter(
		Collection<Object> output,
		BufferProvider bufferProvider,
		TypeSerializer<T> serializer) {
	super(bufferProvider);
	this.output = checkNotNull(output);
	this.delegate = new NonReusingDeserializationDelegate<>(checkNotNull(serializer));
}
 

public BufferProvider getBufferProvider() throws IOException {
	if (isReleased.get()) {
		return null;
	}

	return inputGate.getBufferProvider();
}
 

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

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

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

private static SpillableSubpartition createSubpartition(IOManager ioManager) {
	ResultPartition parent = mock(ResultPartition.class);
	BufferProvider bufferProvider = mock(BufferProvider.class);
	when(parent.getBufferProvider()).thenReturn(bufferProvider);
	when(bufferProvider.getMemorySegmentSize()).thenReturn(32 * 1024);
	return new SpillableSubpartition(0, parent, ioManager);
}
 

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