Java Code Examples for org.apache.flink.runtime.io.network.buffer.BufferPool#requestBuffer()

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

/**
 * 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;
}
 

/**
 * Tests to verify the behaviours of recycling floating and exclusive buffers if the number of available
 * buffers equals to required buffers.
 */
@Test
public void testAvailableBuffersEqualToRequiredBuffers() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(16, 32);
	final int numExclusiveBuffers = 2;
	final int numFloatingBuffers = 14;

	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	inputGate.setInputChannel(inputChannel.partitionId.getPartitionId(), inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = spy(networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers));
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);
		inputChannel.requestSubpartition(0);

		// Prepare the exclusive and floating buffers to verify recycle logic later
		final Buffer exclusiveBuffer = inputChannel.requestBuffer();
		assertNotNull(exclusiveBuffer);
		final Buffer floatingBuffer = bufferPool.requestBuffer();
		assertNotNull(floatingBuffer);
		verify(bufferPool, times(1)).requestBuffer();

		// Receive the producer's backlog
		inputChannel.onSenderBacklog(12);

		// The channel requests (backlog + numExclusiveBuffers) floating buffers from local pool
		// and gets enough floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one floating buffer
		floatingBuffer.recycleBuffer();

		// The floating buffer is returned to local buffer directly because the channel is not waiting
		// for floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 1 buffer available in local pool",
			1, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one exclusive buffer
		exclusiveBuffer.recycleBuffer();

		// Return one extra floating buffer to the local pool because the number of available buffers
		// already equals to required buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 2 buffers available in local pool",
			2, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, null, thrown, inputChannel);
	}
}
 

/**
 * Tests to verify the behaviours of recycling floating and exclusive buffers if the number of available
 * buffers is more than required buffers by decreasing the sender's backlog.
 */
@Test
public void testAvailableBuffersMoreThanRequiredBuffers() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(16, 32);
	final int numExclusiveBuffers = 2;
	final int numFloatingBuffers = 14;

	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	inputGate.setInputChannel(inputChannel.partitionId.getPartitionId(), inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = spy(networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers));
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);
		inputChannel.requestSubpartition(0);

		// Prepare the exclusive and floating buffers to verify recycle logic later
		final Buffer exclusiveBuffer = inputChannel.requestBuffer();
		assertNotNull(exclusiveBuffer);

		final Buffer floatingBuffer = bufferPool.requestBuffer();
		assertNotNull(floatingBuffer);

		verify(bufferPool, times(1)).requestBuffer();

		// Receive the producer's backlog
		inputChannel.onSenderBacklog(12);

		// The channel gets enough floating buffers from local pool
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Decrease the backlog to make the number of available buffers more than required buffers
		inputChannel.onSenderBacklog(10);

		// Only the number of required buffers is changed by (backlog + numExclusiveBuffers)
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one exclusive buffer
		exclusiveBuffer.recycleBuffer();

		// Return one extra floating buffer to the local pool because the number of available buffers
		// is more than required buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 1 buffer available in local pool",
			1, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one floating buffer
		floatingBuffer.recycleBuffer();

		// The floating buffer is returned to local pool directly because the channel is not waiting for
		// floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 2 buffers available in local pool",
			2, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, null, thrown, inputChannel);
	}
}
 

/**
 * Tests to verify the behaviours of recycling floating and exclusive buffers if the number of available
 * buffers equals to required buffers.
 */
@Test
public void testAvailableBuffersEqualToRequiredBuffers() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(16, 32, 2);
	final int numFloatingBuffers = 14;

	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, networkBufferPool);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = spy(networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers));
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		// Prepare the exclusive and floating buffers to verify recycle logic later
		final Buffer exclusiveBuffer = inputChannel.requestBuffer();
		assertNotNull(exclusiveBuffer);
		final Buffer floatingBuffer = bufferPool.requestBuffer();
		assertNotNull(floatingBuffer);
		verify(bufferPool, times(1)).requestBuffer();

		// Receive the producer's backlog
		inputChannel.onSenderBacklog(12);

		// The channel requests (backlog + numExclusiveBuffers) floating buffers from local pool
		// and gets enough floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one floating buffer
		floatingBuffer.recycleBuffer();

		// The floating buffer is returned to local buffer directly because the channel is not waiting
		// for floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 1 buffer available in local pool",
			1, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one exclusive buffer
		exclusiveBuffer.recycleBuffer();

		// Return one extra floating buffer to the local pool because the number of available buffers
		// already equals to required buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 2 buffers available in local pool",
			2, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, null, thrown, inputChannel);
	}
}
 

/**
 * Tests to verify the behaviours of recycling floating and exclusive buffers if the number of available
 * buffers is more than required buffers by decreasing the sender's backlog.
 */
