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

private CheckpointedInputGate createInputGate(
		int numberOfChannels,
		AbstractInvokable toNotify) throws IOException {
	final NettyShuffleEnvironment environment = new NettyShuffleEnvironmentBuilder().build();
	SingleInputGate gate = new SingleInputGateBuilder()
		.setNumberOfChannels(numberOfChannels)
		.setupBufferPoolFactory(environment)
		.build();
	gate.setInputChannels(
		IntStream.range(0, numberOfChannels)
		.mapToObj(channelIndex ->
			InputChannelBuilder.newBuilder()
				.setChannelIndex(channelIndex)
				.setupFromNettyShuffleEnvironment(environment)
				.setConnectionManager(new TestingConnectionManager())
				.buildRemoteChannel(gate))
		.toArray(RemoteInputChannel[]::new));
	sequenceNumbers = new int[numberOfChannels];

	gate.setup();
	gate.requestPartitions();

	return createCheckpointedInputGate(gate, toNotify);
}
 

@Test
public void testPreviousHandlerReset() throws Exception {
	SingleInputGate inputGate = new SingleInputGateBuilder().setNumberOfChannels(2).build();
	inputGate.setInputChannels(new TestInputChannel(inputGate, 0), new TestInputChannel(inputGate, 1));
	TestInvokable target = new TestInvokable();
	CheckpointBarrierAligner alignedHandler = new CheckpointBarrierAligner("test", target, inputGate);
	CheckpointBarrierUnaligner unalignedHandler = new CheckpointBarrierUnaligner(TestSubtaskCheckpointCoordinator.INSTANCE, "test", target, inputGate);
	AlternatingCheckpointBarrierHandler barrierHandler = new AlternatingCheckpointBarrierHandler(alignedHandler, unalignedHandler, target);

	for (int i = 0; i < 4; i++) {
		int channel = i % 2;
		CheckpointType type = channel == 0 ? CHECKPOINT : SAVEPOINT;
		barrierHandler.processBarrier(new CheckpointBarrier(i, 0, new CheckpointOptions(type, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, channel));
		assertEquals(type.isSavepoint(), alignedHandler.isCheckpointPending());
		assertNotEquals(alignedHandler.isCheckpointPending(), unalignedHandler.isCheckpointPending());

		if (!type.isSavepoint()) {
			assertFalse(barrierHandler.getAllBarriersReceivedFuture(i).isDone());
			assertInflightDataEquals(unalignedHandler, barrierHandler, i, inputGate.getNumberOfInputChannels());
		}
	}
}
 

@Test
public void testOutOfOrderBarrier() throws Exception {
	SingleInputGate inputGate = new SingleInputGateBuilder().setNumberOfChannels(2).build();
	inputGate.setInputChannels(new TestInputChannel(inputGate, 0), new TestInputChannel(inputGate, 1));
	TestInvokable target = new TestInvokable();
	CheckpointBarrierAligner alignedHandler = new CheckpointBarrierAligner("test", target, inputGate);
	CheckpointBarrierUnaligner unalignedHandler = new CheckpointBarrierUnaligner(TestSubtaskCheckpointCoordinator.INSTANCE, "test", target, inputGate);
	AlternatingCheckpointBarrierHandler barrierHandler = new AlternatingCheckpointBarrierHandler(alignedHandler, unalignedHandler, target);

	long checkpointId = 10;
	long outOfOrderSavepointId = 5;
	long initialAlignedCheckpointId = alignedHandler.getLatestCheckpointId();

	barrierHandler.processBarrier(new CheckpointBarrier(checkpointId, 0, new CheckpointOptions(CHECKPOINT, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, 0));
	barrierHandler.processBarrier(new CheckpointBarrier(outOfOrderSavepointId, 0, new CheckpointOptions(SAVEPOINT, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, 1));

	assertEquals(checkpointId, barrierHandler.getLatestCheckpointId());
	assertInflightDataEquals(unalignedHandler, barrierHandler, checkpointId, inputGate.getNumberOfInputChannels());
	assertEquals(initialAlignedCheckpointId, alignedHandler.getLatestCheckpointId());
}
 

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

/**
 * 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, 2);
	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder().buildRemoteChannel(inputGate);
	try {
		inputGate.setInputChannels(inputChannel);
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

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

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

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertEquals(2, inputChannel.getSenderBacklog());
	} finally {
		releaseResource(inputGate, networkBufferPool);
	}
}
 

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

@Test
public void testHasInflightDataBeforeProcessBarrier() throws Exception {
	SingleInputGate inputGate = new SingleInputGateBuilder().setNumberOfChannels(2).build();
	inputGate.setInputChannels(new TestInputChannel(inputGate, 0), new TestInputChannel(inputGate, 1));
	TestInvokable target = new TestInvokable();
	CheckpointBarrierAligner alignedHandler = new CheckpointBarrierAligner("test", target, inputGate);
	CheckpointBarrierUnaligner unalignedHandler = new CheckpointBarrierUnaligner(TestSubtaskCheckpointCoordinator.INSTANCE, "test", target, inputGate);
	AlternatingCheckpointBarrierHandler barrierHandler = new AlternatingCheckpointBarrierHandler(alignedHandler, unalignedHandler, target);

	final long id = 1;
	unalignedHandler.processBarrier(new CheckpointBarrier(id, 0, new CheckpointOptions(CHECKPOINT, CheckpointStorageLocationReference.getDefault())), new InputChannelInfo(0, 0));

	assertInflightDataEquals(unalignedHandler, barrierHandler, id, inputGate.getNumberOfInputChannels());
	assertFalse(barrierHandler.getAllBarriersReceivedFuture(id).isDone());
}
 

private void testBarrierHandling(CheckpointType checkpointType) throws Exception {
	final long barrierId = 123L;
	TestInvokable target = new TestInvokable();
	SingleInputGate gate = new SingleInputGateBuilder().setNumberOfChannels(2).build();
	TestInputChannel fast = new TestInputChannel(gate, 0, false, true);
	TestInputChannel slow = new TestInputChannel(gate, 1, false, true);
	gate.setInputChannels(fast, slow);
	AlternatingCheckpointBarrierHandler barrierHandler = barrierHandler(gate, target);
	CheckpointedInputGate checkpointedGate = new CheckpointedInputGate(gate, barrierHandler  /* offset */);

	sendBarrier(barrierId, checkpointType, fast, checkpointedGate);

	assertEquals(checkpointType.isSavepoint(), target.triggeredCheckpoints.isEmpty());
	assertEquals(checkpointType.isSavepoint(), barrierHandler.isBlocked(fast.getChannelInfo()));
	assertFalse(barrierHandler.isBlocked(slow.getChannelInfo()));

	sendBarrier(barrierId, checkpointType, slow, checkpointedGate);

	assertEquals(singletonList(barrierId), target.triggeredCheckpoints);
	for (InputChannel channel : gate.getInputChannels().values()) {
		assertFalse(barrierHandler.isBlocked(channel.getChannelInfo()));
		assertEquals(
			String.format("channel %d should be resumed", channel.getChannelIndex()),
			checkpointType.isSavepoint(),
			((TestInputChannel) channel).isResumed());
	}
}
 

private static CheckpointedInputGate buildGate(TestInvokable target, int numChannels) {
	SingleInputGate gate = new SingleInputGateBuilder().setNumberOfChannels(numChannels).build();
	TestInputChannel[] channels = new TestInputChannel[numChannels];
	for (int i = 0; i < numChannels; i++) {
		channels[i] = new TestInputChannel(gate, i, false, true);
	}
	gate.setInputChannels(channels);
	return new CheckpointedInputGate(gate, barrierHandler(gate, target));
}
 

/**
 * Verifies that {@link BufferResponse} of compressed {@link Buffer} can be handled correctly.
 */
@Test
public void testReceiveCompressedBuffer() throws Exception {
	int bufferSize = 1024;
	String compressionCodec = "LZ4";
	BufferCompressor compressor = new BufferCompressor(bufferSize, compressionCodec);
	BufferDecompressor decompressor = new BufferDecompressor(bufferSize, compressionCodec);
	NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, bufferSize, 2);
	SingleInputGate inputGate = new SingleInputGateBuilder()
		.setBufferDecompressor(decompressor)
		.setSegmentProvider(networkBufferPool)
		.build();
	RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, null);
	inputGate.setInputChannels(inputChannel);

	try {
		BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

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

		Buffer buffer = compressor.compressToOriginalBuffer(TestBufferFactory.createBuffer(bufferSize));
		BufferResponse bufferResponse = createBufferResponse(
			buffer,
			0,
			inputChannel.getInputChannelId(),
			2,
			new NetworkBufferAllocator(handler));
		assertTrue(bufferResponse.isCompressed);
		handler.channelRead(null, bufferResponse);

		Buffer receivedBuffer = inputChannel.getNextReceivedBuffer();
		assertNotNull(receivedBuffer);
		assertTrue(receivedBuffer.isCompressed());
		receivedBuffer.recycleBuffer();
	} finally {
		releaseResource(inputGate, networkBufferPool);
	}
}
 

