org.apache.flink.core.memory.MemorySegmentFactory Java Examples

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

@SuppressWarnings("unchecked")
private <T extends Comparable<T>> void assertNormalizableKey(NormalizableKey<T> key1, NormalizableKey<T> key2, int len) {
	
	byte[] normalizedKeys = new byte[32];
	MemorySegment wrapper = MemorySegmentFactory.wrap(normalizedKeys);
	
	key1.copyNormalizedKey(wrapper, 0, len);
	key2.copyNormalizedKey(wrapper, len, len);
	
	for (int i = 0; i < len; i++) {
		int comp;
		int normKey1 = normalizedKeys[i] & 0xFF;
		int normKey2 = normalizedKeys[len + i] & 0xFF;
		
		if ((comp = (normKey1 - normKey2)) != 0) {
			if (Math.signum(key1.compareTo((T) key2)) != Math.signum(comp)) {
				Assert.fail("Normalized key comparison differs from actual key comparision");
			}
			return;
		}
	}
	if (key1.compareTo((T) key2) != 0 && key1.getMaxNormalizedKeyLen() <= len) {
		Assert.fail("Normalized key was not able to distinguish keys, " +
				"although it should as the length of it sufficies to uniquely identify them");
	}
}
 

public CompressedBlockChannelReader(
		IOManager ioManager,
		ID channel,
		LinkedBlockingQueue<MemorySegment> blockQueue,
		BlockCompressionFactory codecFactory,
		int preferBlockSize,
		int segmentSize) throws IOException {
	this.reader = ioManager.createBufferFileReader(channel, this);
	this.blockQueue = blockQueue;
	copyCompress = preferBlockSize > segmentSize * 2;
	int blockSize = copyCompress ? preferBlockSize : segmentSize;
	this.decompressor = codecFactory.getDecompressor();
	cause = new AtomicReference<>();

	if (copyCompress) {
		this.buf = new byte[blockSize];
		this.bufWrapper = ByteBuffer.wrap(buf);
	}

	BlockCompressor compressor = codecFactory.getCompressor();
	for (int i = 0; i < 2; i++) {
		MemorySegment segment = MemorySegmentFactory.wrap(new byte[compressor.getMaxCompressedSize(blockSize)]);
		reader.readInto(new NetworkBuffer(segment, this));
	}
}
 

/**
 * Test state map iterator illegal next call.
 */
@Test
public void testNamespaceNodeIteratorIllegalNextInvocation() {
	SkipListKeySerializer<Integer, Long> skipListKeySerializer =
		new SkipListKeySerializer<>(IntSerializer.INSTANCE, LongSerializer.INSTANCE);
	byte[] namespaceBytes = skipListKeySerializer.serializeNamespace(namespace);
	MemorySegment namespaceSegment = MemorySegmentFactory.wrap(namespaceBytes);
	Iterator<Long> iterator = stateMap.new NamespaceNodeIterator(namespaceSegment, 0, namespaceBytes.length);
	while (iterator.hasNext()) {
		iterator.next();
	}
	try {
		iterator.next();
		fail("Should have thrown NoSuchElementException.");
	} catch (NoSuchElementException e) {
		// expected
	}
}
 

void ensureAtLeastOneSegmentPresent() {
  if (!pool.isEmpty()) {
    //
    // the next allocation would succeeded because the pool is not empty
    //
    return;
  }
  if (hasRemainingCapacity()) {
    //
    // the next allocation would succeeded because the total allocated size is within the allowed
    // range
    //
    return;
  }
  //
  // we overdraft momentarily.
  //
  MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
  totalAllocatedMemory += PAGE_SIZE;
  pool.add(segment);
}
 

@Test
public void testByteBufferWrapping() {
	try {
		MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(1024);

		ByteBuffer buf1 = seg.wrap(13, 47);
		assertEquals(13, buf1.position());
		assertEquals(60, buf1.limit());
		assertEquals(47, buf1.remaining());

		ByteBuffer buf2 = seg.wrap(500, 267);
		assertEquals(500, buf2.position());
		assertEquals(767, buf2.limit());
		assertEquals(267, buf2.remaining());

		ByteBuffer buf3 = seg.wrap(0, 1024);
		assertEquals(0, buf3.position());
		assertEquals(1024, buf3.limit());
		assertEquals(1024, buf3.remaining());
	}
	catch (Exception e) {
		e.printStackTrace();
		Assert.fail(e.getMessage());
	}
}
 

