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

public ResultPartition build() {
	ResultPartitionFactory resultPartitionFactory = new ResultPartitionFactory(
		partitionManager,
		channelManager,
		networkBufferPool,
		blockingSubpartitionType,
		networkBuffersPerChannel,
		floatingNetworkBuffersPerGate,
		networkBufferSize,
		releasedOnConsumption);

	FunctionWithException<BufferPoolOwner, BufferPool, IOException> factory = bufferPoolFactory.orElseGet(() ->
		resultPartitionFactory.createBufferPoolFactory(numberOfSubpartitions, partitionType));

	return resultPartitionFactory.create(
		"Result Partition task",
		partitionId,
		partitionType,
		numberOfSubpartitions,
		numTargetKeyGroups,
		factory);
}
 

public ResultPartition(
	String owningTaskName,
	ResultPartitionID partitionId,
	ResultPartitionType partitionType,
	ResultSubpartition[] subpartitions,
	int numTargetKeyGroups,
	ResultPartitionManager partitionManager,
	FunctionWithException<BufferPoolOwner, BufferPool, IOException> bufferPoolFactory) {

	this.owningTaskName = checkNotNull(owningTaskName);
	this.partitionId = checkNotNull(partitionId);
	this.partitionType = checkNotNull(partitionType);
	this.subpartitions = checkNotNull(subpartitions);
	this.numTargetKeyGroups = numTargetKeyGroups;
	this.partitionManager = checkNotNull(partitionManager);
	this.bufferPoolFactory = bufferPoolFactory;
}
 

@Override
public Float getValue() {
	int usedBuffers = 0;
	int bufferPoolSize = 0;

	for (ResultPartition resultPartition : resultPartitions) {
		BufferPool bufferPool = resultPartition.getBufferPool();

		if (bufferPool != null) {
			usedBuffers += bufferPool.bestEffortGetNumOfUsedBuffers();
			bufferPoolSize += bufferPool.getNumBuffers();
		}
	}

	if (bufferPoolSize != 0) {
		return ((float) usedBuffers) / bufferPoolSize;
	} else {
		return 0.0f;
	}
}
 

ReleaseOnConsumptionResultPartition(
		String owningTaskName,
		int partitionIndex,
		ResultPartitionID partitionId,
		ResultPartitionType partitionType,
		ResultSubpartition[] subpartitions,
		int numTargetKeyGroups,
		ResultPartitionManager partitionManager,
		@Nullable BufferCompressor bufferCompressor,
		FunctionWithException<BufferPoolOwner, BufferPool, IOException> bufferPoolFactory) {
	super(
		owningTaskName,
		partitionIndex,
		partitionId,
		partitionType,
		subpartitions,
		numTargetKeyGroups,
		partitionManager,
		bufferCompressor,
		bufferPoolFactory);

	this.consumedSubpartitions = new boolean[subpartitions.length];
	this.numUnconsumedSubpartitions = subpartitions.length;
}
 

Buffer requestBufferBlocking() throws IOException, InterruptedException {
	synchronized (bufferQueue) {
		Buffer buffer;
		while ((buffer = bufferQueue.takeBuffer()) == null) {
			if (inputChannel.isReleased()) {
				throw new CancelTaskException("Input channel [" + inputChannel.channelInfo + "] has already been released.");
			}
			if (!isWaitingForFloatingBuffers) {
				BufferPool bufferPool = inputChannel.inputGate.getBufferPool();
				buffer = bufferPool.requestBuffer();
				if (buffer == null && shouldContinueRequest(bufferPool)) {
					continue;
				}
			}

			if (buffer != null) {
				return buffer;
			}
			bufferQueue.wait();
		}
		return buffer;
	}
}
 

/**
 * The minimum pool size should be <code>numberOfSubpartitions + 1</code> for two considerations:
 *
 * <p>1. StreamTask can only process input if there is at-least one available buffer on output side, so it might cause
 * stuck problem if the minimum pool size is exactly equal to the number of subpartitions, because every subpartition
 * might maintain a partial unfilled buffer.
 *
 * <p>2. Increases one more buffer for every output LocalBufferPool to void performance regression if processing input is
 * based on at-least one buffer available on output side.
 */
