Java Code Examples for org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler#INSTANCE

private static BufferOrEvent createBuffer(int channel) {
	final int size = sizeCounter++;
	byte[] bytes = new byte[size];
	RND.nextBytes(bytes);

	MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	memory.put(0, bytes);

	Buffer buf = new NetworkBuffer(memory, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);

	// retain an additional time so it does not get disposed after being read by the input gate
	buf.retainBuffer();

	return new BufferOrEvent(buf, new InputChannelInfo(0, channel));
}
 

@Test
@SuppressWarnings("ConstantConditions")
public void testSmallFilesNotWritten() throws Exception {
	int threshold = 100;
	File checkpointsDir = temporaryFolder.newFolder("checkpointsDir");
	File sharedStateDir = temporaryFolder.newFolder("sharedStateDir");
	FsCheckpointStreamFactory checkpointStreamFactory = new FsCheckpointStreamFactory(getSharedInstance(), fromLocalFile(checkpointsDir), fromLocalFile(sharedStateDir), threshold, threshold);
	ChannelStateWriteResult result = new ChannelStateWriteResult();
	ChannelStateCheckpointWriter writer = createWriter(result, checkpointStreamFactory.createCheckpointStateOutputStream(EXCLUSIVE));
	NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(threshold / 2, null), FreeingBufferRecycler.INSTANCE);
	writer.writeInput(new InputChannelInfo(1, 2), buffer);
	writer.completeOutput();
	writer.completeInput();
	assertTrue(result.isDone());
	assertEquals(0, checkpointsDir.list().length);
	assertEquals(0, sharedStateDir.list().length);
}
 

@Nullable
static Buffer sliceNextBuffer(ByteBuffer memory) {
	final int remaining = memory.remaining();

	// we only check the correct case where data is exhausted
	// all other cases can only occur if our write logic is wrong and will already throw
	// buffer underflow exceptions which will cause the read to fail.
	if (remaining == 0) {
		return null;
	}

	final boolean isEvent = memory.getShort() == HEADER_VALUE_IS_EVENT;
	final boolean isCompressed = memory.getShort() == BUFFER_IS_COMPRESSED;
	final int size = memory.getInt();

	memory.limit(memory.position() + size);
	ByteBuffer buf = memory.slice();
	memory.position(memory.limit());
	memory.limit(memory.capacity());

	MemorySegment memorySegment = MemorySegmentFactory.wrapOffHeapMemory(buf);

	Buffer.DataType dataType = isEvent ? Buffer.DataType.EVENT_BUFFER : Buffer.DataType.DATA_BUFFER;
	return new NetworkBuffer(memorySegment, FreeingBufferRecycler.INSTANCE, dataType, isCompressed, size);
}
 

private static BufferOrEvent createBuffer(int channel) {
	final int size = sizeCounter++;
	byte[] bytes = new byte[size];
	RND.nextBytes(bytes);

	MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	memory.put(0, bytes);

	Buffer buf = new NetworkBuffer(memory, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);

	// retain an additional time so it does not get disposed after being read by the input gate
	buf.retainBuffer();

	return new BufferOrEvent(buf, channel);
}
 

@Test
public void testRecyclingBuffers() throws Exception {
	ChannelStateCheckpointWriter writer = createWriter(new ChannelStateWriteResult());
	NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(10, null), FreeingBufferRecycler.INSTANCE);
	writer.writeInput(new InputChannelInfo(1, 2), buffer);
	assertTrue(buffer.isRecycled());
}
 

private static BufferOrEvent createBuffer(int channel, int size) {
	byte[] bytes = new byte[size];
	RND.nextBytes(bytes);

	MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	memory.put(0, bytes);

	Buffer buf = new NetworkBuffer(memory, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);

	// retain an additional time so it does not get disposed after being read by the input gate
	buf.retainBuffer();

	return new BufferOrEvent(buf, channel);
}
 

public static BufferOrEvent generateRandomBuffer(int size, int channelIndex) {
	MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	for (int i = 0; i < size; i++) {
		seg.put(i, (byte) i);
	}

	Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);
	return new BufferOrEvent(buf, channelIndex);
}
 

public static BufferOrEvent generateRandomBuffer(int size, int channelIndex) {
	MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	for (int i = 0; i < size; i++) {
		seg.put(i, (byte) i);
	}

	Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);
	return new BufferOrEvent(buf, channelIndex);
}
 

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 Buffer createDataBuffer(int size, Buffer.DataType dataType) {
	MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(size);
	NetworkBuffer buffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, dataType);
	for (int i = 0; i < size / 4; ++i) {
		buffer.writeInt(i);
	}

	return buffer;
}
 

