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

NettyShuffleEnvironment(
		ResourceID taskExecutorResourceId,
		NettyShuffleEnvironmentConfiguration config,
		NetworkBufferPool networkBufferPool,
		ConnectionManager connectionManager,
		ResultPartitionManager resultPartitionManager,
		FileChannelManager fileChannelManager,
		ResultPartitionFactory resultPartitionFactory,
		SingleInputGateFactory singleInputGateFactory,
		Executor ioExecutor) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.config = config;
	this.networkBufferPool = networkBufferPool;
	this.connectionManager = connectionManager;
	this.resultPartitionManager = resultPartitionManager;
	this.inputGatesById = new ConcurrentHashMap<>(10);
	this.fileChannelManager = fileChannelManager;
	this.resultPartitionFactory = resultPartitionFactory;
	this.singleInputGateFactory = singleInputGateFactory;
	this.ioExecutor = ioExecutor;
	this.isClosed = false;
}
 

@Before
public void setUp() throws Exception {
	networkBufferPool = new NetworkBufferPool(100, 8192);
	testBufferPool = networkBufferPool.createBufferPool(1, Integer.MAX_VALUE);

	final Configuration configuration = new Configuration();
	configuration.setInteger(WebOptions.BACKPRESSURE_NUM_SAMPLES, BACKPRESSURE_NUM_SAMPLES);

	testingMiniCluster = new TestingMiniCluster(new TestingMiniClusterConfiguration.Builder()
		.setNumTaskManagers(JOB_PARALLELISM)
		.setConfiguration(configuration)
		.build());
	testingMiniCluster.start();

	dispatcherGateway = testingMiniCluster.getDispatcherGatewayFuture().get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
}
 

/**
 * Assign the exclusive buffers to all remote input channels directly for credit-based mode.
 *
 * @param networkBufferPool The global pool to request and recycle exclusive buffers
 * @param networkBuffersPerChannel The number of exclusive buffers for each channel
 */
public void assignExclusiveSegments(NetworkBufferPool networkBufferPool, int networkBuffersPerChannel) throws IOException {
	checkState(this.isCreditBased, "Bug in input gate setup logic: exclusive buffers only exist with credit-based flow control.");
	checkState(this.networkBufferPool == null, "Bug in input gate setup logic: global buffer pool has" +
		"already been set for this input gate.");

	this.networkBufferPool = checkNotNull(networkBufferPool);
	this.networkBuffersPerChannel = networkBuffersPerChannel;

	synchronized (requestLock) {
		for (InputChannel inputChannel : inputChannels.values()) {
			if (inputChannel instanceof RemoteInputChannel) {
				((RemoteInputChannel) inputChannel).assignExclusiveSegments(
					networkBufferPool.requestMemorySegments(networkBuffersPerChannel));
			}
		}
	}
}
 

@Test
public void testInitializeEmptyState() throws Exception {
	final int totalBuffers = 2;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 1, 1);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final ChannelStateReader stateReader = ChannelStateReader.NO_OP;
	try {
		partition.setup();
		partition.readRecoveredState(stateReader);

		for (ResultSubpartition subpartition : partition.getAllPartitions()) {
			// no buffers are added into the queue for empty states
			assertEquals(0, subpartition.getTotalNumberOfBuffers());
		}

		// destroy the local pool to verify that all the requested buffers by partition are recycled
		partition.getBufferPool().lazyDestroy();
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
	} finally {
		// cleanup
		globalPool.destroyAllBufferPools();
		globalPool.destroy();
	}
}
 

public NetworkEnvironment(
	int numBuffers,
	int memorySegmentSize,
	int partitionRequestInitialBackoff,
	int partitionRequestMaxBackoff,
	int networkBuffersPerChannel,
	int extraNetworkBuffersPerGate,
	boolean enableCreditBased) {
	this(
		new NetworkBufferPool(numBuffers, memorySegmentSize),
		new LocalConnectionManager(),
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new KvStateRegistry(),
		null,
		null,
		IOManager.IOMode.SYNC,
		partitionRequestInitialBackoff,
		partitionRequestMaxBackoff,
		networkBuffersPerChannel,
		extraNetworkBuffersPerGate,
		enableCreditBased);
}
 