@Test
public void compareTo() {
	assertEquals(0, fromString("   ").compareTo(blankString(3)));
	assertTrue(fromString("").compareTo(fromString("a")) < 0);
	assertTrue(fromString("abc").compareTo(fromString("ABC")) > 0);
	assertTrue(fromString("abc0").compareTo(fromString("abc")) > 0);
	assertEquals(0, fromString("abcabcabc").compareTo(fromString("abcabcabc")));
	assertTrue(fromString("aBcabcabc").compareTo(fromString("Abcabcabc")) > 0);
	assertTrue(fromString("Abcabcabc").compareTo(fromString("abcabcabC")) < 0);
	assertTrue(fromString("abcabcabc").compareTo(fromString("abcabcabC")) > 0);

	assertTrue(fromString("abc").compareTo(fromString("世界")) < 0);
	assertTrue(fromString("你好").compareTo(fromString("世界")) > 0);
	assertTrue(fromString("你好123").compareTo(fromString("你好122")) > 0);

	MemorySegment segment1 = MemorySegmentFactory.allocateUnpooledSegment(1024);
	MemorySegment segment2 = MemorySegmentFactory.allocateUnpooledSegment(1024);
	SortUtil.putStringNormalizedKey(fromString("abcabcabc"), segment1, 0, 9);
	SortUtil.putStringNormalizedKey(fromString("abcabcabC"), segment2, 0, 9);
	assertTrue(segment1.compare(segment2, 0, 0, 9) > 0);
	SortUtil.putStringNormalizedKey(fromString("abcab"), segment1, 0, 9);
	assertTrue(segment1.compare(segment2, 0, 0, 9) < 0);
}
 

public CompressedHeaderlessChannelWriterOutputView(
		BufferFileWriter writer,
		BlockCompressionFactory compressionCodecFactory,
		int compressionBlockSize) {
	super(compressionBlockSize, 0);

	this.compressionBlockSize = compressionBlockSize;
	buffer = MemorySegmentFactory.wrap(new byte[compressionBlockSize]);
	compressor = compressionCodecFactory.getCompressor();
	for (int i = 0; i < 2; i++) {
		compressedBuffers.add(MemorySegmentFactory.wrap(
				new byte[compressor.getMaxCompressedSize(compressionBlockSize)]));
	}
	this.writer = writer;

	try {
		advance();
	} catch (IOException ioex) {
		throw new RuntimeException(ioex);
	}
}
 

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

/**
 * Get a binary row consisting of 6 segments.
 * The bytes of the returned row is the same with the given input binary row.
 */
public static BinaryRowData getMultiSeg160BytesBinaryRow(BinaryRowData row160) {
	BinaryRowData multiSegRow160 = new BinaryRowData(2);
	MemorySegment[] segments = new MemorySegment[6];
	int baseOffset = 8;
	int posInSeg = baseOffset;
	int remainSize = 160;
	for (int i = 0; i < segments.length; i++) {
		segments[i] = MemorySegmentFactory.wrap(new byte[32]);
		int copy = Math.min(32 - posInSeg, remainSize);
		row160.getSegments()[0].copyTo(160 - remainSize, segments[i], posInSeg, copy);
		remainSize -= copy;
		posInSeg = 0;
	}
	multiSegRow160.pointTo(segments, baseOffset, 160);
	assertEquals(row160, multiSegRow160);
	return multiSegRow160;
}
 

private static void writeLongs(BoundedBlockingSubpartition partition, long nums) throws IOException {
	final MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE);

	long l = 0;
	while (nums > 0) {
		int pos = 0;
		for (; nums > 0 && pos <= memory.size() - 8; pos += 8) {
			memory.putLongBigEndian(pos, l++);
			nums--;
		}

		partition.add(new BufferConsumer(memory, (ignored) -> {}, pos, true));

		// we need to flush after every buffer as long as the add() contract is that
		// buffer are immediately added and can be filled further after that (for low latency
		// streaming data exchanges)
		partition.flush();
	}
}
 

@Test
public void testSerStringToKryo() throws IOException {
	KryoSerializer<BinaryString> serializer = new KryoSerializer<>(
		BinaryString.class, new ExecutionConfig());

	BinaryString string = BinaryString.fromString("hahahahaha");
	RandomAccessOutputView out = new RandomAccessOutputView(
		new MemorySegment[]{MemorySegmentFactory.wrap(new byte[1024])}, 64);
	serializer.serialize(string, out);

	RandomAccessInputView input = new RandomAccessInputView(
		new ArrayList<>(Collections.singletonList(out.getCurrentSegment())), 64, 64);
	BinaryString newStr = serializer.deserialize(input);

	assertEquals(string, newStr);
}
 