@Test
public void testAvailableBuffersMoreThanRequiredBuffers() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(16, 32, 2);
	final int numFloatingBuffers = 14;

	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, networkBufferPool);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = spy(networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers));
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		// Prepare the exclusive and floating buffers to verify recycle logic later
		final Buffer exclusiveBuffer = inputChannel.requestBuffer();
		assertNotNull(exclusiveBuffer);

		final Buffer floatingBuffer = bufferPool.requestBuffer();
		assertNotNull(floatingBuffer);

		verify(bufferPool, times(1)).requestBuffer();

		// Receive the producer's backlog
		inputChannel.onSenderBacklog(12);

		// The channel gets enough floating buffers from local pool
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Decrease the backlog to make the number of available buffers more than required buffers
		inputChannel.onSenderBacklog(10);

		// Only the number of required buffers is changed by (backlog + numExclusiveBuffers)
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one exclusive buffer
		exclusiveBuffer.recycleBuffer();

		// Return one extra floating buffer to the local pool because the number of available buffers
		// is more than required buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 1 buffer available in local pool",
			1, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one floating buffer
		floatingBuffer.recycleBuffer();

		// The floating buffer is returned to local pool directly because the channel is not waiting for
		// floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel);
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 2 buffers available in local pool",
			2, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, null, thrown, inputChannel);
	}
}
 

/**
 * Tests to verify the behaviours of recycling floating and exclusive buffers if the number of available
 * buffers equals to required buffers.
 */
@Test
public void testAvailableBuffersEqualToRequiredBuffers() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(16, 32, 2);
	final int numFloatingBuffers = 14;

	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	inputGate.setInputChannels(inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = spy(networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers));
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		// Prepare the exclusive and floating buffers to verify recycle logic later
		final Buffer exclusiveBuffer = inputChannel.requestBuffer();
		assertNotNull(exclusiveBuffer);
		final Buffer floatingBuffer = bufferPool.requestBuffer();
		assertNotNull(floatingBuffer);
		verify(bufferPool, times(1)).requestBuffer();

		// Receive the producer's backlog
		inputChannel.onSenderBacklog(12);

		// The channel requests (backlog + numExclusiveBuffers) floating buffers from local pool
		// and gets enough floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one floating buffer
		floatingBuffer.recycleBuffer();

		// The floating buffer is returned to local buffer directly because the channel is not waiting
		// for floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 1 buffer available in local pool",
			1, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one exclusive buffer
		exclusiveBuffer.recycleBuffer();

		// Return one extra floating buffer to the local pool because the number of available buffers
		// already equals to required buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 2 buffers available in local pool",
			2, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, null, thrown, inputChannel);
	}
}
 

/**
 * Tests to verify the behaviours of recycling floating and exclusive buffers if the number of available
 * buffers is more than required buffers by decreasing the sender's backlog.
 */
@Test
public void testAvailableBuffersMoreThanRequiredBuffers() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(16, 32, 2);
	final int numFloatingBuffers = 14;

	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	inputGate.setInputChannels(inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = spy(networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers));
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		// Prepare the exclusive and floating buffers to verify recycle logic later
		final Buffer exclusiveBuffer = inputChannel.requestBuffer();
		assertNotNull(exclusiveBuffer);

		final Buffer floatingBuffer = bufferPool.requestBuffer();
		assertNotNull(floatingBuffer);

		verify(bufferPool, times(1)).requestBuffer();

		// Receive the producer's backlog
		inputChannel.onSenderBacklog(12);

		// The channel gets enough floating buffers from local pool
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 14 buffers required in the channel",
			14, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Decrease the backlog to make the number of available buffers more than required buffers
		inputChannel.onSenderBacklog(10);

		// Only the number of required buffers is changed by (backlog + numExclusiveBuffers)
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 0 buffers available in local pool",
			0, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one exclusive buffer
		exclusiveBuffer.recycleBuffer();

		// Return one extra floating buffer to the local pool because the number of available buffers
		// is more than required buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 1 buffer available in local pool",
			1, bufferPool.getNumberOfAvailableMemorySegments());

		// Recycle one floating buffer
		floatingBuffer.recycleBuffer();

		// The floating buffer is returned to local pool directly because the channel is not waiting for
		// floating buffers
		verify(bufferPool, times(14)).requestBuffer();
		verify(bufferPool, times(0)).addBufferListener(inputChannel.getBufferManager());
		assertEquals("There should be 14 buffers available in the channel",
			14, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 12 buffers required in the channel",
			12, inputChannel.getNumberOfRequiredBuffers());
		assertEquals("There should be 2 buffers available in local pool",
			2, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, null, thrown, inputChannel);
	}
}