private void testReadEmptyStateOrThrowException(boolean isRemote, ChannelStateReader reader) throws Exception {
	// setup
	final int totalBuffers = 10;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 32, 2);
	final SingleInputGate inputGate = createInputGate(globalPool);
	final RecoveredInputChannel inputChannel = createRecoveredChannel(isRemote, inputGate);

	try {
		inputGate.setInputChannels(inputChannel);
		inputGate.setup();

		// it would throw expected exception for the case of testReadStateWithException.
		inputChannel.readRecoveredState(reader);

		// the channel only has one EndOfChannelStateEvent in the queue for the case of testReadEmptyState.
		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertFalse(inputChannel.getNextBuffer().isPresent());
		assertTrue(inputChannel.getStateConsumedFuture().isDone());
	} finally {
		// cleanup and verify no buffer leak
		inputGate.close();
		globalPool.destroyAllBufferPools();
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
		globalPool.destroy();
	}
}
 

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

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

/**
 * Tests that the buffer is recycled correctly if exception is thrown during
 * {@link ChannelStateReader#readOutputData(ResultSubpartitionInfo, BufferBuilder)}.
 */
@Test
public void testReadRecoveredStateWithException() throws Exception {
	final int totalBuffers = 2;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 1, 1);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final ChannelStateReader stateReader = new ChannelStateReaderWithException();

	try {
		partition.setup();
		partition.readRecoveredState(stateReader);
	} catch (IOException e) {
		assertThat("should throw custom exception message", e.getMessage().contains("test"));
	} finally {
		globalPool.destroyAllBufferPools();
		// verify whether there are any buffers leak
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
		globalPool.destroy();
	}
}
 

NettyShuffleEnvironment(
		ResourceID taskExecutorResourceId,
		NettyShuffleEnvironmentConfiguration config,
		NetworkBufferPool networkBufferPool,
		ConnectionManager connectionManager,
		ResultPartitionManager resultPartitionManager,
		FileChannelManager fileChannelManager,
		ResultPartitionFactory resultPartitionFactory,
		SingleInputGateFactory singleInputGateFactory) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.config = config;
	this.networkBufferPool = networkBufferPool;
	this.connectionManager = connectionManager;
	this.resultPartitionManager = resultPartitionManager;
	this.inputGatesById = new ConcurrentHashMap<>(10);
	this.fileChannelManager = fileChannelManager;
	this.resultPartitionFactory = resultPartitionFactory;
	this.singleInputGateFactory = singleInputGateFactory;
	this.isClosed = false;
}
 

public SingleInputGateFactory(
		@Nonnull ResourceID taskExecutorResourceId,
		@Nonnull NettyShuffleEnvironmentConfiguration networkConfig,
		@Nonnull ConnectionManager connectionManager,
		@Nonnull ResultPartitionManager partitionManager,
		@Nonnull TaskEventPublisher taskEventPublisher,
		@Nonnull NetworkBufferPool networkBufferPool) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.partitionRequestInitialBackoff = networkConfig.partitionRequestInitialBackoff();
	this.partitionRequestMaxBackoff = networkConfig.partitionRequestMaxBackoff();
	this.networkBuffersPerChannel = networkConfig.networkBuffersPerChannel();
	this.floatingNetworkBuffersPerGate = networkConfig.floatingNetworkBuffersPerGate();
	this.blockingShuffleCompressionEnabled = networkConfig.isBlockingShuffleCompressionEnabled();
	this.compressionCodec = networkConfig.getCompressionCodec();
	this.networkBufferSize = networkConfig.networkBufferSize();
	this.connectionManager = connectionManager;
	this.partitionManager = partitionManager;
	this.taskEventPublisher = taskEventPublisher;
	this.networkBufferPool = networkBufferPool;
}
 

