Java Code Examples for org.apache.flink.runtime.io.network.buffer.BufferBuilder#createBufferConsumer()

/**
 * The request could be from broadcast or non-broadcast modes like {@link #randomEmit(IOReadableWritable)}.
 *
 * <p>For non-broadcast, the created {@link BufferConsumer} is only for the target channel.
 *
 * <p>For broadcast, all the channels share the same requested {@link BufferBuilder} and the created
 * {@link BufferConsumer} is copied for every channel.
 */
@Override
public BufferBuilder requestNewBufferBuilder(int targetChannel) throws IOException, InterruptedException {
	checkState(bufferBuilder == null || bufferBuilder.isFinished());

	BufferBuilder builder = super.requestNewBufferBuilder(targetChannel);
	if (randomTriggered) {
		addBufferConsumer(randomTriggeredConsumer = builder.createBufferConsumer(), targetChannel);
	} else {
		try (BufferConsumer bufferConsumer = builder.createBufferConsumer()) {
			for (int channel = 0; channel < numberOfChannels; channel++) {
				addBufferConsumer(bufferConsumer.copy(), channel);
			}
		}
	}

	bufferBuilder = builder;
	return builder;
}
 

@Override
public void readRecoveredState(ChannelStateReader stateReader) throws IOException, InterruptedException {
	boolean recycleBuffer = true;
	for (ReadResult readResult = ReadResult.HAS_MORE_DATA; readResult == ReadResult.HAS_MORE_DATA;) {
		BufferBuilder bufferBuilder = parent.getBufferPool().requestBufferBuilderBlocking(subpartitionInfo.getSubPartitionIdx());
		BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
		try {
			readResult = stateReader.readOutputData(subpartitionInfo, bufferBuilder);

			// check whether there are some states data filled in this time
			if (bufferConsumer.isDataAvailable()) {
				add(bufferConsumer, false, false);
				recycleBuffer = false;
				bufferBuilder.finish();
			}
		} finally {
			if (recycleBuffer) {
				bufferConsumer.close();
			}
		}
	}
}
 

@Test
public void testReadToBufferBuilder() throws IOException {
	byte[] data = generateData(100);
	BufferBuilder bufferBuilder = new BufferBuilder(HeapMemorySegment.FACTORY.allocateUnpooledSegment(data.length, null), FreeingBufferRecycler.INSTANCE);
	BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();

	new ChannelStateSerializerImpl().readData(new ByteArrayInputStream(data), wrap(bufferBuilder), Integer.MAX_VALUE);

	assertFalse(bufferBuilder.isFinished());

	bufferBuilder.finish();
	Buffer buffer = bufferConsumer.build();

	assertEquals(data.length, buffer.readableBytes());
	byte[] actual = new byte[buffer.readableBytes()];
	buffer.asByteBuf().readBytes(actual);
	assertArrayEquals(data, actual);
}
 

private static BufferAndSerializerResult setNextBufferForSerializer(
		RecordSerializer<SerializationTestType> serializer,
		int segmentSize) throws IOException {
	// create a bufferBuilder with some random starting offset to properly test handling buffer slices in the
	// deserialization code.
	int startingOffset = segmentSize > 2 ? RANDOM.nextInt(segmentSize / 2) : 0;
	BufferBuilder bufferBuilder = createFilledBufferBuilder(segmentSize + startingOffset, startingOffset);
	BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
	bufferConsumer.build().recycleBuffer();

	return new BufferAndSerializerResult(
		bufferBuilder,
		bufferConsumer,
		serializer.copyToBufferBuilder(bufferBuilder));
}
 

@Override
public BufferConsumerAndChannel getNextBufferConsumer() throws Exception {
	if (channelIndexes.size() > 0) {
		final int channelIndex = channelIndexes.remove(0);
		BufferBuilder bufferBuilder = bufferProvider.requestBufferBuilderBlocking();
		bufferBuilder.appendAndCommit(ByteBuffer.wrap(new byte[4]));
		bufferBuilder.finish();
		return new BufferConsumerAndChannel(bufferBuilder.createBufferConsumer(), channelIndex);
	}

	return null;
}
 

private static BufferAndSerializerResult setNextBufferForSerializer(
		RecordSerializer<SerializationTestType> serializer,
		int segmentSize) throws IOException {
	// create a bufferBuilder with some random starting offset to properly test handling buffer slices in the
	// deserialization code.
	int startingOffset = segmentSize > 2 ? RANDOM.nextInt(segmentSize / 2) : 0;
	BufferBuilder bufferBuilder = createFilledBufferBuilder(segmentSize + startingOffset, startingOffset);
	BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
	bufferConsumer.build().recycleBuffer();

	return new BufferAndSerializerResult(
		bufferBuilder,
		bufferConsumer,
		serializer.copyToBufferBuilder(bufferBuilder));
}
 

@Override
public BufferConsumerAndChannel getNextBufferConsumer() throws Exception {
	if (channelIndexes.size() > 0) {
		final int channelIndex = channelIndexes.remove(0);
		BufferBuilder bufferBuilder = bufferProvider.requestBufferBuilderBlocking();
		bufferBuilder.appendAndCommit(ByteBuffer.wrap(new byte[4]));
		bufferBuilder.finish();
		return new BufferConsumerAndChannel(bufferBuilder.createBufferConsumer(), channelIndex);
	}

	return null;
}
 