@VisibleForTesting
FunctionWithException<BufferPoolOwner, BufferPool, IOException> createBufferPoolFactory(
		int numberOfSubpartitions,
		ResultPartitionType type) {
	return bufferPoolOwner -> {
		int maxNumberOfMemorySegments = type.isBounded() ?
			numberOfSubpartitions * networkBuffersPerChannel + floatingNetworkBuffersPerGate : Integer.MAX_VALUE;
		// If the partition type is back pressure-free, we register with the buffer pool for
		// callbacks to release memory.
		return bufferPoolFactory.createBufferPool(
			numberOfSubpartitions + 1,
			maxNumberOfMemorySegments,
			type.hasBackPressure() ? null : bufferPoolOwner,
			numberOfSubpartitions,
			maxBuffersPerChannel);
	};
}
 

/**
 * Requests floating buffers from the buffer pool based on the given required amount, and returns the actual
 * requested amount. If the required amount is not fully satisfied, it will register as a listener.
 */
int requestFloatingBuffers(int numRequired) throws IOException {
	int numRequestedBuffers = 0;
	synchronized (bufferQueue) {
		// Similar to notifyBufferAvailable(), make sure that we never add a buffer after channel
		// released all buffers via releaseAllResources().
		if (inputChannel.isReleased()) {
			return numRequestedBuffers;
		}

		numRequiredBuffers = numRequired;

		while (bufferQueue.getAvailableBufferSize() < numRequiredBuffers && !isWaitingForFloatingBuffers) {
			BufferPool bufferPool = inputChannel.inputGate.getBufferPool();
			Buffer buffer = bufferPool.requestBuffer();
			if (buffer != null) {
				bufferQueue.addFloatingBuffer(buffer);
				numRequestedBuffers++;
			} else if (bufferPool.addBufferListener(this)) {
				isWaitingForFloatingBuffers = true;
				break;
			}
		}
	}
	return numRequestedBuffers;
}
 

@Override
public Float getValue() {
	int usedBuffers = 0;
	int bufferPoolSize = 0;

	for (ResultPartition resultPartition : resultPartitions) {
		BufferPool bufferPool = resultPartition.getBufferPool();

		if (bufferPool != null) {
			usedBuffers += bufferPool.bestEffortGetNumOfUsedBuffers();
			bufferPoolSize += bufferPool.getNumBuffers();
		}
	}

	if (bufferPoolSize != 0) {
		return ((float) usedBuffers) / bufferPoolSize;
	} else {
		return 0.0f;
	}
}
 

public ResultPartition build() {
	ResultPartitionFactory resultPartitionFactory = new ResultPartitionFactory(
		partitionManager,
		channelManager,
		networkBufferPool,
		blockingSubpartitionType,
		networkBuffersPerChannel,
		floatingNetworkBuffersPerGate,
		networkBufferSize,
		releasedOnConsumption,
		blockingShuffleCompressionEnabled,
		compressionCodec,
		maxBuffersPerChannel);

	FunctionWithException<BufferPoolOwner, BufferPool, IOException> factory = bufferPoolFactory.orElseGet(() ->
		resultPartitionFactory.createBufferPoolFactory(numberOfSubpartitions, partitionType));

	return resultPartitionFactory.create(
		"Result Partition task",
		partitionIndex,
		partitionId,
		partitionType,
		numberOfSubpartitions,
		numTargetKeyGroups,
		factory);
}
 

/**
 * Registers a buffer pool with this result partition.
 *
 * <p>There is one pool for each result partition, which is shared by all its sub partitions.
 *
 * <p>The pool is registered with the partition *after* it as been constructed in order to conform
 * to the life-cycle of task registrations in the {@link TaskManager}.
 */
public void registerBufferPool(BufferPool bufferPool) {
	checkArgument(bufferPool.getNumberOfRequiredMemorySegments() >= getNumberOfSubpartitions(),
			"Bug in result partition setup logic: Buffer pool has not enough guaranteed buffers for this result partition.");

	checkState(this.bufferPool == null, "Bug in result partition setup logic: Already registered buffer pool.");

	this.bufferPool = checkNotNull(bufferPool);
}
 

