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

/**
 * Tests that input gate can successfully convert unknown input channels into local and remote
 * channels.
 */
@Test
public void testUpdateUnknownInputChannel() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();

	final ResultPartition localResultPartition = new ResultPartitionBuilder()
		.setResultPartitionManager(network.getResultPartitionManager())
		.setupBufferPoolFactoryFromNettyShuffleEnvironment(network)
		.build();

	final ResultPartition remoteResultPartition = new ResultPartitionBuilder()
		.setResultPartitionManager(network.getResultPartitionManager())
		.setupBufferPoolFactoryFromNettyShuffleEnvironment(network)
		.build();

	localResultPartition.setup();
	remoteResultPartition.setup();

	final SingleInputGate inputGate = createInputGate(network, 2, ResultPartitionType.PIPELINED);

	try {
		final ResultPartitionID localResultPartitionId = localResultPartition.getPartitionId();
		addUnknownInputChannel(network, inputGate, localResultPartitionId, 0);

		final ResultPartitionID remoteResultPartitionId = remoteResultPartition.getPartitionId();
		addUnknownInputChannel(network, inputGate, remoteResultPartitionId, 1);

		inputGate.setup();

		assertThat(inputGate.getInputChannels().get(remoteResultPartitionId.getPartitionId()),
			is(instanceOf((UnknownInputChannel.class))));
		assertThat(inputGate.getInputChannels().get(localResultPartitionId.getPartitionId()),
			is(instanceOf((UnknownInputChannel.class))));

		ResourceID localLocation = ResourceID.generate();

		// Trigger updates to remote input channel from unknown input channel
		inputGate.updateInputChannel(
			localLocation,
			createRemoteWithIdAndLocation(remoteResultPartitionId.getPartitionId(), ResourceID.generate()));

		assertThat(inputGate.getInputChannels().get(remoteResultPartitionId.getPartitionId()),
			is(instanceOf((RemoteInputChannel.class))));
		assertThat(inputGate.getInputChannels().get(localResultPartitionId.getPartitionId()),
			is(instanceOf((UnknownInputChannel.class))));

		// Trigger updates to local input channel from unknown input channel
		inputGate.updateInputChannel(
			localLocation,
			createRemoteWithIdAndLocation(localResultPartitionId.getPartitionId(), localLocation));

		assertThat(inputGate.getInputChannels().get(remoteResultPartitionId.getPartitionId()),
			is(instanceOf((RemoteInputChannel.class))));
		assertThat(inputGate.getInputChannels().get(localResultPartitionId.getPartitionId()),
			is(instanceOf((LocalInputChannel.class))));
	} finally {
		inputGate.close();
		network.close();
	}
}
 

@Test
public void testQueuedBuffers() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();

	final ResultPartition resultPartition = new ResultPartitionBuilder()
		.setResultPartitionManager(network.getResultPartitionManager())
		.setupBufferPoolFactoryFromNettyShuffleEnvironment(network)
		.build();

	final SingleInputGate inputGate = createInputGate(network, 2, ResultPartitionType.PIPELINED);

	final ResultPartitionID localResultPartitionId = resultPartition.getPartitionId();

	final RemoteInputChannel remoteInputChannel = InputChannelBuilder.newBuilder()
		.setChannelIndex(1)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildRemoteAndSetToGate(inputGate);

	InputChannelBuilder.newBuilder()
		.setChannelIndex(0)
		.setPartitionId(localResultPartitionId)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildLocalAndSetToGate(inputGate);

	try {
		resultPartition.setup();
		inputGate.setup();

		remoteInputChannel.onBuffer(TestBufferFactory.createBuffer(1), 0, 0);
		assertEquals(1, inputGate.getNumberOfQueuedBuffers());

		resultPartition.addBufferConsumer(BufferBuilderTestUtils.createFilledBufferConsumer(1), 0);
		assertEquals(2, inputGate.getNumberOfQueuedBuffers());
	} finally {
		resultPartition.release();
		inputGate.close();
		network.close();
	}
}
 

