Java Code Examples for org.apache.flink.runtime.io.network.buffer.Buffer#recycleBuffer()

/**
 * Tests {@link EventSerializer#isEvent(Buffer, Class)}
 * whether it peaks into the buffer only, i.e. after the call, the buffer
 * is still de-serializable.
 */
@Test
public void testIsEventPeakOnly() throws Exception {
	final Buffer serializedEvent =
		EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	try {
		final ClassLoader cl = getClass().getClassLoader();
		assertTrue(
			EventSerializer.isEvent(serializedEvent, EndOfPartitionEvent.class));
		EndOfPartitionEvent event = (EndOfPartitionEvent) EventSerializer
			.fromBuffer(serializedEvent, cl);
		assertEquals(EndOfPartitionEvent.INSTANCE, event);
	} finally {
		serializedEvent.recycleBuffer();
	}
}
 

/**
 * Tests {@link EventSerializer#isEvent(Buffer, Class)}
 * whether it peaks into the buffer only, i.e. after the call, the buffer
 * is still de-serializable.
 */
@Test
public void testIsEventPeakOnly() throws Exception {
	final Buffer serializedEvent =
		EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	try {
		final ClassLoader cl = getClass().getClassLoader();
		assertTrue(
			EventSerializer.isEvent(serializedEvent, EndOfPartitionEvent.class));
		EndOfPartitionEvent event = (EndOfPartitionEvent) EventSerializer
			.fromBuffer(serializedEvent, cl);
		assertEquals(EndOfPartitionEvent.INSTANCE, event);
	} finally {
		serializedEvent.recycleBuffer();
	}
}
 

private void onRecoveredStateBuffer(Buffer buffer) {
	boolean recycleBuffer = true;
	try {
		final boolean wasEmpty;
		synchronized (receivedBuffers) {
			// Similar to notifyBufferAvailable(), make sure that we never add a buffer
			// after releaseAllResources() released all buffers from receivedBuffers.
			if (isReleased) {
				return;
			}

			wasEmpty = receivedBuffers.isEmpty();
			receivedBuffers.add(buffer);
			recycleBuffer = false;
		}

		if (wasEmpty) {
			notifyChannelNonEmpty();
		}
	} finally {
		if (recycleBuffer) {
			buffer.recycleBuffer();
		}
	}
}
 

public void clearBuffers() {
	for (RecordDeserializer<?> deserializer : recordDeserializers) {
		Buffer buffer = deserializer.getCurrentBuffer();
		if (buffer != null && !buffer.isRecycled()) {
			buffer.recycleBuffer();
		}
		deserializer.clear();
	}
}
 

/**
 * 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 static BufferOrEvent parseBuffer(BufferConsumer bufferConsumer, int targetChannel) throws IOException {
	Buffer buffer = buildSingleBuffer(bufferConsumer);
	if (buffer.isBuffer()) {
		return new BufferOrEvent(buffer, targetChannel);
	} else {
		// is event:
		AbstractEvent event = EventSerializer.fromBuffer(buffer, RecordWriterTest.class.getClassLoader());
		buffer.recycleBuffer(); // the buffer is not needed anymore
		return new BufferOrEvent(event, targetChannel);
	}
}
 

@Override
public void invoke() throws Exception {
	InputGate gate = getEnvironment().getInputGate(0);
	gate.requestPartitions();
	Buffer buffer = gate.getNext().orElseThrow(IllegalStateException::new).getBuffer();
	if (buffer != null) {
		buffer.recycleBuffer();
	}
}
 

@Override
public void close() throws IOException {
	// clear the buffers first. this part should not ever fail
	for (RecordDeserializer<?> deserializer : recordDeserializers) {
		Buffer buffer = deserializer.getCurrentBuffer();
		if (buffer != null && !buffer.isRecycled()) {
			buffer.recycleBuffer();
		}
		deserializer.clear();
	}

	// cleanup the barrier handler resources
	barrierHandler.cleanup();
}
 

/**
 * Writes the given block asynchronously.
 *
 * @param buffer
 * 		the buffer to be written (will be recycled when done)
 *
 * @throws IOException
 * 		thrown if adding the write operation fails
 */