public RandomGeneratingInputGate(BufferPool[] bufferPools, BarrierGenerator[] barrierGens) {
	this.numberOfChannels = bufferPools.length;
	this.currentBarriers = new int[numberOfChannels];
	this.channelBlocked = new boolean[numberOfChannels];
	this.bufferPools = bufferPools;
	this.barrierGens = barrierGens;
	availabilityHelper.resetAvailable();
}
 

public SingleInputGate(
	String owningTaskName,
	IntermediateDataSetID consumedResultId,
	final ResultPartitionType consumedPartitionType,
	int consumedSubpartitionIndex,
	int numberOfInputChannels,
	PartitionProducerStateProvider partitionProducerStateProvider,
	boolean isCreditBased,
	SupplierWithException<BufferPool, IOException> bufferPoolFactory) {

	this.owningTaskName = checkNotNull(owningTaskName);

	this.consumedResultId = checkNotNull(consumedResultId);
	this.consumedPartitionType = checkNotNull(consumedPartitionType);
	this.bufferPoolFactory = checkNotNull(bufferPoolFactory);

	checkArgument(consumedSubpartitionIndex >= 0);
	this.consumedSubpartitionIndex = consumedSubpartitionIndex;

	checkArgument(numberOfInputChannels > 0);
	this.numberOfInputChannels = numberOfInputChannels;

	this.inputChannels = new HashMap<>(numberOfInputChannels);
	this.channelsWithEndOfPartitionEvents = new BitSet(numberOfInputChannels);
	this.enqueuedInputChannelsWithData = new BitSet(numberOfInputChannels);

	this.partitionProducerStateProvider = checkNotNull(partitionProducerStateProvider);

	this.isCreditBased = isCreditBased;

	this.closeFuture = new CompletableFuture<>();
}
 

/**
 * Creates an input gate and all of its input channels.
 */
public SingleInputGate create(
		@Nonnull String owningTaskName,
		@Nonnull InputGateDeploymentDescriptor igdd,
		@Nonnull PartitionProducerStateProvider partitionProducerStateProvider,
		@Nonnull InputChannelMetrics metrics) {
	SupplierWithException<BufferPool, IOException> bufferPoolFactory = createBufferPoolFactory(
		networkBufferPool,
		isCreditBased,
		networkBuffersPerChannel,
		floatingNetworkBuffersPerGate,
		igdd.getShuffleDescriptors().length,
		igdd.getConsumedPartitionType());

	SingleInputGate inputGate = new SingleInputGate(
		owningTaskName,
		igdd.getConsumedResultId(),
		igdd.getConsumedPartitionType(),
		igdd.getConsumedSubpartitionIndex(),
		igdd.getShuffleDescriptors().length,
		partitionProducerStateProvider,
		isCreditBased,
		bufferPoolFactory);

	createInputChannels(owningTaskName, igdd, inputGate, metrics);
	return inputGate;
}
 

@Test
public void testDoublePartitionRequest() throws Exception {
	final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new NettyPartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final int numExclusiveBuffers = 2;
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, numExclusiveBuffers);
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client, networkBufferPool);

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

@VisibleForTesting
public void setupPartition(ResultPartition partition) throws IOException {
	BufferPool bufferPool = null;

	try {
		int maxNumberOfMemorySegments = partition.getPartitionType().isBounded() ?
			partition.getNumberOfSubpartitions() * networkBuffersPerChannel +
				extraNetworkBuffersPerGate : Integer.MAX_VALUE;
		// If the partition type is back pressure-free, we register with the buffer pool for
		// callbacks to release memory.
		bufferPool = networkBufferPool.createBufferPool(partition.getNumberOfSubpartitions(),
			maxNumberOfMemorySegments,
			partition.getPartitionType().hasBackPressure() ? Optional.empty() : Optional.of(partition));

		partition.registerBufferPool(bufferPool);

		resultPartitionManager.registerResultPartition(partition);
	} catch (Throwable t) {
		if (bufferPool != null) {
			bufferPool.lazyDestroy();
		}

		if (t instanceof IOException) {
			throw (IOException) t;
		} else {
			throw new IOException(t.getMessage(), t);
		}
	}

	taskEventDispatcher.registerPartition(partition.getPartitionId());
}
 

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

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

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

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

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