@Test
public void testEmptyString() {
	BinaryStringData str2 = fromString("hahahahah");
	BinaryStringData str3;
	{
		MemorySegment[] segments = new MemorySegment[2];
		segments[0] = MemorySegmentFactory.wrap(new byte[10]);
		segments[1] = MemorySegmentFactory.wrap(new byte[10]);
		str3 = BinaryStringData.fromAddress(segments, 15, 0);
	}

	assertTrue(BinaryStringData.EMPTY_UTF8.compareTo(str2) < 0);
	assertTrue(str2.compareTo(BinaryStringData.EMPTY_UTF8) > 0);

	assertEquals(0, BinaryStringData.EMPTY_UTF8.compareTo(str3));
	assertEquals(0, str3.compareTo(BinaryStringData.EMPTY_UTF8));

	assertNotEquals(BinaryStringData.EMPTY_UTF8, str2);
	assertNotEquals(str2, BinaryStringData.EMPTY_UTF8);

	assertEquals(BinaryStringData.EMPTY_UTF8, str3);
	assertEquals(str3, BinaryStringData.EMPTY_UTF8);
}
 

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

void ensureAtLeastOneSegmentPresent() {
  if (!pool.isEmpty()) {
    //
    // the next allocation would succeeded because the pool is not empty
    //
    return;
  }
  if (hasRemainingCapacity()) {
    //
    // the next allocation would succeeded because the total allocated size is within the allowed
    // range
    //
    return;
  }
  //
  // we overdraft momentarily.
  //
  MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
  totalAllocatedMemory += PAGE_SIZE;
  pool.add(segment);
}
 

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

public MemorySegment setupNormalizedKeysMemSegment(T[] data, int normKeyLen, TypeComparator<T> comparator) {
	MemorySegment memSeg = MemorySegmentFactory.allocateUnpooledSegment(2048);

	// Setup normalized Keys in the memory segment
	int offset = 0;
	for (T e : data) {
		comparator.putNormalizedKey(e, memSeg, offset, normKeyLen);
		offset += normKeyLen;
	}
	return memSeg;
}
 

@Test
public void testGenericObject() throws Exception {

	GenericTypeInfo<MyObj> info = new GenericTypeInfo<>(MyObj.class);
	TypeSerializer<MyObj> genericSerializer = info.createSerializer(new ExecutionConfig());

	BinaryRow row = new BinaryRow(4);
	BinaryRowWriter writer = new BinaryRowWriter(row);
	writer.writeInt(0, 0);

	BinaryGeneric<MyObj> myObj1 = new BinaryGeneric<>(new MyObj(0, 1), genericSerializer);
	writer.writeGeneric(1, myObj1);
	BinaryGeneric<MyObj> myObj2 = new BinaryGeneric<>(new MyObj(123, 5.0), genericSerializer);
	myObj2.ensureMaterialized();
	writer.writeGeneric(2, myObj2);
	BinaryGeneric<MyObj> myObj3 = new BinaryGeneric<>(new MyObj(1, 1), genericSerializer);
	writer.writeGeneric(3, myObj3);
	writer.complete();

	assertTestGenericObjectRow(row, genericSerializer);

	// getBytes from var-length memorySegments.
	BinaryRowSerializer serializer = new BinaryRowSerializer(4);
	MemorySegment[] memorySegments = new MemorySegment[3];
	ArrayList<MemorySegment> memorySegmentList = new ArrayList<>();
	for (int i = 0; i < 3; i++) {
		memorySegments[i] = MemorySegmentFactory.wrap(new byte[64]);
		memorySegmentList.add(memorySegments[i]);
	}
	RandomAccessOutputView out = new RandomAccessOutputView(memorySegments, 64);
	serializer.serializeToPages(row, out);

	BinaryRow mapRow = serializer.mapFromPages(new RandomAccessInputView(memorySegmentList, 64));
	assertTestGenericObjectRow(mapRow, genericSerializer);
}
 

/**
 * Creates a {@link BinaryStringData} instance from the given UTF-8 bytes with offset and number of bytes.
 */
public static BinaryStringData fromBytes(byte[] bytes, int offset, int numBytes) {
	return new BinaryStringData(
		new MemorySegment[] {MemorySegmentFactory.wrap(bytes)},
		offset,
		numBytes);
}
 