public SingleInputGateFactory(
		@Nonnull ResourceID taskExecutorResourceId,
		@Nonnull NettyShuffleEnvironmentConfiguration networkConfig,
		@Nonnull ConnectionManager connectionManager,
		@Nonnull ResultPartitionManager partitionManager,
		@Nonnull TaskEventPublisher taskEventPublisher,
		@Nonnull NetworkBufferPool networkBufferPool) {
	this.taskExecutorResourceId = taskExecutorResourceId;
	this.isCreditBased = networkConfig.isCreditBased();
	this.partitionRequestInitialBackoff = networkConfig.partitionRequestInitialBackoff();
	this.partitionRequestMaxBackoff = networkConfig.partitionRequestMaxBackoff();
	this.networkBuffersPerChannel = networkConfig.networkBuffersPerChannel();
	this.floatingNetworkBuffersPerGate = networkConfig.floatingNetworkBuffersPerGate();
	this.connectionManager = connectionManager;
	this.partitionManager = partitionManager;
	this.taskEventPublisher = taskEventPublisher;
	this.networkBufferPool = networkBufferPool;
}
 

@Test
public void testResumeConsumption() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = createPartitionRequestClient(channel, handler);

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

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		inputChannel.resumeConsumption();
		channel.runPendingTasks();
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(ResumeConsumption.class));
		assertEquals(inputChannel.getInputChannelId(), ((ResumeConsumption) readFromOutbound).receiverId);

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.close();

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

private void testConcurrentReadStateAndProcessAndRelease(boolean isRemote) throws Exception {
	// setup
	final int totalBuffers = 10;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 32, 2);
	final SingleInputGate inputGate = createInputGate(globalPool);
	final RecoveredInputChannel inputChannel = createRecoveredChannel(isRemote, inputGate);

	// the number of states is more that the total buffer amount
	final int totalStates = 15;
	final int[] states = {1, 2, 3, 4};
	final ChannelStateReader reader = new ResultPartitionTest.FiniteChannelStateReader(totalStates, states);
	final ExecutorService executor = Executors.newFixedThreadPool(2);
	Throwable thrown = null;
	try {
		inputGate.setInputChannels(inputChannel);
		inputGate.setup();

		final Callable<Void> processTask = processRecoveredBufferTask(inputChannel, totalStates, states, true);
		final Callable<Void> readStateTask = readRecoveredStateTask(inputChannel, reader, true);
		final Callable<Void> releaseTask = releaseChannelTask(inputChannel);

		submitTasksAndWaitForResults(executor, new Callable[] {readStateTask, processTask, releaseTask});
	} catch (Throwable t) {
		thrown = t;
	} finally {
		// cleanup
		cleanup(globalPool, executor, null, thrown, inputChannel);
	}
}
 

private static void releaseResource(SingleInputGate inputGate, NetworkBufferPool networkBufferPool) throws IOException {
	// Release all the buffer resources
	inputGate.close();

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

/**
 * Tests that input gate requests and assigns network buffers for remote input channel.
 */
@Test
public void testRequestBuffersWithRemoteInputChannel() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();
	final SingleInputGate inputGate = createInputGate(network, 1, ResultPartitionType.PIPELINED_BOUNDED);
	int buffersPerChannel = 2;
	int extraNetworkBuffersPerGate = 8;

	try {
		RemoteInputChannel remote =
			InputChannelBuilder.newBuilder()
				.setupFromNettyShuffleEnvironment(network)
				.setConnectionManager(new TestingConnectionManager())
				.buildRemoteChannel(inputGate);
		inputGate.setInputChannels(remote);
		inputGate.setup();

		NetworkBufferPool bufferPool = network.getNetworkBufferPool();
		// only the exclusive buffers should be assigned/available now
		assertEquals(buffersPerChannel, remote.getNumberOfAvailableBuffers());

		assertEquals(bufferPool.getTotalNumberOfMemorySegments() - buffersPerChannel,
			bufferPool.getNumberOfAvailableMemorySegments());
		// note: exclusive buffers are not handed out into LocalBufferPool and are thus not counted
		assertEquals(extraNetworkBuffersPerGate, bufferPool.countBuffers());
	} finally {
		inputGate.close();
		network.close();
	}
}
 

/**
 * Tests that input gate requests and assigns network buffers when unknown input channel
 * updates to remote input channel.
 */
