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

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

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 testAddWithFailingWriter() throws Exception {
	AsynchronousBufferFileWriter writer =
		new AsynchronousBufferFileWriter(ioManager.createChannel(), new RequestQueue<>());
	writer.close();

	exception.expect(IOException.class);

	Buffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(4096),
		FreeingBufferRecycler.INSTANCE);
	try {
		writer.writeBlock(buffer);
	} finally {
		if (!buffer.isRecycled()) {
			buffer.recycleBuffer();
			Assert.fail("buffer not recycled");
		}
		assertEquals("Shouln't increment number of outstanding requests.", 0, writer.getNumberOfOutstandingRequests());
	}
}
 

@Test
public void readPrematureEndOfFile2() throws Exception {
	final FileChannel fc = tmpFileChannel();
	final Buffer buffer = createTestBuffer();
	final MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledOffHeapMemory(buffer.getSize(), null);

	BufferReaderWriterUtil.writeToByteChannel(fc, buffer, BufferReaderWriterUtil.allocatedWriteBufferArray());
	fc.truncate(2); // less than a header size
	fc.position(0);

	try {
		BufferReaderWriterUtil.readFromByteChannel(
				fc, BufferReaderWriterUtil.allocatedHeaderBuffer(), readBuffer, FreeingBufferRecycler.INSTANCE);
		fail();
	}
	catch (IOException e) {
		// expected
	}
}
 

@Test
public void readPrematureEndOfFile1() throws Exception {
	final FileChannel fc = tmpFileChannel();
	final Buffer buffer = createTestBuffer();
	final MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledOffHeapMemory(buffer.getSize(), null);

	BufferReaderWriterUtil.writeToByteChannel(fc, buffer, BufferReaderWriterUtil.allocatedWriteBufferArray());
	fc.truncate(fc.position() - 1);
	fc.position(0);

	try {
		BufferReaderWriterUtil.readFromByteChannel(
				fc, BufferReaderWriterUtil.allocatedHeaderBuffer(), readBuffer, FreeingBufferRecycler.INSTANCE);
		fail();
	}
	catch (IOException e) {
		// expected
	}
}
 

@Test
public void readPrematureEndOfFile2() throws Exception {
	final FileChannel fc = tmpFileChannel();
	final Buffer buffer = createTestBuffer();
	final MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledOffHeapMemory(buffer.getSize(), null);

	BufferReaderWriterUtil.writeToByteChannel(fc, buffer, BufferReaderWriterUtil.allocatedWriteBufferArray());
	fc.truncate(2); // less than a header size
	fc.position(0);

	try {
		BufferReaderWriterUtil.readFromByteChannel(
				fc, BufferReaderWriterUtil.allocatedHeaderBuffer(), readBuffer, FreeingBufferRecycler.INSTANCE);
		fail();
	}
	catch (IOException e) {
		// expected
	}
}
 

@Test
public void readPrematureEndOfFile1() throws Exception {
	final FileChannel fc = tmpFileChannel();
	final Buffer buffer = createTestBuffer();
	final MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledOffHeapMemory(buffer.getSize(), null);

	BufferReaderWriterUtil.writeToByteChannel(fc, buffer, BufferReaderWriterUtil.allocatedWriteBufferArray());
	fc.truncate(fc.position() - 1);
	fc.position(0);

	try {
		BufferReaderWriterUtil.readFromByteChannel(
				fc, BufferReaderWriterUtil.allocatedHeaderBuffer(), readBuffer, FreeingBufferRecycler.INSTANCE);
		fail();
	}
	catch (IOException e) {
		// expected
	}
}
 

@Test
public void testAddWithFailingWriter() throws Exception {
	AsynchronousBufferFileWriter writer =
		new AsynchronousBufferFileWriter(ioManager.createChannel(), new RequestQueue<>());
	writer.close();

	exception.expect(IOException.class);

	Buffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(4096),
		FreeingBufferRecycler.INSTANCE);
	try {
		writer.writeBlock(buffer);
	} finally {
		if (!buffer.isRecycled()) {
			buffer.recycleBuffer();
			Assert.fail("buffer not recycled");
		}
		assertEquals("Shouln't increment number of outstanding requests.", 0, writer.getNumberOfOutstandingRequests());
	}
}
 

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