private static BufferOrEvent createBuffer(int channel) {
	return new BufferOrEvent(
			new NetworkBuffer(MemorySegmentFactory.wrap(new byte[]{1, 2}), FreeingBufferRecycler.INSTANCE), channel);
}
 

private static BufferOrEvent createBuffer(int channel) {
	return new BufferOrEvent(
			new NetworkBuffer(MemorySegmentFactory.wrap(new byte[]{1, 2}), FreeingBufferRecycler.INSTANCE), new InputChannelInfo(0, channel));
}
 

@Override
public BufferOrEvent getNext() throws IOException {
	if (buffer.remaining() < HEADER_LENGTH) {
		buffer.compact();

		while (buffer.position() < HEADER_LENGTH) {
			if (fileChannel.read(buffer) == -1) {
				if (buffer.position() == 0) {
					// no trailing data
					return null;
				} else {
					throw new IOException("Found trailing incomplete buffer or event");
				}
			}
		}

		buffer.flip();
	}

	final int channel = buffer.getInt();
	final int length = buffer.getInt();
	final boolean isBuffer = buffer.get() == 0;

	if (isBuffer) {
		// deserialize buffer
		if (length > pageSize) {
			throw new IOException(String.format(
					"Spilled buffer (%d bytes) is larger than page size of (%d bytes)", length, pageSize));
		}

		MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(pageSize);

		int segPos = 0;
		int bytesRemaining = length;

		while (true) {
			int toCopy = Math.min(buffer.remaining(), bytesRemaining);
			if (toCopy > 0) {
				seg.put(segPos, buffer, toCopy);
				segPos += toCopy;
				bytesRemaining -= toCopy;
			}

			if (bytesRemaining == 0) {
				break;
			}
			else {
				buffer.clear();
				if (fileChannel.read(buffer) == -1) {
					throw new IOException("Found trailing incomplete buffer");
				}
				buffer.flip();
			}
		}

		Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
		buf.setSize(length);

		return new BufferOrEvent(buf, channel, true);
	}
	else {
		// deserialize event
		if (length > buffer.capacity() - HEADER_LENGTH) {
			throw new IOException("Event is too large");
		}

		if (buffer.remaining() < length) {
			buffer.compact();

			while (buffer.position() < length) {
				if (fileChannel.read(buffer) == -1) {
					throw new IOException("Found trailing incomplete event");
				}
			}

			buffer.flip();
		}

		int oldLimit = buffer.limit();
		buffer.limit(buffer.position() + length);
		AbstractEvent evt = EventSerializer.fromSerializedEvent(buffer, getClass().getClassLoader());
		buffer.limit(oldLimit);

		return new BufferOrEvent(evt, channel, true, length);
	}
}
 

private NetworkBuffer buffer() {
	return new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(10), FreeingBufferRecycler.INSTANCE);
}
 

private void write(ChannelStateCheckpointWriter writer, InputChannelInfo channelInfo, byte[] data) throws Exception {
	MemorySegment segment = wrap(data);
	NetworkBuffer buffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, Buffer.DataType.DATA_BUFFER, segment.size());
	writer.writeInput(channelInfo, buffer);
}
 

private static Buffer wrapWithBuffer(byte[] data) {
	NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(data.length, null), FreeingBufferRecycler.INSTANCE);
	buffer.writeBytes(data);
	return buffer;
}
 

public static BufferConsumer toBufferConsumer(AbstractEvent event) throws IOException {
	final ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(event);

	MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());

	return new BufferConsumer(data, FreeingBufferRecycler.INSTANCE, false);
}
 

/**
 * Allocates an un-pooled network buffer with the specific size.
 *
 * @param size The requested buffer size.
 * @param dataType The data type this buffer represents.
 * @return The un-pooled network buffer.
 */
Buffer allocateUnPooledNetworkBuffer(int size, Buffer.DataType dataType) {
	byte[] byteArray = new byte[size];
	MemorySegment memSeg = MemorySegmentFactory.wrap(byteArray);

	return new NetworkBuffer(memSeg, FreeingBufferRecycler.INSTANCE, dataType);
}
 

public static Buffer toBuffer(AbstractEvent event) throws IOException {
	final ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(event);

	MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());

	final Buffer buffer = new NetworkBuffer(data, FreeingBufferRecycler.INSTANCE, false);
	buffer.setSize(serializedEvent.remaining());

	return buffer;
}
 

public static BufferConsumer toBufferConsumer(AbstractEvent event) throws IOException {
	final ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(event);

	MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());

	return new BufferConsumer(data, FreeingBufferRecycler.INSTANCE, false);
}