/**
 * 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 NettyPartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, 2);
	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);

		// 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,
			new NetworkBufferAllocator(handler));
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(2, inputChannel.getUnannouncedCredit());

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

		// 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 {
		releaseResource(inputGate, networkBufferPool);
		channel.close();
	}
}
 

private void testReadBufferResponseWithReleasingOrRemovingChannel(
	boolean isRemoved,
	boolean readBeforeReleasingOrRemoving) throws Exception {

	int bufferSize = 1024;

	NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, bufferSize, 2);
	SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	RemoteInputChannel inputChannel = new InputChannelBuilder()
		.buildRemoteChannel(inputGate);
	inputGate.setInputChannels(inputChannel);
	inputGate.assignExclusiveSegments();

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

	try {
		if (!readBeforeReleasingOrRemoving) {
			// Release the channel.
			inputGate.close();
			if (isRemoved) {
				handler.removeInputChannel(inputChannel);
			}
		}

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

		if (readBeforeReleasingOrRemoving) {
			// Release the channel.
			inputGate.close();
			if (isRemoved) {
				handler.removeInputChannel(inputChannel);
			}
		}

		handler.channelRead(null, bufferResponse);

		assertEquals(0, inputChannel.getNumberOfQueuedBuffers());
		if (!readBeforeReleasingOrRemoving) {
			assertNull(bufferResponse.getBuffer());
		} else {
			assertNotNull(bufferResponse.getBuffer());
			assertTrue(bufferResponse.getBuffer().isRecycled());
		}

		embeddedChannel.runScheduledPendingTasks();
		NettyMessage.CancelPartitionRequest cancelPartitionRequest = embeddedChannel.readOutbound();
		assertNotNull(cancelPartitionRequest);
		assertEquals(inputChannel.getInputChannelId(), cancelPartitionRequest.receiverId);
	} finally {
		releaseResource(inputGate, networkBufferPool);
		embeddedChannel.close();
	}
}
 

public static void setupInputGate(SingleInputGate gate, InputChannel... channels) throws IOException {
	gate.setInputChannels(channels);
	gate.setup();
	gate.requestPartitions();
}
 

@Test
public void testSnapshotWithTwoInputGates() throws Exception {
	SingleInputGate inputGate1 = new SingleInputGateBuilder().setSingleInputGateIndex(0).build();
	RemoteInputChannel channel1 = InputChannelBuilder.newBuilder().buildRemoteChannel(inputGate1);
	inputGate1.setInputChannels(channel1);

	SingleInputGate inputGate2 = new SingleInputGateBuilder().setSingleInputGateIndex(1).build();
	RemoteInputChannel channel2 = InputChannelBuilder.newBuilder().buildRemoteChannel(inputGate2);
	inputGate2.setInputChannels(channel2);

	CheckpointBarrierUnaligner unaligner = new CheckpointBarrierUnaligner(
		TestSubtaskCheckpointCoordinator.INSTANCE,
		"test",
		new DummyCheckpointInvokable(),
		inputGate1,
		inputGate2);
	inputGate1.registerBufferReceivedListener(unaligner.getBufferReceivedListener().get());
	inputGate2.registerBufferReceivedListener(unaligner.getBufferReceivedListener().get());

	StreamTaskNetworkInput<Long> input1 = createInput(unaligner, inputGate1);
	StreamTaskNetworkInput<Long> input2 = createInput(unaligner, inputGate2);

	CheckpointBarrier barrier = new CheckpointBarrier(0, 0L, CheckpointOptions.forCheckpointWithDefaultLocation());
	channel1.onBuffer(EventSerializer.toBuffer(barrier), 0, 0);
	channel1.onBuffer(BufferBuilderTestUtils.buildSomeBuffer(1), 1, 0);

	// all records on inputGate2 are now in-flight
	channel2.onBuffer(BufferBuilderTestUtils.buildSomeBuffer(2), 0, 0);
	channel2.onBuffer(BufferBuilderTestUtils.buildSomeBuffer(3), 1, 0);

	// now snapshot all inflight buffers
	RecordingChannelStateWriter channelStateWriter = new RecordingChannelStateWriter();
	channelStateWriter.start(0, CheckpointOptions.forCheckpointWithDefaultLocation());
	CompletableFuture<Void> completableFuture1 = input1.prepareSnapshot(channelStateWriter, 0);
	CompletableFuture<Void> completableFuture2 = input2.prepareSnapshot(channelStateWriter, 0);

	// finish unaligned checkpoint on input side
	channel2.onBuffer(EventSerializer.toBuffer(barrier), 2, 0);

	// futures should be completed
	completableFuture1.join();
	completableFuture2.join();

	assertEquals(channelStateWriter.getAddedInput().get(channel1.getChannelInfo()), Collections.emptyList());
	List<Buffer> storedBuffers = channelStateWriter.getAddedInput().get(channel2.getChannelInfo());
	assertEquals(Arrays.asList(2, 3), storedBuffers.stream().map(Buffer::getSize).collect(Collectors.toList()));
}