Java Code Examples for org.apache.flink.runtime.io.network.partition.consumer.RemoteInputChannel#getInputChannelId()

@Before
public void setup() throws IOException, InterruptedException {
	networkBufferPool = new NetworkBufferPool(8, BUFFER_SIZE, 8);
	inputGate = createSingleInputGate(1, networkBufferPool);
	RemoteInputChannel inputChannel = createRemoteInputChannel(
		inputGate,
		new TestingPartitionRequestClient());
	inputChannel.requestSubpartition(0);
	inputGate.setInputChannels(inputChannel);
	inputGate.assignExclusiveSegments();

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

	channel = new EmbeddedChannel(
		new NettyMessageEncoder(), // For outbound messages
		new NettyMessageClientDecoderDelegate(handler)); // For inbound messages

	inputChannelId = inputChannel.getInputChannelId();
}
 

@Before
public void setup() throws IOException, InterruptedException {
	CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	networkBufferPool = new NetworkBufferPool(
		NUMBER_OF_BUFFER_RESPONSES,
		BUFFER_SIZE,
		NUMBER_OF_BUFFER_RESPONSES);
	channel = new EmbeddedChannel(new NettyMessageClientDecoderDelegate(handler));

	inputGate = createSingleInputGate(1, networkBufferPool);
	RemoteInputChannel inputChannel = createRemoteInputChannel(
		inputGate,
		new TestingPartitionRequestClient());
	inputGate.setInputChannels(inputChannel);
	inputGate.assignExclusiveSegments();
	inputChannel.requestSubpartition(0);
	handler.addInputChannel(inputChannel);
	inputChannelId = inputChannel.getInputChannelId();

	SingleInputGate releasedInputGate = createSingleInputGate(1, networkBufferPool);
	RemoteInputChannel releasedInputChannel = new InputChannelBuilder()
		.buildRemoteChannel(inputGate);
	releasedInputGate.close();
	handler.addInputChannel(releasedInputChannel);
	releasedInputChannelId = releasedInputChannel.getInputChannelId();
}
 

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

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

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

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

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