@Override
public void writeBlock(Buffer buffer) throws IOException {
	try {
		// if successfully added, the buffer will be recycled after the write operation
		addRequest(new BufferWriteRequest(this, buffer));
	} catch (Throwable e) {
		// if not added, we need to recycle here
		buffer.recycleBuffer();
		ExceptionUtils.rethrowIOException(e);
	}

}
 

private void writeAndCloseBufferConsumer(BufferConsumer bufferConsumer) throws IOException {
	try {
		final Buffer buffer = bufferConsumer.build();
		try {
			if (canBeCompressed(buffer)) {
				final Buffer compressedBuffer = parent.bufferCompressor.compressToIntermediateBuffer(buffer);
				data.writeBuffer(compressedBuffer);
				if (compressedBuffer != buffer) {
					compressedBuffer.recycleBuffer();
				}
			} else {
				data.writeBuffer(buffer);
			}

			numBuffersAndEventsWritten++;
			if (buffer.isBuffer()) {
				numDataBuffersWritten++;
			}
		}
		finally {
			buffer.recycleBuffer();
		}
	}
	finally {
		bufferConsumer.close();
	}
}
 

@Override
public void readInto(Buffer buffer) throws IOException {
	if (fileChannel.size() - fileChannel.position() > 0) {
		hasReachedEndOfFile = reader.readBufferFromFileChannel(buffer);
	}
	else {
		buffer.recycleBuffer();
	}
}
 

@VisibleForTesting
long spillFinishedBufferConsumers(boolean forceFinishRemainingBuffers) throws IOException {
	long spilledBytes = 0;

	while (!buffers.isEmpty()) {
		BufferConsumer bufferConsumer = buffers.getFirst();
		Buffer buffer = bufferConsumer.build();
		updateStatistics(buffer);
		int bufferSize = buffer.getSize();
		spilledBytes += bufferSize;

		// NOTE we may be in the process of finishing the subpartition where any buffer should
		// be treated as if it was finished!
		if (bufferConsumer.isFinished() || forceFinishRemainingBuffers) {
			if (bufferSize > 0) {
				spillWriter.writeBlock(buffer);
			} else {
				// If we skip a buffer for the spill writer, we need to adapt the backlog accordingly
				decreaseBuffersInBacklog(buffer);
				buffer.recycleBuffer();
			}
			bufferConsumer.close();
			buffers.poll();
		} else {
			// If there is already data, we need to spill it anyway, since we do not get this
			// slice from the buffer consumer again during the next build.
			// BEWARE: by doing so, we increase the actual number of buffers in the spill writer!
			if (bufferSize > 0) {
				spillWriter.writeBlock(buffer);
				increaseBuffersInBacklog(bufferConsumer);
			} else {
				buffer.recycleBuffer();
			}

			return spilledBytes;
		}
	}
	return spilledBytes;
}
 

@Override
protected MemorySegment nextSegment(MemorySegment current) throws IOException {
	if (cause.get() != null) {
		throw cause.get();
	}

	// check for end-of-stream
	if (this.numBlocksRemaining <= 0) {
		this.reader.close();
		throw new EOFException();
	}

	try {
		Buffer buffer;
		while ((buffer = retBuffers.poll(1, TimeUnit.SECONDS)) == null) {
			if (cause.get() != null) {
				throw cause.get();
			}
		}
		this.currentSegmentLimit = decompressor.decompress(
				buffer.getMemorySegment().getArray(), 0, buffer.getSize(),
				uncompressedBuffer.getArray(), 0
		);

		buffer.recycleBuffer();
		this.numBlocksRemaining--;
		return uncompressedBuffer;
	}
	catch (InterruptedException e) {
		throw new IOException(e);
	}
}
 

private BufferOrEvent transformEvent(
		Buffer buffer,
		boolean moreAvailable,
		InputChannel currentChannel) throws IOException, InterruptedException {
	final AbstractEvent event;
	try {
		event = EventSerializer.fromBuffer(buffer, getClass().getClassLoader());
	} finally {
		buffer.recycleBuffer();
	}

	if (event.getClass() == EndOfPartitionEvent.class) {
		channelsWithEndOfPartitionEvents.set(currentChannel.getChannelIndex());

		if (channelsWithEndOfPartitionEvents.cardinality() == numberOfInputChannels) {
			// Because of race condition between:
			// 1. releasing inputChannelsWithData lock in this method and reaching this place
			// 2. empty data notification that re-enqueues a channel
			// we can end up with moreAvailable flag set to true, while we expect no more data.
			checkState(!moreAvailable || !pollNext().isPresent());
			moreAvailable = false;
			hasReceivedAllEndOfPartitionEvents = true;
			markAvailable();
		}

		currentChannel.releaseAllResources();
	}

	return new BufferOrEvent(event, currentChannel.getChannelInfo(), moreAvailable, buffer.getSize());
}
 

