Java Code Examples for org.apache.flink.core.memory.MemorySegment

public List<MemorySegment> close() throws IOException {
	if (LOG.isDebugEnabled()) {
		LOG.debug("Spilling Resettable Iterator closing. Stored " + this.elementCount + " records.");
	}

	this.inView = null;
	
	final List<MemorySegment> memory = this.buffer.close();
	memory.addAll(this.memorySegments);
	this.memorySegments.clear();
	
	if (this.releaseMemoryOnClose) {
		this.memoryManager.release(memory);
		return Collections.emptyList();
	} else {
		return memory;
	}
}
 

/**
 * Write the meta and data for the key to the given node.
 *
 * @param node          the node for the key to write.
 * @param level         level of this node.
 * @param keySegment    memory segment storing the key.
 * @param keyOffset     offset of key in memory segment.
 * @param keyLen        length of the key.
 * @param valuePointer  pointer to value.
 * @param nextNode      next node on level 0.
 */
private void doWriteKey(
	long node,
	int level,
	MemorySegment keySegment,
	int keyOffset,
	int keyLen,
	long valuePointer,
	long nextNode) {
	Node nodeStorage = getNodeSegmentAndOffset(node);
	MemorySegment segment = nodeStorage.nodeSegment;
	int offsetInSegment = nodeStorage.nodeOffset;

	SkipListUtils.putLevelAndNodeStatus(segment, offsetInSegment, level, NodeStatus.PUT);
	SkipListUtils.putKeyLen(segment, offsetInSegment, keyLen);
	SkipListUtils.putValuePointer(segment, offsetInSegment, valuePointer);
	SkipListUtils.putNextKeyPointer(segment, offsetInSegment, nextNode);
	SkipListUtils.putKeyData(segment, offsetInSegment, keySegment, keyOffset, keyLen, level);
}
 

/**
 * This is the method called by the partitions to request memory to serialize records.
 * It automatically spills partitions, if memory runs out. 
 * 
 * @return The next available memory segment.
 */
@Override
public MemorySegment nextSegment() {
	final MemorySegment seg = getNextBuffer();
	if (seg != null) {
		return seg;
	} else {
		try {
			spillPartition();
		} catch (IOException ioex) {
			throw new RuntimeException("Error spilling Hash Join Partition" + (ioex.getMessage() == null ?
				"." : ": " + ioex.getMessage()), ioex);
		}
		
		MemorySegment fromSpill = getNextBuffer();
		if (fromSpill == null) {
			throw new RuntimeException("BUG in Hybrid Hash Join: Spilling did not free a buffer.");
		} else {
			return fromSpill;
		}
	}
}
 