private static List<MemorySegment> getMemory(int numPages, int pageSize) {
	List<MemorySegment> memory = new ArrayList<MemorySegment>();
	
	for (int i = 0; i < numPages; i++) {
		memory.add(MemorySegmentFactory.allocateUnpooledSegment(pageSize));
	}
	
	return memory;
}
 

@Override
public MemorySegment nextSegment() {
	allocated++;
	if (allocated <= segments.size()) {
		return segments.get(allocated - 1);
	} else {
		return MemorySegmentFactory.wrap(new byte[pageSize()]);
	}
}
 

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

private BinaryArray deserializeReuse(BinaryArray reuse, DataInputView source) throws IOException {
	int length = source.readInt();
	byte[] bytes = new byte[length];
	source.readFully(bytes);
	reuse.pointTo(MemorySegmentFactory.wrap(bytes), 0, bytes.length);
	return reuse;
}
 

public BinaryRowWriter(BinaryRow row, int initialSize) {
	this.nullBitsSizeInBytes = BinaryRow.calculateBitSetWidthInBytes(row.getArity());
	this.fixedSize = row.getFixedLengthPartSize();
	this.cursor = fixedSize;

	this.segment = MemorySegmentFactory.wrap(new byte[fixedSize + initialSize]);
	this.row = row;
	this.row.pointTo(segment, 0, segment.size());
}
 

@Override
public void materialize() {
	byte[] bytes = StringUtf8Utils.encodeUTF8(javaObject);
	segments = new MemorySegment[] {MemorySegmentFactory.wrap(bytes)};
	offset = 0;
	sizeInBytes = bytes.length;
}
 

private void putKeySpace(KeySpace keySpace, MemorySegment memorySegment, int offset) {
	SkipListUtils.putLevelAndNodeStatus(memorySegment, offset, keySpace.level, keySpace.status);
	SkipListUtils.putKeyLen(memorySegment, offset, keySpace.keyData.length);
	SkipListUtils.putValuePointer(memorySegment, offset, keySpace.valuePointer);
	SkipListUtils.putNextKeyPointer(memorySegment, offset, keySpace.nextKeyPointer);
	for (int i = 1; i <= keySpace.nextIndexNodes.length; i++) {
		SkipListUtils.putNextIndexNode(memorySegment, offset, i, keySpace.nextIndexNodes[i - 1]);
	}
	for (int i = 1; i <= keySpace.prevIndexNodes.length; i++) {
		SkipListUtils.putPrevIndexNode(memorySegment, offset, keySpace.level, i, keySpace.prevIndexNodes[i - 1]);
	}
	SkipListUtils.putKeyData(memorySegment, offset, MemorySegmentFactory.wrap(keySpace.keyData),
		0, keySpace.keyData.length, keySpace.level);
}
 

private static List<MemorySegment> getMemory(int numPages, int pageSize) {
	List<MemorySegment> memory = new ArrayList<>();

	for (int i = 0; i < numPages; i++) {
		memory.add(MemorySegmentFactory.allocateUnpooledSegment(pageSize));
	}

	return memory;
}
 

public BufferDecompressor(int bufferSize, String factoryName) {
	checkArgument(bufferSize > 0);
	checkNotNull(factoryName);

	// the decompressed data size should be never larger than the configured buffer size
	final byte[] heapBuffer = new byte[bufferSize];
	this.internalBuffer = new NetworkBuffer(MemorySegmentFactory.wrap(heapBuffer), FreeingBufferRecycler.INSTANCE);
	this.blockDecompressor = BlockCompressionFactory.createBlockCompressionFactory(factoryName).getDecompressor();
}
 

private void verifyGetValueSpace(ValueSpace valueSpace, MemorySegment memorySegment, int offset) {
	assertEquals(valueSpace.version, SkipListUtils.getValueVersion(memorySegment, offset));
	assertEquals(valueSpace.keyPointer, SkipListUtils.getKeyPointer(memorySegment, offset));
	assertEquals(valueSpace.nextValuePointer, SkipListUtils.getNextValuePointer(memorySegment, offset));
	assertEquals(valueSpace.valueData.length, SkipListUtils.getValueLen(memorySegment, offset));
	int valueDataOffset = SkipListUtils.getValueMetaLen();
	MemorySegment valueDataSegment = MemorySegmentFactory.wrap(valueSpace.valueData);
	assertEquals(0,
		memorySegment.compare(valueDataSegment, offset + valueDataOffset, 0, valueSpace.valueData.length));
}