@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, int pageSize) {
	final int size = sizeCounter++;
	byte[] bytes = new byte[size];
	RND.nextBytes(bytes);

	MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(pageSize);
	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 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);
}
 

@Test
public void testAddWithFailingWriter() throws Exception {
	AsynchronousBufferFileWriter writer =
		new AsynchronousBufferFileWriter(ioManager.createChannel(), new RequestQueue<>());
	writer.close();

	exception.expect(IOException.class);

	Buffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(4096),
		FreeingBufferRecycler.INSTANCE);
	try {
		writer.writeBlock(buffer);
	} finally {
		if (!buffer.isRecycled()) {
			buffer.recycleBuffer();
			Assert.fail("buffer not recycled");
		}
		assertEquals("Shouln't increment number of outstanding requests.", 0, writer.getNumberOfOutstandingRequests());
	}
}
 

@Override
public Buffer next() {
	byte[] buffer = new byte[bufferSize];
	int bytesRead = read(buffer);
	checkState(bytesRead >= 0, "unexpected end of file, file = " + file.getFile() + ", offset=" + offset);
	offset += bytesRead;
	bytesToRead -= bytesRead;
	return new NetworkBuffer(wrap(buffer), FreeingBufferRecycler.INSTANCE, DATA_BUFFER, bytesRead);
}
 

@Test
public void writeReadFileChannel() throws Exception {
	final FileChannel fc = tmpFileChannel();
	final Buffer buffer = createTestBuffer();
	final MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledOffHeapMemory(buffer.getSize(), null);

	BufferReaderWriterUtil.writeToByteChannel(fc, buffer, BufferReaderWriterUtil.allocatedWriteBufferArray());
	fc.position(0);

	Buffer result = BufferReaderWriterUtil.readFromByteChannel(
			fc, BufferReaderWriterUtil.allocatedHeaderBuffer(), readBuffer, FreeingBufferRecycler.INSTANCE);

	validateTestBuffer(result);
}
 

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

@Test
public void testReadWritten() throws Exception {
	long checkpointId = 1L;

	InputChannelInfo inputChannelInfo = new InputChannelInfo(2, 3);
	byte[] inputChannelInfoData = randomBytes(1024);

	ResultSubpartitionInfo resultSubpartitionInfo = new ResultSubpartitionInfo(4, 5);
	byte[] resultSubpartitionInfoData = randomBytes(1024);

	ChannelStateWriteResult handles = write(
		checkpointId,
		singletonMap(inputChannelInfo, inputChannelInfoData),
		singletonMap(resultSubpartitionInfo, resultSubpartitionInfoData)
	);

	assertArrayEquals(inputChannelInfoData, read(
		toTaskStateSnapshot(handles),
		inputChannelInfoData.length,
		(reader, mem) -> reader.readInputData(inputChannelInfo, new NetworkBuffer(mem, FreeingBufferRecycler.INSTANCE))
	));

	assertArrayEquals(resultSubpartitionInfoData, read(
		toTaskStateSnapshot(handles),
		resultSubpartitionInfoData.length,
		(reader, mem) -> reader.readOutputData(resultSubpartitionInfo, new BufferBuilder(mem, FreeingBufferRecycler.INSTANCE))
	));
}
 

private void write(byte[] data, ChannelStateSerializerImpl serializer, OutputStream baos) throws IOException {
	DataOutputStream out = new DataOutputStream(baos);
	serializer.writeHeader(out);
	NetworkBuffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(data.length), FreeingBufferRecycler.INSTANCE);
	try {
		buffer.writeBytes(data);
		serializer.writeData(out, buffer);
		out.flush();
	} finally {
		buffer.release();
	}
}
 

private void readAndCheck(byte[] data, ChannelStateSerializerImpl serializer, ByteArrayInputStream is) throws IOException {
	serializer.readHeader(is);
	int size = serializer.readLength(is);
	assertEquals(data.length, size);
	NetworkBuffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(data.length), FreeingBufferRecycler.INSTANCE);
	try {
		int read = serializer.readData(is, wrap(buffer), size);
		assertEquals(size, read);
		assertArrayEquals(data, readBytes(buffer));
	} finally {
		buffer.release();
	}
}
 

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