@Test
public void testRequestBuffersWithUnknownInputChannel() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();
	final SingleInputGate inputGate = createInputGate(network, 1, ResultPartitionType.PIPELINED_BOUNDED);
	int buffersPerChannel = 2;
	int extraNetworkBuffersPerGate = 8;

	try {
		final ResultPartitionID resultPartitionId = new ResultPartitionID();
		InputChannel inputChannel = buildUnknownInputChannel(network, inputGate, resultPartitionId, 0);

		inputGate.setInputChannels(inputChannel);
		inputGate.setup();
		NetworkBufferPool bufferPool = network.getNetworkBufferPool();

		assertEquals(bufferPool.getTotalNumberOfMemorySegments(),
			bufferPool.getNumberOfAvailableMemorySegments());
		// note: exclusive buffers are not handed out into LocalBufferPool and are thus not counted
		assertEquals(extraNetworkBuffersPerGate, bufferPool.countBuffers());

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

		RemoteInputChannel remote = (RemoteInputChannel) inputGate.getInputChannels()
			.get(resultPartitionId.getPartitionId());
		// only the exclusive buffers should be assigned/available now
		assertEquals(buffersPerChannel, remote.getNumberOfAvailableBuffers());

		assertEquals(bufferPool.getTotalNumberOfMemorySegments() - buffersPerChannel,
			bufferPool.getNumberOfAvailableMemorySegments());
		// note: exclusive buffers are not handed out into LocalBufferPool and are thus not counted
		assertEquals(extraNetworkBuffersPerGate, bufferPool.countBuffers());
	} finally {
		inputGate.close();
		network.close();
	}
}
 

private void testConcurrentReadStateAndRelease(boolean isRemote) throws Exception {
	// setup
	final int totalBuffers = 10;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 32, 2);
	final SingleInputGate inputGate = createInputGate(globalPool);
	final RecoveredInputChannel inputChannel = createRecoveredChannel(isRemote, inputGate);

	// the number of states is more that the total buffer amount
	final int totalStates = 15;
	final int[] states = {1, 2, 3, 4};
	final ChannelStateReader reader = new ResultPartitionTest.FiniteChannelStateReader(totalStates, states);
	final ExecutorService executor = Executors.newFixedThreadPool(2);

	Throwable thrown = null;
	try {
		inputGate.setInputChannels(inputChannel);
		inputGate.setup();

		submitTasksAndWaitForResults(
			executor,
			new Callable[] {readRecoveredStateTask(inputChannel, reader, true), releaseChannelTask(inputChannel)});
	} catch (Throwable t) {
		thrown = t;
	} finally {
		// cleanup
		cleanup(globalPool, executor, null, thrown, inputChannel);
	}
}
 

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

private void testPartitionBufferPool(ResultPartitionType type) throws Exception {
	//setup
	final int networkBuffersPerChannel = 2;
	final int floatingNetworkBuffersPerGate = 8;
	final NetworkBufferPool globalPool = new NetworkBufferPool(20, 1, 1);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setResultPartitionType(type)
		.setFileChannelManager(fileChannelManager)
		.setNetworkBuffersPerChannel(networkBuffersPerChannel)
		.setFloatingNetworkBuffersPerGate(floatingNetworkBuffersPerGate)
		.setNetworkBufferPool(globalPool)
		.build();

	try {
		partition.setup();
		BufferPool bufferPool = partition.getBufferPool();
		// verify the amount of buffers in created local pool
		assertEquals(partition.getNumberOfSubpartitions() + 1, bufferPool.getNumberOfRequiredMemorySegments());
		if (type.isBounded()) {
			final int maxNumBuffers = networkBuffersPerChannel * partition.getNumberOfSubpartitions() + floatingNetworkBuffersPerGate;
			assertEquals(maxNumBuffers, bufferPool.getMaxNumberOfMemorySegments());
		} else {
			assertEquals(Integer.MAX_VALUE, bufferPool.getMaxNumberOfMemorySegments());
		}

	} finally {
		// cleanup
		globalPool.destroyAllBufferPools();
		globalPool.destroy();
	}
}
 

/**
 * Tests that the RecordWriter is available iif the respective LocalBufferPool has at-least one available buffer.
 */