@Test
public void testEmitRecordWithPartitionStateRecovery() throws Exception {
	final int totalBuffers = 10; // enough for both states and normal records
	final int totalStates = 2;
	final int[] states = {1, 2, 3, 4};
	final int[] records = {5, 6, 7, 8};
	final int bufferSize = states.length * Integer.BYTES;

	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, bufferSize, 1);
	final ChannelStateReader stateReader = new ResultPartitionTest.FiniteChannelStateReader(totalStates, states);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final RecordWriter<IntValue> recordWriter = new RecordWriterBuilder<IntValue>().build(partition);

	try {
		partition.setup();
		partition.readRecoveredState(stateReader);

		for (int record: records) {
			// the record length 4 is also written into buffer for every emit
			recordWriter.broadcastEmit(new IntValue(record));
		}

		// every buffer can contain 2 int records with 2 int length(4)
		final int[][] expectedRecordsInBuffer = {{4, 5, 4, 6}, {4, 7, 4, 8}};

		for (ResultSubpartition subpartition : partition.getAllPartitions()) {
			// create the view to consume all the buffers with states and records
			final ResultSubpartitionView view = new PipelinedSubpartitionView(
				(PipelinedSubpartition) subpartition,
				new NoOpBufferAvailablityListener());

			int numConsumedBuffers = 0;
			ResultSubpartition.BufferAndBacklog bufferAndBacklog;
			while ((bufferAndBacklog = view.getNextBuffer()) != null) {
				Buffer buffer = bufferAndBacklog.buffer();
				int[] expected = numConsumedBuffers < totalStates ? states : expectedRecordsInBuffer[numConsumedBuffers - totalStates];
				BufferBuilderAndConsumerTest.assertContent(
					buffer,
					partition.getBufferPool()
						.getSubpartitionBufferRecyclers()[subpartition.getSubPartitionIndex()],
					expected);

				buffer.recycleBuffer();
				numConsumedBuffers++;
			}

			assertEquals(totalStates + expectedRecordsInBuffer.length, numConsumedBuffers);
		}
	} finally {
		// cleanup
		globalPool.destroyAllBufferPools();
		globalPool.destroy();
	}
}
 

/**
 * Tests that input gate can successfully convert unknown input channels into local and remote
 * channels.
 */
@Test
public void testUpdateUnknownInputChannel() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();

	final ResultPartition localResultPartition = new ResultPartitionBuilder()
		.setResultPartitionManager(network.getResultPartitionManager())
		.setupBufferPoolFactoryFromNettyShuffleEnvironment(network)
		.build();

	final ResultPartition remoteResultPartition = new ResultPartitionBuilder()
		.setResultPartitionManager(network.getResultPartitionManager())
		.setupBufferPoolFactoryFromNettyShuffleEnvironment(network)
		.build();

	localResultPartition.setup();
	remoteResultPartition.setup();

	final SingleInputGate inputGate = createInputGate(network, 2, ResultPartitionType.PIPELINED);
	final InputChannel[] inputChannels = new InputChannel[2];

	try {
		final ResultPartitionID localResultPartitionId = localResultPartition.getPartitionId();
		inputChannels[0] = buildUnknownInputChannel(network, inputGate, localResultPartitionId, 0);

		final ResultPartitionID remoteResultPartitionId = remoteResultPartition.getPartitionId();
		inputChannels[1] = buildUnknownInputChannel(network, inputGate, remoteResultPartitionId, 1);

		inputGate.setInputChannels(inputChannels);
		inputGate.setup();

		assertThat(inputGate.getInputChannels().get(remoteResultPartitionId.getPartitionId()),
			is(instanceOf((UnknownInputChannel.class))));
		assertThat(inputGate.getInputChannels().get(localResultPartitionId.getPartitionId()),
			is(instanceOf((UnknownInputChannel.class))));

		ResourceID localLocation = ResourceID.generate();

		// Trigger updates to remote input channel from unknown input channel
		inputGate.updateInputChannel(
			localLocation,
			createRemoteWithIdAndLocation(remoteResultPartitionId.getPartitionId(), ResourceID.generate()));

		assertThat(inputGate.getInputChannels().get(remoteResultPartitionId.getPartitionId()),
			is(instanceOf((RemoteInputChannel.class))));
		assertThat(inputGate.getInputChannels().get(localResultPartitionId.getPartitionId()),
			is(instanceOf((UnknownInputChannel.class))));

		// Trigger updates to local input channel from unknown input channel
		inputGate.updateInputChannel(
			localLocation,
			createRemoteWithIdAndLocation(localResultPartitionId.getPartitionId(), localLocation));

		assertThat(inputGate.getInputChannels().get(remoteResultPartitionId.getPartitionId()),
			is(instanceOf((RemoteInputChannel.class))));
		assertThat(inputGate.getInputChannels().get(localResultPartitionId.getPartitionId()),
			is(instanceOf((LocalInputChannel.class))));
	} finally {
		inputGate.close();
		network.close();
	}
}
 