@VisibleForTesting
FunctionWithException<BufferPoolOwner, BufferPool, IOException> createBufferPoolFactory(
		int numberOfSubpartitions,
		ResultPartitionType type) {
	return p -> {
		int maxNumberOfMemorySegments = type.isBounded() ?
			numberOfSubpartitions * networkBuffersPerChannel + floatingNetworkBuffersPerGate : Integer.MAX_VALUE;
		// If the partition type is back pressure-free, we register with the buffer pool for
		// callbacks to release memory.
		return bufferPoolFactory.createBufferPool(numberOfSubpartitions,
			maxNumberOfMemorySegments,
			type.hasBackPressure() ? Optional.empty() : Optional.of(p));
	};
}
 

ReleaseOnConsumptionResultPartition(
		String owningTaskName,
		ResultPartitionID partitionId,
		ResultPartitionType partitionType,
		ResultSubpartition[] subpartitions,
		int numTargetKeyGroups,
		ResultPartitionManager partitionManager,
		FunctionWithException<BufferPoolOwner, BufferPool, IOException> bufferPoolFactory) {
	super(owningTaskName, partitionId, partitionType, subpartitions, numTargetKeyGroups, partitionManager, bufferPoolFactory);

	this.consumedSubpartitions = new boolean[subpartitions.length];
	this.numUnconsumedSubpartitions = subpartitions.length;
}
 

@Override
public int calculateUsedBuffers(SingleInputGate inputGate) {
	int availableFloatingBuffers = 0;
	BufferPool bufferPool = inputGate.getBufferPool();
	if (bufferPool != null) {
		int requestedFloatingBuffers = bufferPool.bestEffortGetNumOfUsedBuffers();
		for (InputChannel ic : inputGate.getInputChannels().values()) {
			if (ic instanceof RemoteInputChannel) {
				availableFloatingBuffers += ((RemoteInputChannel) ic).unsynchronizedGetFloatingBuffersAvailable();
			}
		}
		return Math.max(0, requestedFloatingBuffers - availableFloatingBuffers);
	}
	return 0;
}
 

/**
 * 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);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder()
		.setMemorySegmentProvider(networkBufferPool)
		.buildRemoteAndSetToGate(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

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

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

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
	} 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);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder()
		.setMemorySegmentProvider(networkBufferPool)
		.buildRemoteAndSetToGate(inputGate);
	try {
		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);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertEquals(2, inputChannel.getSenderBacklog());
	} finally {
		// Release all the buffer resources
		inputGate.close();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

public SingleInputGate(
	String owningTaskName,
	int gateIndex,
	IntermediateDataSetID consumedResultId,
	final ResultPartitionType consumedPartitionType,
	int consumedSubpartitionIndex,
	int numberOfInputChannels,
	PartitionProducerStateProvider partitionProducerStateProvider,
	SupplierWithException<BufferPool, IOException> bufferPoolFactory,
	@Nullable BufferDecompressor bufferDecompressor,
	MemorySegmentProvider memorySegmentProvider) {

	this.owningTaskName = checkNotNull(owningTaskName);
	Preconditions.checkArgument(0 <= gateIndex, "The gate index must be positive.");
	this.gateIndex = gateIndex;

	this.consumedResultId = checkNotNull(consumedResultId);
	this.consumedPartitionType = checkNotNull(consumedPartitionType);
	this.bufferPoolFactory = checkNotNull(bufferPoolFactory);

	checkArgument(consumedSubpartitionIndex >= 0);
	this.consumedSubpartitionIndex = consumedSubpartitionIndex;

	checkArgument(numberOfInputChannels > 0);
	this.numberOfInputChannels = numberOfInputChannels;

	this.inputChannels = new HashMap<>(numberOfInputChannels);
	this.channels = new InputChannel[numberOfInputChannels];
	this.channelsWithEndOfPartitionEvents = new BitSet(numberOfInputChannels);
	this.enqueuedInputChannelsWithData = new BitSet(numberOfInputChannels);

	this.partitionProducerStateProvider = checkNotNull(partitionProducerStateProvider);

	this.bufferDecompressor = bufferDecompressor;
	this.memorySegmentProvider = checkNotNull(memorySegmentProvider);

	this.closeFuture = new CompletableFuture<>();
}
 

@VisibleForTesting
static SupplierWithException<BufferPool, IOException> createBufferPoolFactory(
		BufferPoolFactory bufferPoolFactory,
		int networkBuffersPerChannel,
		int floatingNetworkBuffersPerGate,
		int size,
		ResultPartitionType type) {
	// Note that we should guarantee at-least one floating buffer for local channel state recovery.
	return () -> bufferPoolFactory.createBufferPool(1, floatingNetworkBuffersPerGate);
}
 

/**
 * Creates an input gate and all of its input channels.
 */