@Test
public void testIsAvailableOrNot() throws Exception {
	// setup
	final NetworkBufferPool globalPool = new NetworkBufferPool(10, 128, 2);
	final BufferPool localPool = globalPool.createBufferPool(1, 1, null, 1, Integer.MAX_VALUE);
	final ResultPartitionWriter resultPartition = new ResultPartitionBuilder()
		.setBufferPoolFactory(p -> localPool)
		.build();
	resultPartition.setup();
	final ResultPartitionWriter partitionWrapper = new ConsumableNotifyingResultPartitionWriterDecorator(
		new NoOpTaskActions(),
		new JobID(),
		resultPartition,
		new NoOpResultPartitionConsumableNotifier());
	final RecordWriter recordWriter = createRecordWriter(partitionWrapper);

	try {
		// record writer is available because of initial available global pool
		assertTrue(recordWriter.getAvailableFuture().isDone());

		// request one buffer from the local pool to make it unavailable afterwards
		final BufferBuilder bufferBuilder = resultPartition.getBufferBuilder(0);
		assertNotNull(bufferBuilder);
		assertFalse(recordWriter.getAvailableFuture().isDone());

		// recycle the buffer to make the local pool available again
		final Buffer buffer = BufferBuilderTestUtils.buildSingleBuffer(bufferBuilder);
		buffer.recycleBuffer();
		assertTrue(recordWriter.getAvailableFuture().isDone());
		assertEquals(recordWriter.AVAILABLE, recordWriter.getAvailableFuture());

	} finally {
		localPool.lazyDestroy();
		globalPool.destroy();
	}
}
 

private RecordWriter createRecordWriter(NetworkBufferPool globalPool) throws Exception {
	final BufferPool localPool = globalPool.createBufferPool(1, 1, null, 1, Integer.MAX_VALUE);
	final ResultPartitionWriter partition = new ResultPartitionBuilder()
		.setBufferPoolFactory(p -> localPool)
		.build();
	partition.setup();

	return new RecordWriterBuilder().build(partition);
}
 

private static void instantiateNetworkMetrics(
	MetricGroup metrics,
	final NetworkEnvironment network) {

	final NetworkBufferPool networkBufferPool = network.getNetworkBufferPool();
	metrics.<Integer, Gauge<Integer>>gauge("TotalMemorySegments", networkBufferPool::getTotalNumberOfMemorySegments);
	metrics.<Integer, Gauge<Integer>>gauge("AvailableMemorySegments", networkBufferPool::getNumberOfAvailableMemorySegments);
}
 

private static void registerShuffleMetrics(
		String groupName,
		MetricGroup metricGroup,
		NetworkBufferPool networkBufferPool) {
	MetricGroup networkGroup = metricGroup.addGroup(groupName);
	networkGroup.<Integer, Gauge<Integer>>gauge(METRIC_TOTAL_MEMORY_SEGMENT,
		networkBufferPool::getTotalNumberOfMemorySegments);
	networkGroup.<Integer, Gauge<Integer>>gauge(METRIC_AVAILABLE_MEMORY_SEGMENT,
		networkBufferPool::getNumberOfAvailableMemorySegments);
}
 

public static void registerShuffleMetrics(
		MetricGroup metricGroup,
		NetworkBufferPool networkBufferPool) {
	checkNotNull(metricGroup);
	checkNotNull(networkBufferPool);

	//noinspection deprecation
	registerShuffleMetrics(METRIC_GROUP_NETWORK_DEPRECATED, metricGroup, networkBufferPool);
	registerShuffleMetrics(METRIC_GROUP_NETTY, metricGroup.addGroup(METRIC_GROUP_SHUFFLE), networkBufferPool);
}
 

public SingleInputGateBenchmarkFactory(
		ResourceID taskExecutorResourceId,
		NettyShuffleEnvironmentConfiguration networkConfig,
		ConnectionManager connectionManager,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		NetworkBufferPool networkBufferPool) {
	super(
		taskExecutorResourceId,
		networkConfig,
		connectionManager,
		partitionManager,
		taskEventPublisher,
		networkBufferPool);
}
 

private static ResultPartition createResultPartition(
		boolean releasePartitionOnConsumption,
		ResultPartitionType partitionType) {
	ResultPartitionFactory factory = new ResultPartitionFactory(
		new ResultPartitionManager(),
		fileChannelManager,
		new NetworkBufferPool(1, SEGMENT_SIZE, 1),
		BoundedBlockingSubpartitionType.AUTO,
		1,
		1,
		SEGMENT_SIZE,
		releasePartitionOnConsumption);

	final ResultPartitionDeploymentDescriptor descriptor = new ResultPartitionDeploymentDescriptor(
		new PartitionDescriptor(
			new IntermediateDataSetID(),
			new IntermediateResultPartitionID(),
			partitionType,
			1,
			0),
		NettyShuffleDescriptorBuilder.newBuilder().buildLocal(),
		1,
		true
	);

	return factory.create("test", descriptor);
}
 

