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

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 testFairConsumptionLocalChannelsPreFilled() throws Exception {
	final int numberOfChannels = 37;
	final int buffersPerChannel = 27;

	final ResultPartition resultPartition = mock(ResultPartition.class);
	final BufferConsumer bufferConsumer = createFilledBufferConsumer(42);

	// ----- create some source channels and fill them with buffers -----

	final PipelinedSubpartition[] sources = new PipelinedSubpartition[numberOfChannels];

	for (int i = 0; i < numberOfChannels; i++) {
		PipelinedSubpartition partition = new PipelinedSubpartition(0, resultPartition);

		for (int p = 0; p < buffersPerChannel; p++) {
			partition.add(bufferConsumer.copy());
		}

		partition.finish();
		sources[i] = partition;
	}

	// ----- create reading side -----

	ResultPartitionManager resultPartitionManager = createResultPartitionManager(sources);

	final SingleInputGate gate = createFairnessVerifyingInputGate(numberOfChannels);

	for (int i = 0; i < numberOfChannels; i++) {
		createLocalInputChannel(gate, i, resultPartitionManager);
	}

	gate.setup();

	// read all the buffers and the EOF event
	for (int i = numberOfChannels * (buffersPerChannel + 1); i > 0; --i) {
		assertNotNull(gate.getNext());

		int min = Integer.MAX_VALUE;
		int max = 0;

		for (PipelinedSubpartition source : sources) {
			int size = source.getCurrentNumberOfBuffers();
			min = Math.min(min, size);
			max = Math.max(max, size);
		}

		assertTrue(max == min || max == (min + 1));
	}

	assertFalse(gate.getNext().isPresent());
}
 

@Test
public void testFairConsumptionLocalChannels() throws Exception {
	final int numberOfChannels = 37;
	final int buffersPerChannel = 27;

	final ResultPartition resultPartition = mock(ResultPartition.class);
	try (BufferConsumer bufferConsumer = createFilledBufferConsumer(42)) {

		// ----- create some source channels and fill them with one buffer each -----

		final PipelinedSubpartition[] sources = new PipelinedSubpartition[numberOfChannels];

		for (int i = 0; i < numberOfChannels; i++) {
			sources[i] = new PipelinedSubpartition(0, resultPartition);
		}

		// ----- create reading side -----

		ResultPartitionManager resultPartitionManager = createResultPartitionManager(sources);

		final SingleInputGate gate = createFairnessVerifyingInputGate(numberOfChannels);

		for (int i = 0; i < numberOfChannels; i++) {
			createLocalInputChannel(gate, i, resultPartitionManager);
		}

		// seed one initial buffer
		sources[12].add(bufferConsumer.copy());

		gate.setup();

		// read all the buffers and the EOF event
		for (int i = 0; i < numberOfChannels * buffersPerChannel; i++) {
			assertNotNull(gate.getNext());

			int min = Integer.MAX_VALUE;
			int max = 0;

			for (PipelinedSubpartition source : sources) {
				int size = source.getCurrentNumberOfBuffers();
				min = Math.min(min, size);
				max = Math.max(max, size);
			}

			assertTrue(max == min || max == min + 1);

			if (i % (2 * numberOfChannels) == 0) {
				// add three buffers to each channel, in random order
				fillRandom(sources, 3, bufferConsumer);
			}
		}
		// there is still more in the queues
	}
}
 

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