/**
 * Releases all exclusive and floating buffers, closes the partition request client.
 */
@Override
void releaseAllResources() throws IOException {
	if (isReleased.compareAndSet(false, true)) {

		// Gather all exclusive buffers and recycle them to global pool in batch, because
		// we do not want to trigger redistribution of buffers after each recycle.
		final List<MemorySegment> exclusiveRecyclingSegments = new ArrayList<>();

		synchronized (receivedBuffers) {
			Buffer buffer;
			while ((buffer = receivedBuffers.poll()) != null) {
				if (buffer.getRecycler() == this) {
					exclusiveRecyclingSegments.add(buffer.getMemorySegment());
				} else {
					buffer.recycleBuffer();
				}
			}
		}
		synchronized (bufferQueue) {
			bufferQueue.releaseAll(exclusiveRecyclingSegments);
		}

		if (exclusiveRecyclingSegments.size() > 0) {
			inputGate.returnExclusiveSegments(exclusiveRecyclingSegments);
		}

		// The released flag has to be set before closing the connection to ensure that
		// buffers received concurrently with closing are properly recycled.
		if (partitionRequestClient != null) {
			partitionRequestClient.close(this);
		} else {
			connectionManager.closeOpenChannelConnections(connectionId);
		}
	}
}
 

public void onBuffer(Buffer buffer, int sequenceNumber, int backlog) throws IOException {
	boolean recycleBuffer = true;

	try {
		if (expectedSequenceNumber != sequenceNumber) {
			onError(new BufferReorderingException(expectedSequenceNumber, sequenceNumber));
			return;
		}

		final boolean wasEmpty;
		final CheckpointBarrier notifyReceivedBarrier;
		final Buffer notifyReceivedBuffer;
		final BufferReceivedListener listener = inputGate.getBufferReceivedListener();
		synchronized (receivedBuffers) {
			// Similar to notifyBufferAvailable(), make sure that we never add a buffer
			// after releaseAllResources() released all buffers from receivedBuffers
			// (see above for details).
			if (isReleased.get()) {
				return;
			}

			wasEmpty = receivedBuffers.isEmpty();
			receivedBuffers.add(buffer);

			if (listener != null && buffer.isBuffer() && receivedCheckpointId < lastRequestedCheckpointId) {
				notifyReceivedBuffer = buffer.retainBuffer();
			} else {
				notifyReceivedBuffer = null;
			}
			notifyReceivedBarrier = listener != null ? parseCheckpointBarrierOrNull(buffer) : null;
		}
		recycleBuffer = false;

		++expectedSequenceNumber;

		if (wasEmpty) {
			notifyChannelNonEmpty();
		}

		if (backlog >= 0) {
			onSenderBacklog(backlog);
		}

		if (notifyReceivedBarrier != null) {
			receivedCheckpointId = notifyReceivedBarrier.getId();
			if (notifyReceivedBarrier.isCheckpoint()) {
				listener.notifyBarrierReceived(notifyReceivedBarrier, channelInfo);
			}
		} else if (notifyReceivedBuffer != null) {
			listener.notifyBufferReceived(notifyReceivedBuffer, channelInfo);
		}
	} finally {
		if (recycleBuffer) {
			buffer.recycleBuffer();
		}
	}
}
 

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