/**
 * Tests that input gate requests and assigns network buffers for remote input channel.
 */
@Test
public void testRequestBuffersWithRemoteInputChannel() throws Exception {
	final NettyShuffleEnvironment network = createNettyShuffleEnvironment();
	final SingleInputGate inputGate = createInputGate(network, 1, ResultPartitionType.PIPELINED_BOUNDED);
	int buffersPerChannel = 2;
	int extraNetworkBuffersPerGate = 8;

	try {
		RemoteInputChannel remote =
			InputChannelBuilder.newBuilder()
				.setupFromNettyShuffleEnvironment(network)
				.setConnectionManager(new TestingConnectionManager())
				.buildRemoteAndSetToGate(inputGate);
		inputGate.setup();

		NetworkBufferPool bufferPool = network.getNetworkBufferPool();
		if (enableCreditBasedFlowControl) {
			// only the exclusive buffers should be assigned/available now
			assertEquals(buffersPerChannel, remote.getNumberOfAvailableBuffers());

			assertEquals(bufferPool.getTotalNumberOfMemorySegments() - buffersPerChannel,
				bufferPool.getNumberOfAvailableMemorySegments());
			// note: exclusive buffers are not handed out into LocalBufferPool and are thus not counted
			assertEquals(extraNetworkBuffersPerGate, bufferPool.countBuffers());
		} else {
			assertEquals(buffersPerChannel + extraNetworkBuffersPerGate, bufferPool.countBuffers());
		}
	} finally {
		inputGate.close();
		network.close();
	}
}
 

/**
 * Tests that the process of reading recovered state executes concurrently with channel
 * buffer processing, based on the condition of the total number of states is more that
 * the total buffer amount, to confirm that the lifecycle(recycle) of exclusive/floating
 * buffers works well.
 */
private void testConcurrentReadStateAndProcess(boolean isRemote) throws Exception {
	// setup
	final int totalBuffers = 10;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 32, 2);
	final SingleInputGate inputGate = createInputGate(globalPool);
	final RecoveredInputChannel inputChannel = createRecoveredChannel(isRemote, inputGate);

	// the number of states is more that the total buffer amount
	final int totalStates = 15;
	final int[] states = {1, 2, 3, 4};
	final ChannelStateReader reader = new ResultPartitionTest.FiniteChannelStateReader(totalStates, states);
	final ExecutorService executor = Executors.newFixedThreadPool(2);

	Throwable thrown = null;
	try {
		inputGate.setInputChannels(inputChannel);
		inputGate.setup();

		final Callable<Void> processTask = processRecoveredBufferTask(inputChannel, totalStates, states, false);
		final Callable<Void> readStateTask = readRecoveredStateTask(inputChannel, reader, false);

		submitTasksAndWaitForResults(executor, new Callable[] {readStateTask, processTask});
	} catch (Throwable t) {
		thrown = t;
	} finally {
		// cleanup
		cleanup(globalPool, executor, null, thrown, inputChannel);
	}
}
 

private static ResultPartition createResultPartition(
		boolean releasePartitionOnConsumption,
		ResultPartitionType partitionType) {
	ResultPartitionFactory factory = new ResultPartitionFactory(
		new ResultPartitionManager(),
		fileChannelManager,
		new NetworkBufferPool(1, SEGMENT_SIZE, 1),
		BoundedBlockingSubpartitionType.AUTO,
		1,
		1,
		SEGMENT_SIZE,
		releasePartitionOnConsumption,
		false,
		"LZ4",
		Integer.MAX_VALUE);

	final ResultPartitionDeploymentDescriptor descriptor = new ResultPartitionDeploymentDescriptor(
		PartitionDescriptorBuilder
			.newBuilder()
			.setPartitionType(partitionType)
			.build(),
		NettyShuffleDescriptorBuilder.newBuilder().buildLocal(),
		1,
		true
	);

	return factory.create("test", 0, descriptor);
}