private static Buffer appendLeftOverBytes(Buffer buffer, byte[] leftOverBytes) {
	BufferBuilder bufferBuilder = new BufferBuilder(
		MemorySegmentFactory.allocateUnpooledSegment(buffer.readableBytes() + leftOverBytes.length),
		FreeingBufferRecycler.INSTANCE);
	try (BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer()) {
		bufferBuilder.append(buffer.getNioBufferReadable());
		bufferBuilder.appendAndCommit(ByteBuffer.wrap(leftOverBytes));
		return bufferConsumer.build();
	}
}
 

private static BufferAndSerializerResult setNextBufferForSerializer(
		RecordSerializer<SerializationTestType> serializer,
		int segmentSize) throws IOException {
	// create a bufferBuilder with some random starting offset to properly test handling buffer slices in the
	// deserialization code.
	int startingOffset = segmentSize > 2 ? RANDOM.nextInt(segmentSize / 2) : 0;
	BufferBuilder bufferBuilder = createFilledBufferBuilder(segmentSize + startingOffset, startingOffset);
	BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
	bufferConsumer.build().recycleBuffer();

	return new BufferAndSerializerResult(
		bufferBuilder,
		bufferConsumer,
		serializer.copyToBufferBuilder(bufferBuilder));
}
 

@Override
public BufferConsumerAndChannel getNextBufferConsumer() throws Exception {
	if (channelIndexes.size() > 0) {
		final int channelIndex = channelIndexes.remove(0);
		BufferBuilder bufferBuilder = bufferProvider.requestBufferBuilderBlocking();
		BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
		bufferBuilder.appendAndCommit(ByteBuffer.wrap(new byte[4]));
		bufferBuilder.finish();
		return new BufferConsumerAndChannel(bufferConsumer, channelIndex);
	}

	return null;
}
 

private IteratorWrappingTestSingleInputGate<T> wrapIterator(MutableObjectIterator<T> iterator) throws IOException, InterruptedException {
	inputIterator = iterator;
	serializer = new SpanningRecordSerializer<T>();

	// The input iterator can produce an infinite stream. That's why we have to serialize each
	// record on demand and cannot do it upfront.
	final BufferAndAvailabilityProvider answer = new BufferAndAvailabilityProvider() {

		private boolean hasData = inputIterator.next(reuse) != null;

		@Override
		public Optional<BufferAndAvailability> getBufferAvailability() throws IOException {
			if (hasData) {
				serializer.serializeRecord(reuse);
				BufferBuilder bufferBuilder = createBufferBuilder(bufferSize);
				BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
				serializer.copyToBufferBuilder(bufferBuilder);

				hasData = inputIterator.next(reuse) != null;

				// Call getCurrentBuffer to ensure size is set
				return Optional.of(new BufferAndAvailability(bufferConsumer.build(), true, 0));
			} else {
				inputChannel.setReleased();

				return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
					false,
					0));
			}
		}
	};

	inputChannel.addBufferAndAvailability(answer);

	inputGate.setInputChannels(inputChannel);

	return this;
}
 

private BufferOrEvent createDataBuffer() throws IOException {
	BufferBuilder bufferBuilder = BufferBuilderTestUtils.createEmptyBufferBuilder(PAGE_SIZE);
	BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
	serializeRecord(42L, bufferBuilder);
	serializeRecord(44L, bufferBuilder);

	return new BufferOrEvent(bufferConsumer.build(), new InputChannelInfo(0, 0), false);
}
 

@SuppressWarnings("unchecked")
private void setupInputChannels() {

	for (int i = 0; i < numInputChannels; i++) {
		final int channelIndex = i;
		final RecordSerializer<SerializationDelegate<Object>> recordSerializer = new SpanningRecordSerializer<SerializationDelegate<Object>>();
		final SerializationDelegate<Object> delegate = (SerializationDelegate<Object>) (SerializationDelegate<?>)
			new SerializationDelegate<>(new StreamElementSerializer<T>(serializer));

		inputQueues[channelIndex] = new ConcurrentLinkedQueue<>();
		inputChannels[channelIndex] = new TestInputChannel(inputGate, i);

		final BufferAndAvailabilityProvider answer = () -> {
			ConcurrentLinkedQueue<InputValue<Object>> inputQueue = inputQueues[channelIndex];
			InputValue<Object> input;
			boolean moreAvailable;
			synchronized (inputQueue) {
				input = inputQueue.poll();
				moreAvailable = !inputQueue.isEmpty();
			}
			if (input != null && input.isStreamEnd()) {
				inputChannels[channelIndex].setReleased();
				return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE), moreAvailable, 0));
			} else if (input != null && input.isStreamRecord()) {
				Object inputElement = input.getStreamRecord();

				delegate.setInstance(inputElement);
				recordSerializer.serializeRecord(delegate);
				BufferBuilder bufferBuilder = createBufferBuilder(bufferSize);
				BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
				recordSerializer.copyToBufferBuilder(bufferBuilder);
				bufferBuilder.finish();

				// Call getCurrentBuffer to ensure size is set
				return Optional.of(new BufferAndAvailability(bufferConsumer.build(), moreAvailable, 0));
			} else if (input != null && input.isEvent()) {
				AbstractEvent event = input.getEvent();
				if (event instanceof EndOfPartitionEvent) {
					inputChannels[channelIndex].setReleased();
				}

				return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(event), moreAvailable, 0));
			} else {
				return Optional.empty();
			}
		};

		inputChannels[channelIndex].addBufferAndAvailability(answer);
	}
	inputGate.setInputChannels(inputChannels);
}