@Test
public void writeReadFileChannel() throws Exception {
	final FileChannel fc = tmpFileChannel();
	final Buffer buffer = createTestBuffer();
	final MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledOffHeapMemory(buffer.getSize(), null);

	BufferReaderWriterUtil.writeToByteChannel(fc, buffer, BufferReaderWriterUtil.allocatedWriteBufferArray());
	fc.position(0);

	Buffer result = BufferReaderWriterUtil.readFromByteChannel(
			fc, BufferReaderWriterUtil.allocatedHeaderBuffer(), readBuffer, FreeingBufferRecycler.INSTANCE);

	validateTestBuffer(result);
}
 

private void testEncodeDecodeBuffer(boolean testReadOnlyBuffer) {
	NetworkBuffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(1024), FreeingBufferRecycler.INSTANCE);

	for (int i = 0; i < 1024; i += 4) {
		buffer.writeInt(i);
	}

	Buffer testBuffer = testReadOnlyBuffer ? buffer.readOnlySlice() : buffer;

	NettyMessage.BufferResponse expected = new NettyMessage.BufferResponse(
		testBuffer, random.nextInt(), new InputChannelID(), random.nextInt());
	NettyMessage.BufferResponse actual = encodeAndDecode(expected);

	// Netty 4.1 is not copying the messages, but retaining slices of them. BufferResponse actual is in this case
	// holding a reference to the buffer. Buffer will be recycled only once "actual" will be released.
	assertFalse(buffer.isRecycled());
	assertFalse(testBuffer.isRecycled());

	final ByteBuf retainedSlice = actual.getNettyBuffer();

	// Ensure not recycled and same size as original buffer
	assertEquals(1, retainedSlice.refCnt());
	assertEquals(1024, retainedSlice.readableBytes());

	for (int i = 0; i < 1024; i += 4) {
		assertEquals(i, retainedSlice.readInt());
	}

	// Release the retained slice
	actual.releaseBuffer();
	assertEquals(0, retainedSlice.refCnt());
	assertTrue(buffer.isRecycled());
	assertTrue(testBuffer.isRecycled());

	assertEquals(expected.sequenceNumber, actual.sequenceNumber);
	assertEquals(expected.receiverId, actual.receiverId);
	assertEquals(expected.backlog, actual.backlog);
}
 

@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 int header = memory.getInt();
	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);

	return bufferFromMemorySegment(
			memorySegment,
			FreeingBufferRecycler.INSTANCE,
			size,
			header == HEADER_VALUE_IS_EVENT);
}
 

private Buffer decompress(Buffer buffer) {
	MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE);
	Buffer compressedBuffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE);
	buffer.asByteBuf().readBytes(compressedBuffer.asByteBuf(), buffer.readableBytes());
	compressedBuffer.setCompressed(true);
	return DECOMPRESSOR.decompressToOriginalBuffer(compressedBuffer);
}
 

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

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

/**
 * Tests that the compressed buffer will be decompressed after calling {@link SingleInputGate#getNext()}.
 */
@Test
public void testGetCompressedBuffer() throws Exception {
	int bufferSize = 1024;
	String compressionCodec = "LZ4";
	BufferCompressor compressor = new BufferCompressor(bufferSize, compressionCodec);
	BufferDecompressor decompressor = new BufferDecompressor(bufferSize, compressionCodec);

	try (SingleInputGate inputGate = new SingleInputGateBuilder().setBufferDecompressor(decompressor).build()) {
		TestInputChannel inputChannel = new TestInputChannel(inputGate, 0);

		MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(bufferSize);
		for (int i = 0; i < bufferSize; i += 8) {
			segment.putLongLittleEndian(i, i);
		}
		Buffer uncompressedBuffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE);
		uncompressedBuffer.setSize(bufferSize);
		Buffer compressedBuffer = compressor.compressToOriginalBuffer(uncompressedBuffer);
		assertTrue(compressedBuffer.isCompressed());

		inputChannel.read(compressedBuffer);
		inputGate.setInputChannels(inputChannel);
		inputGate.notifyChannelNonEmpty(inputChannel);

		Optional<BufferOrEvent> bufferOrEvent = inputGate.getNext();
		assertTrue(bufferOrEvent.isPresent());
		assertTrue(bufferOrEvent.get().isBuffer());
		ByteBuffer buffer = bufferOrEvent.get().getBuffer().getNioBufferReadable().order(ByteOrder.LITTLE_ENDIAN);
		for (int i = 0; i < bufferSize; i += 8) {
			assertEquals(i, buffer.getLong());
		}
	}
}