@Test
public void testConcurrentModificationOnce() {
	try {
		final int numSegments = 10000;
		final int segmentSize = 4096;

		MemoryManager memMan = new MemoryManager(numSegments * segmentSize, 1, segmentSize, MemoryType.HEAP, true);

		ArrayList<MemorySegment> segs = new ListWithConcModExceptionOnFirstAccess<>();
		memMan.allocatePages(this, segs, numSegments);

		memMan.release(segs);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

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

/**
 * Creates a new partition, in memory, with one buffer. 
 * 
 * @param serializer Serializer for T.
 * @param partitionNumber The number of the partition.
 * @param memSource memory pool
 * @param pageSize segment size in bytes
 * @param pageSizeInBits
 */
public InMemoryPartition(TypeSerializer<T> serializer, int partitionNumber,
		ListMemorySegmentSource memSource, int pageSize, int pageSizeInBits)
{
	this.overflowSegments = new MemorySegment[2];
	this.numOverflowSegments = 0;
	this.nextOverflowBucket = 0;
	
	this.serializer = serializer;
	this.partitionPages = new ArrayList<MemorySegment>(64);
	this.availableMemory = memSource;
	
	this.partitionNumber = partitionNumber;
	
	// add the first segment
	this.partitionPages.add(memSource.nextSegment());
	// empty partitions have no garbage
	this.compacted = true;
	
	this.pageSize = pageSize;
	
	this.pageSizeInBits = pageSizeInBits;
	
	this.writeView = new WriteView(this.partitionPages, memSource, pageSize, pageSizeInBits);
	this.readView = new ReadView(this.partitionPages, pageSize, pageSizeInBits);
}
 

@Test
public void allocateTooMuch() {
	try {
		final AbstractInvokable mockInvoke = new DummyInvokable();

		List<MemorySegment> segs = this.memoryManager.allocatePages(mockInvoke, NUM_PAGES);

		try {
			this.memoryManager.allocatePages(mockInvoke, 1);
			Assert.fail("Expected MemoryAllocationException.");
		} catch (MemoryAllocationException maex) {
			// expected
		}

		Assert.assertTrue("The previously allocated segments were not valid any more.",
																allMemorySegmentsValid(segs));

		this.memoryManager.releaseAll(mockInvoke);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

public void testNormalizedKeysEquals(boolean halfLength) {
	try {
		TypeComparator<T> comparator = getComparator(true);
		T[] data = getSortedData();
		int normKeyLen = getNormKeyLen(halfLength, data, comparator);

		MemorySegment memSeg1 = setupNormalizedKeysMemSegment(data, normKeyLen, comparator);
		MemorySegment memSeg2 = setupNormalizedKeysMemSegment(data, normKeyLen, comparator);

		for (int i = 0; i < data.length; i++) {
			assertTrue(memSeg1.compare(memSeg2, i * normKeyLen, i * normKeyLen, normKeyLen) == 0);
		}
	} catch (Exception e) {
		System.err.println(e.getMessage());
		e.printStackTrace();
		fail("Exception in test: " + e.getMessage());
	}
}
 

@Test
public void testSetAndGet() {
	MemorySegment segment = MemorySegmentFactory.wrap(new byte[80]);
	BinaryRowData row = new BinaryRowData(9);
	row.pointTo(segment, 0, 80);
	row.setNullAt(0);
	row.setInt(1, 11);
	row.setLong(2, 22);
	row.setDouble(3, 33);
	row.setBoolean(4, true);
	row.setShort(5, (short) 55);
	row.setByte(6, (byte) 66);
	row.setFloat(7, 77f);

	assertEquals(33d, (long) row.getDouble(3), 0);
	assertEquals(11, row.getInt(1));
	assertTrue(row.isNullAt(0));
	assertEquals(55, row.getShort(5));
	assertEquals(22, row.getLong(2));
	assertTrue(row.getBoolean(4));
	assertEquals((byte) 66, row.getByte(6));
	assertEquals(77f, row.getFloat(7), 0);
}
 

public List<MemorySegment> close() throws IOException {
	if (LOG.isDebugEnabled()) {
		LOG.debug("Spilling Resettable Iterator closing. Stored " + this.elementCount + " records.");
	}

	this.inView = null;
	
	final List<MemorySegment> memory = this.buffer.close();
	memory.addAll(this.memorySegments);
	this.memorySegments.clear();
	
	if (this.releaseMemoryOnClose) {
		this.memoryManager.release(memory);
		return Collections.emptyList();
	} else {
		return memory;
	}
}
 

private NotificationResult fireBufferAvailableNotification(BufferListener listener, MemorySegment segment) {
	// We do not know which locks have been acquired before the recycle() or are needed in the
	// notification and which other threads also access them.
	// -> call notifyBufferAvailable() outside of the synchronized block to avoid a deadlock (FLINK-9676)
	NotificationResult notificationResult = listener.notifyBufferAvailable(new NetworkBuffer(segment, this));
	if (notificationResult.needsMoreBuffers()) {
		synchronized (availableMemorySegments) {
			if (isDestroyed) {
				// cleanup tasks how they would have been done if we only had one synchronized block
				listener.notifyBufferDestroyed();
			} else {
				registeredListeners.add(listener);
			}
		}
	}
	return notificationResult;
}
 

@Override
public void putNormalizedKey(IntPair record, MemorySegment target, int offset, int len) {
	// see IntValue for a documentation of the logic
	final int value = record.getKey() - Integer.MIN_VALUE;
	
	if (len == 4) {
		target.putIntBigEndian(offset, value);
	}
	else if (len <= 0) {
	}
	else if (len < 4) {
		for (int i = 0; len > 0; len--, i++) {
			target.put(offset + i, (byte) ((value >>> ((3-i)<<3)) & 0xff));
		}
	}
	else {
		target.putIntBigEndian(offset, value);
		for (int i = 4; i < len; i++) {
			target.put(offset + i, (byte) 0);
		}
	}
}
 

void release(MemorySegment segment) {
  if (totalAllocatedMemory > inMemoryBufferSize) {
    //
    // we previously overdraft.
    //
    segment.free();
    totalAllocatedMemory -= PAGE_SIZE;
    return;
  }
  pool.add(segment);
}
 

@Test
public void testHashAndCopy() throws IOException {
	MemorySegment[] segments = new MemorySegment[3];
	for (int i = 0; i < 3; i++) {
		segments[i] = MemorySegmentFactory.wrap(new byte[64]);
	}
	RandomAccessOutputView out = new RandomAccessOutputView(segments, 64);
	BinaryRowDataSerializer serializer = new BinaryRowDataSerializer(2);

	BinaryRowData row = new BinaryRowData(2);
	BinaryRowWriter writer = new BinaryRowWriter(row);
	writer.writeString(0, fromString("hahahahahahahahahahahahahahahahahahahhahahahahahahahahah"));
	writer.writeString(1, fromString("hahahahahahahahahahahahahahahahahahahhahahahahahahahahaa"));
	writer.complete();
	serializer.serializeToPages(row, out);

	ArrayList<MemorySegment> segmentList = new ArrayList<>(Arrays.asList(segments));
	RandomAccessInputView input = new RandomAccessInputView(segmentList, 64, 64);

	BinaryRowData mapRow = serializer.mapFromPages(input);
	assertEquals(row, mapRow);
	assertEquals(row.getString(0), mapRow.getString(0));
	assertEquals(row.getString(1), mapRow.getString(1));
	assertNotEquals(row.getString(0), mapRow.getString(1));

	// test if the hash code before and after serialization are the same
	assertEquals(row.hashCode(), mapRow.hashCode());
	assertEquals(row.getString(0).hashCode(), mapRow.getString(0).hashCode());
	assertEquals(row.getString(1).hashCode(), mapRow.getString(1).hashCode());

	// test if the copy method produce a row with the same contents
	assertEquals(row.copy(), mapRow.copy());
	assertEquals(
		((BinaryStringData) row.getString(0)).copy(),
		((BinaryStringData) mapRow.getString(0)).copy());
	assertEquals(
		((BinaryStringData) row.getString(1)).copy(),
		((BinaryStringData) mapRow.getString(1)).copy());
}
 

@Before
public void setUp() throws Exception {
	final MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE);
	buffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, true, 0);
	for (int i = 0; i < DATA_SIZE; ++i) {
		buffer.writeByte(i);
	}
}
 

private static MemorySegment[] splitBytes(byte[] bytes, int baseOffset) {
	int newSize = (bytes.length + 1) / 2 + baseOffset;
	MemorySegment[] ret = new MemorySegment[2];
	ret[0] = MemorySegmentFactory.wrap(new byte[newSize]);
	ret[1] = MemorySegmentFactory.wrap(new byte[newSize]);

	ret[0].put(baseOffset, bytes, 0, newSize - baseOffset);
	ret[1].put(0, bytes, newSize - baseOffset, bytes.length - (newSize - baseOffset));
	return ret;
}
 

static int fillBufferWithAscendingNumbers(Buffer buffer, int currentNumber, int size) {
	checkArgument(size % 4 == 0);

	MemorySegment segment = buffer.getMemorySegment();

	for (int i = 0; i < size; i += 4) {
		segment.putInt(i, currentNumber++);
	}
	buffer.setSize(size);

	return currentNumber;
}
 

public void sortAndSpill(
		ArrayList<MemorySegment> recordBufferSegments,
		long numElements,
		MemorySegmentPool pool) throws IOException {

	// 1. sort buffer
	BinaryKVInMemorySortBuffer buffer =
			BinaryKVInMemorySortBuffer.createBuffer(
					nKeyComputer, keySerializer, valueSerializer, comparator,
					recordBufferSegments, numElements, pool);
	this.sorter.sort(buffer);

	// 2. spill
	FileIOChannel.ID channel = enumerator.next();
	channelManager.addChannel(channel);

	AbstractChannelWriterOutputView output = null;
	int bytesInLastBuffer;
	int blockCount;
	try {
		numSpillFiles++;
		output = FileChannelUtil.createOutputView(ioManager, channel, compressionEnable,
				compressionCodecFactory, compressionBlockSize, pageSize);
		buffer.writeToOutput(output);
		spillInBytes += output.getNumBytes();
		spillInCompressedBytes += output.getNumCompressedBytes();
		bytesInLastBuffer = output.close();
		blockCount = output.getBlockCount();
		LOG.info("here spill the {}th kv external buffer data with {} bytes and {} compressed bytes",
				numSpillFiles, spillInBytes, spillInCompressedBytes);
	} catch (IOException e) {
		if (output != null) {
			output.close();
			output.getChannel().deleteChannel();
		}
		throw e;
	}
	channelIDs.add(new ChannelWithMeta(channel, blockCount, bytesInLastBuffer));
}
 

@Override
public E next(E reuse) throws IOException
{
	try {
		return this.accessors.deserialize(reuse, this.inView);
	} catch (EOFException eofex) {
		final List<MemorySegment> freeMem = this.inView.close();
		if (this.freeMemTarget != null) {
			this.freeMemTarget.addAll(freeMem);
		}
		return null;
	}
}
 

private BuildSideBuffer(MemorySegment initialSegment, MemorySegmentSource memSource) {
	super(initialSegment, initialSegment.size(), 0);
	
	this.targetList = new ArrayList<MemorySegment>();
	this.memSource = memSource;
	this.sizeBits = MathUtils.log2strict(initialSegment.size());
}
 

HashTableBloomFilter(MemorySegment[] buffers, long numRecords) {
	checkArgument(buffers != null && buffers.length > 0);
	this.buffers = buffers;
	this.numBuffers = buffers.length;
	checkArgument(MathUtils.isPowerOf2(numBuffers));
	this.numBuffersMask = numBuffers - 1;
	int bufferSize = buffers[0].size();
	this.filter = new BloomFilter((int) (numRecords / numBuffers), buffers[0].size());
	filter.setBitsLocation(buffers[0], 0);

	// We assume that a BloomFilter can contain up to 2.44 elements per byte.
	// fpp roughly equal 0.2
	this.maxSize = (int) ((numBuffers * bufferSize) * 2.44);
}
 

private static void setTwoByteSlowly(
		MemorySegment[] segments, int segSize, int segNum, int segOffset, int b1, int b2) {
	MemorySegment segment = segments[segNum];
	segment.put(segOffset, (byte) (LITTLE_ENDIAN ? b1 : b2));
	segOffset++;
	if (segOffset == segSize) {
		segment = segments[++segNum];
		segOffset = 0;
	}
	segment.put(segOffset, (byte) (LITTLE_ENDIAN ? b2 : b1));
}
 

@Override
public List<MemorySegment> getFreeMemory() {
	if (!this.closed) {
		throw new IllegalStateException("Cannot return memory while InPlaceMutableHashTable is open.");
	}

	return freeMemorySegments;
}
 

@Override
public List<MemorySegment> getFullSegments() {
	synchronized (this.closeLock) {
		if (!this.isClosed() || this.requestsNotReturned.get() > 0) {
			throw new IllegalStateException("Full segments can only be obtained after the reader was properly closed.");
		}
	}
	
	return this.returnBuffers;
}
 

private boolean allMemorySegmentsValid(List<MemorySegment> memSegs) {
	for (MemorySegment seg : memSegs) {
		if (seg.isFreed()) {
			return false;
		}
	}
	return true;
}
 

@Test
public void testCloseAndDeleteOutputView() {
	final IOManager ioManager = new IOManagerAsync();
	try {
		MemoryManager memMan = new MemoryManager(4 * 16*1024, 1, 16*1024, MemoryType.HEAP, true);
		List<MemorySegment> memory = new ArrayList<MemorySegment>();
		memMan.allocatePages(new DummyInvokable(), memory, 4);
		
		FileIOChannel.ID channel = ioManager.createChannel();
		BlockChannelWriter<MemorySegment> writer = ioManager.createBlockChannelWriter(channel);
		
		FileChannelOutputView out = new FileChannelOutputView(writer, memMan, memory, memMan.getPageSize());
		new StringValue("Some test text").write(out);
		
		// close for the first time, make sure all memory returns
		out.close();
		assertTrue(memMan.verifyEmpty());
		
		// close again, should not cause an exception
		out.close();
		
		// delete, make sure file is removed
		out.closeAndDelete();
		assertFalse(new File(channel.getPath()).exists());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
	finally {
		ioManager.shutdown();
	}
}
 

@Override
protected MemorySegment nextSegment(MemorySegment current, int positionInCurrent) throws EOFException {
	final MemorySegment next = this.memorySource.nextSegment();
	if (next != null) {
		this.fullSegments.add(next);
		this.segmentNum++;
		return next;
	} else {
		throw new EOFException("Can't collect further: memorySource depleted");
	}
}
 

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

/**
 * resets overflow bucket counters and returns freed memory and should only be used for resizing
 * 
 * @return freed memory segments
 */
public ArrayList<MemorySegment> resetOverflowBuckets() {
	this.numOverflowSegments = 0;
	this.nextOverflowBucket = 0;
	
	ArrayList<MemorySegment> result = new ArrayList<MemorySegment>(this.overflowSegments.length);
	for(int i = 0; i < this.overflowSegments.length; i++) {
		if(this.overflowSegments[i] != null) {
			result.add(this.overflowSegments[i]);
		}
	}
	this.overflowSegments = new MemorySegment[2];
	return result;
}
 

private void setBucket(
		MemorySegment bucket, MemorySegment[] overflowSegments,
		int bucketInSegmentOffset) {
	this.bucketSegment = bucket;
	this.overflowSegments = overflowSegments;
	this.bucketInSegmentOffset = bucketInSegmentOffset;
	this.pointerOffset = bucketInSegmentOffset + BinaryHashBucketArea.BUCKET_POINTER_START_OFFSET;
	this.countInBucket = bucket.getShort(bucketInSegmentOffset + BinaryHashBucketArea.HEADER_COUNT_OFFSET);
	this.numInBucket = 0;
	// reset probedSet with probedFlags offset in this bucket.
	this.probedSet.setMemorySegment(bucketSegment, this.bucketInSegmentOffset + BinaryHashBucketArea.PROBED_FLAG_OFFSET);
}