public SingleInputGate create(
		@Nonnull String owningTaskName,
		int gateIndex,
		@Nonnull InputGateDeploymentDescriptor igdd,
		@Nonnull PartitionProducerStateProvider partitionProducerStateProvider,
		@Nonnull InputChannelMetrics metrics) {
	SupplierWithException<BufferPool, IOException> bufferPoolFactory = createBufferPoolFactory(
		networkBufferPool,
		networkBuffersPerChannel,
		floatingNetworkBuffersPerGate,
		igdd.getShuffleDescriptors().length,
		igdd.getConsumedPartitionType());

	BufferDecompressor bufferDecompressor = null;
	if (igdd.getConsumedPartitionType().isBlocking() && blockingShuffleCompressionEnabled) {
		bufferDecompressor = new BufferDecompressor(networkBufferSize, compressionCodec);
	}

	SingleInputGate inputGate = new SingleInputGate(
		owningTaskName,
		gateIndex,
		igdd.getConsumedResultId(),
		igdd.getConsumedPartitionType(),
		igdd.getConsumedSubpartitionIndex(),
		igdd.getShuffleDescriptors().length,
		partitionProducerStateProvider,
		bufferPoolFactory,
		bufferDecompressor,
		networkBufferPool);

	createInputChannels(owningTaskName, igdd, inputGate, metrics);
	return inputGate;
}
 

private boolean shouldContinueRequest(BufferPool bufferPool) {
	if (bufferPool.addBufferListener(this)) {
		isWaitingForFloatingBuffers = true;
		numRequiredBuffers = 1;
		return false;
	} else if (bufferPool.isDestroyed()) {
		throw new CancelTaskException("Local buffer pool has already been released.");
	} else {
		return true;
	}
}
 

public TestLocalInputChannelConsumer(
		int subpartitionIndex,
		int numberOfInputChannels,
		int numberOfExpectedBuffersPerChannel,
		BufferPool bufferPool,
		ResultPartitionManager partitionManager,
		TaskEventDispatcher taskEventDispatcher,
		ResultPartitionID[] consumedPartitionIds) throws IOException, InterruptedException {

	checkArgument(numberOfInputChannels >= 1);
	checkArgument(numberOfExpectedBuffersPerChannel >= 1);

	this.inputGate = new SingleInputGateBuilder()
		.setConsumedSubpartitionIndex(subpartitionIndex)
		.setNumberOfChannels(numberOfInputChannels)
		.setBufferPoolFactory(bufferPool)
		.build();

	// Setup input channels
	for (int i = 0; i < numberOfInputChannels; i++) {
		InputChannelBuilder.newBuilder()
			.setChannelIndex(i)
			.setPartitionManager(partitionManager)
			.setPartitionId(consumedPartitionIds[i])
			.setTaskEventPublisher(taskEventDispatcher)
			.buildLocalAndSetToGate(inputGate);
	}

	inputGate.setup();

	this.numberOfInputChannels = numberOfInputChannels;
	this.numberOfExpectedBuffersPerChannel = numberOfExpectedBuffersPerChannel;
}
 

@Override
public int calculateTotalBuffers(SingleInputGate inputGate) {
	BufferPool bufferPool = inputGate.getBufferPool();
	if (bufferPool != null) {
		return inputGate.getBufferPool().getNumBuffers();
	}
	return 0;
}
 

public RandomGeneratingInputGate(BufferPool[] bufferPools, BarrierGenerator[] barrierGens) {
	this.numberOfChannels = bufferPools.length;
	this.currentBarriers = new int[numberOfChannels];
	this.bufferPools = bufferPools;
	this.barrierGens = barrierGens;
	this.isAvailable = AVAILABLE;
}