protected boolean getNextRecord(T target) throws IOException, InterruptedException {
	// The action of partition request was removed from InputGate#setup since FLINK-16536, and this is the only
	// unified way for launching partition request for batch jobs. In order to avoid potential performance concern,
	// we might consider migrating this action back to the setup based on some condition judgement future.
	if (!requestedPartitions) {
		inputGate.requestPartitions();
		requestedPartitions = true;
	}

	if (isFinished) {
		return false;
	}

	while (true) {
		if (currentRecordDeserializer != null) {
			DeserializationResult result = currentRecordDeserializer.getNextRecord(target);

			if (result.isBufferConsumed()) {
				final Buffer currentBuffer = currentRecordDeserializer.getCurrentBuffer();

				currentBuffer.recycleBuffer();
				currentRecordDeserializer = null;
			}

			if (result.isFullRecord()) {
				return true;
			}
		}

		final BufferOrEvent bufferOrEvent = inputGate.getNext().orElseThrow(IllegalStateException::new);

		if (bufferOrEvent.isBuffer()) {
			currentRecordDeserializer = recordDeserializers.get(bufferOrEvent.getChannelInfo());
			currentRecordDeserializer.setNextBuffer(bufferOrEvent.getBuffer());
		}
		else {
			// sanity check for leftover data in deserializers. events should only come between
			// records, not in the middle of a fragment
			if (recordDeserializers.get(bufferOrEvent.getChannelInfo()).hasUnfinishedData()) {
				throw new IOException(
						"Received an event in channel " + bufferOrEvent.getChannelInfo() + " while still having "
						+ "data from a record. This indicates broken serialization logic. "
						+ "If you are using custom serialization code (Writable or Value types), check their "
						+ "serialization routines. In the case of Kryo, check the respective Kryo serializer.");
			}

			if (handleEvent(bufferOrEvent.getEvent())) {
				if (inputGate.isFinished()) {
					isFinished = true;
					return false;
				}
				else if (hasReachedEndOfSuperstep()) {
					return false;
				}
				// else: More data is coming...
			}
		}
	}
}
 

/**
 * Tests that failures are propagated correctly if
 * {@link RemoteInputChannel#notifyBufferAvailable(int)} throws an exception. Also tests that
 * a second listener will be notified in this case.
 */
@Test
public void testFailureInNotifyBufferAvailable() throws Exception {
	// Setup
	final int numExclusiveBuffers = 1;
	final int numFloatingBuffers = 1;
	final int numTotalBuffers = numExclusiveBuffers + numFloatingBuffers;
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(
		numTotalBuffers, 32, numExclusiveBuffers);

	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel successfulRemoteIC = createRemoteInputChannel(inputGate);
	successfulRemoteIC.requestSubpartition(0);

	// late creation -> no exclusive buffers, also no requested subpartition in successfulRemoteIC
	// (to trigger a failure in RemoteInputChannel#notifyBufferAvailable())
	final RemoteInputChannel failingRemoteIC = createRemoteInputChannel(inputGate);

	Buffer buffer = null;
	Throwable thrown = null;
	try {
		final BufferPool bufferPool =
			networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);

		buffer = checkNotNull(bufferPool.requestBuffer());

		// trigger subscription to buffer pool
		failingRemoteIC.onSenderBacklog(1);
		successfulRemoteIC.onSenderBacklog(numExclusiveBuffers + 1);
		// recycling will call RemoteInputChannel#notifyBufferAvailable() which will fail and
		// this exception will be swallowed and set as an error in failingRemoteIC
		buffer.recycleBuffer();
		buffer = null;
		try {
			failingRemoteIC.checkError();
			fail("The input channel should have an error based on the failure in RemoteInputChannel#notifyBufferAvailable()");
		} catch (IOException e) {
			assertThat(e, hasProperty("cause", isA(IllegalStateException.class)));
		}
		// currently, the buffer is still enqueued in the bufferQueue of failingRemoteIC
		assertEquals(0, bufferPool.getNumberOfAvailableMemorySegments());
		buffer = successfulRemoteIC.requestBuffer();
		assertNull("buffer should still remain in failingRemoteIC", buffer);

		// releasing resources in failingRemoteIC should free the buffer again and immediately
		// recycle it into successfulRemoteIC
		failingRemoteIC.releaseAllResources();
		assertEquals(0, bufferPool.getNumberOfAvailableMemorySegments());
		buffer = successfulRemoteIC.requestBuffer();
		assertNotNull("no buffer given to successfulRemoteIC", buffer);
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, buffer, thrown, failingRemoteIC, successfulRemoteIC);
	}
}
 

@Override
public void onBuffer(Buffer buffer) {
	super.onBuffer(buffer);

	buffer.recycleBuffer();
}