@Test
public void testQueuedBuffers() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();

	final ResultPartition resultPartition = new ResultPartitionBuilder()
		.setResultPartitionManager(network.getResultPartitionManager())
		.setupBufferPoolFactoryFromNettyShuffleEnvironment(network)
		.build();

	final SingleInputGate inputGate = createInputGate(network, 2, ResultPartitionType.PIPELINED);

	final ResultPartitionID localResultPartitionId = resultPartition.getPartitionId();
	final InputChannel[] inputChannels = new InputChannel[2];

	final RemoteInputChannel remoteInputChannel = InputChannelBuilder.newBuilder()
		.setChannelIndex(1)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildRemoteChannel(inputGate);
	inputChannels[0] = remoteInputChannel;

	inputChannels[1] = InputChannelBuilder.newBuilder()
		.setChannelIndex(0)
		.setPartitionId(localResultPartitionId)
		.setupFromNettyShuffleEnvironment(network)
		.setConnectionManager(new TestingConnectionManager())
		.buildLocalChannel(inputGate);

	try {
		resultPartition.setup();
		setupInputGate(inputGate, inputChannels);

		remoteInputChannel.onBuffer(createBuffer(1), 0, 0);
		assertEquals(1, inputGate.getNumberOfQueuedBuffers());

		resultPartition.addBufferConsumer(BufferBuilderTestUtils.createFilledFinishedBufferConsumer(1), 0);
		assertEquals(2, inputGate.getNumberOfQueuedBuffers());
	} finally {
		resultPartition.release();
		inputGate.close();
		network.close();
	}
}
 

/**
 * Tests the consumption of multiple subpartitions via local input channels.
 *
 * <p>Multiple producer tasks produce pipelined partitions, which are consumed by multiple
 * tasks via local input channels.
 */
@Test
public void testConcurrentConsumeMultiplePartitions() throws Exception {
	// Config
	final int parallelism = 32;
	final int producerBufferPoolSize = parallelism + 1;
	final int numberOfBuffersPerChannel = 1024;

	checkArgument(parallelism >= 1);
	checkArgument(producerBufferPoolSize >= parallelism);
	checkArgument(numberOfBuffersPerChannel >= 1);

	// Setup
	// One thread per produced partition and one per consumer
	final ExecutorService executor = Executors.newFixedThreadPool(2 * parallelism);

	final NetworkBufferPool networkBuffers = new NetworkBufferPool(
		(parallelism * producerBufferPoolSize) + (parallelism * parallelism),
		TestBufferFactory.BUFFER_SIZE, 1);

	final ResultPartitionManager partitionManager = new ResultPartitionManager();

	final ResultPartitionID[] partitionIds = new ResultPartitionID[parallelism];
	final TestPartitionProducer[] partitionProducers = new TestPartitionProducer[parallelism];

	// Create all partitions
	for (int i = 0; i < parallelism; i++) {
		partitionIds[i] = new ResultPartitionID();

		final ResultPartition partition = new ResultPartitionBuilder()
			.setResultPartitionId(partitionIds[i])
			.setNumberOfSubpartitions(parallelism)
			.setNumTargetKeyGroups(parallelism)
			.setResultPartitionManager(partitionManager)
			.setBufferPoolFactory(p ->
				networkBuffers.createBufferPool(producerBufferPoolSize, producerBufferPoolSize))
			.build();

		// Create a buffer pool for this partition
		partition.setup();

		// Create the producer
		partitionProducers[i] = new TestPartitionProducer(
			partition,
			false,
			new TestPartitionProducerBufferSource(
				parallelism,
				partition.getBufferPool(),
				numberOfBuffersPerChannel)
		);
	}

	// Test
	try {
		// Submit producer tasks
		List<CompletableFuture<?>> results = Lists.newArrayListWithCapacity(
			parallelism + 1);

		for (int i = 0; i < parallelism; i++) {
			results.add(CompletableFuture.supplyAsync(
				CheckedSupplier.unchecked(partitionProducers[i]::call), executor));
		}

		// Submit consumer
		for (int i = 0; i < parallelism; i++) {
			final TestLocalInputChannelConsumer consumer = new TestLocalInputChannelConsumer(
				i,
				parallelism,
				numberOfBuffersPerChannel,
				networkBuffers.createBufferPool(parallelism, parallelism),
				partitionManager,
				new TaskEventDispatcher(),
				partitionIds);

			results.add(CompletableFuture.supplyAsync(CheckedSupplier.unchecked(consumer::call), executor));
		}

		FutureUtils.waitForAll(results).get();
	}
	finally {
		networkBuffers.destroyAllBufferPools();
		networkBuffers.destroy();
		executor.shutdown();
	}
}