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

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

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

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

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

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

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

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

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

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

/**
 * This tests the case where no additional partition buffers are used at the point when spilling
 * is triggered, testing that overflow bucket buffers are taken into account when deciding which
 * partition to spill.
 */
@Test
public void testSpillingFreesOnlyOverflowSegments() {
	final TypeSerializer<ByteValue> serializer = ByteValueSerializer.INSTANCE;
	final TypeComparator<ByteValue> buildComparator = new ValueComparator<>(true, ByteValue.class);
	final TypeComparator<ByteValue> probeComparator = new ValueComparator<>(true, ByteValue.class);
	
	@SuppressWarnings("unchecked")
	final TypePairComparator<ByteValue, ByteValue> pairComparator = Mockito.mock(TypePairComparator.class);
	
	try (final IOManager ioMan = new IOManagerAsync()) {
		final int pageSize = 32*1024;
		final int numSegments = 34;

		List<MemorySegment> memory = getMemory(numSegments, pageSize);

		MutableHashTable<ByteValue, ByteValue> table = new MutableHashTable<>(
				serializer, serializer, buildComparator, probeComparator,
				pairComparator, memory, ioMan, 1, false);

		table.open(new ByteValueIterator(100000000), new ByteValueIterator(1));
		
		table.close();
		
		checkNoTempFilesRemain(ioMan);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * get int from segments.
 *
 * @param segments target segments.
 * @param offset value offset.
 */
public static int getInt(MemorySegment[] segments, int offset) {
	if (inFirstSegment(segments, offset, 4)) {
		return segments[0].getInt(offset);
	} else {
		return getIntMultiSegments(segments, offset);
	}
}
 

/**
 * Gets an iterator over all KV records in this buffer in their logical order.
 *
 * @return An iterator returning the records in their logical order.
 */
public final MutableObjectIterator<Tuple2<BinaryRowData, BinaryRowData>> getIterator() {
	return new MutableObjectIterator<Tuple2<BinaryRowData, BinaryRowData>>() {
		private final int size = size();
		private int current = 0;

		private int currentSegment = 0;
		private int currentOffset = 0;

		private MemorySegment currentIndexSegment = sortIndex.get(0);

		@Override
		public Tuple2<BinaryRowData, BinaryRowData> next(Tuple2<BinaryRowData, BinaryRowData> kv) {
			if (this.current < this.size) {
				this.current++;
				if (this.currentOffset > lastIndexEntryOffset) {
					this.currentOffset = 0;
					this.currentIndexSegment = sortIndex.get(++this.currentSegment);
				}

				long pointer = this.currentIndexSegment.getLong(this.currentOffset);
				this.currentOffset += indexEntrySize;

				try {
					return getRecordFromBuffer(kv.f0, kv.f1, pointer);
				} catch (IOException ioe) {
					throw new RuntimeException(ioe);
				}
			} else {
				return null;
			}
		}

		@Override
		public Tuple2<BinaryRowData, BinaryRowData> next() {
			throw new RuntimeException("Not support!");
		}
	};
}
 

/**
 * unset bit from segments.
 *
 * @param segments target segments.
 * @param baseOffset bits base offset.
 * @param index bit index from base offset.
 */
public static void bitUnSet(MemorySegment[] segments, int baseOffset, int index) {
	if (segments.length == 1) {
		MemorySegment segment = segments[0];
		int offset = baseOffset + byteIndex(index);
		byte current = segment.get(offset);
		current &= ~(1 << (index & BIT_BYTE_INDEX_MASK));
		segment.put(offset, current);
	} else {
		bitUnSetMultiSegments(segments, baseOffset, index);
	}
}
 

@Override
public MemorySegment nextSegment() {
	if (this.segments.size() > 0) {
		return this.segments.remove(this.segments.size() - 1);
	} else {
		return null;
	}
}
 

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

@Override
public <T> void transform(
	K key,
	N namespace,
	T value,
	StateTransformationFunction<S, T> transformation) throws Exception {
	updateStat();
	MemorySegment keySegment = getKeySegment(key, namespace);
	int keyLen = keySegment.size();

	S oldState = getNode(keySegment, 0, keyLen);
	S newState = transformation.apply(oldState, value);
	byte[] stateBytes = skipListValueSerializer.serialize(newState);
	putValue(keySegment, 0, keyLen, stateBytes, false);
}
 

@Override
public MemorySegment nextSegment(MemorySegment current, int posInSegment) throws IOException {
	if (current != null) {
		writer.writeBlock(current);
	}

	final MemorySegment next = this.writer.getNextReturnedBlock();
	this.blockCount++;
	return next;
}
 

@Override
void returnSegmentToPool(MemorySegment segment) {
	if (segment.getClass() == HybridMemorySegment.class) {
		HybridMemorySegment heapSegment = (HybridMemorySegment) segment;
		availableMemory.add(heapSegment.getArray());
		heapSegment.free();
	}
	else {
		throw new IllegalArgumentException("Memory segment is not a " + HybridMemorySegment.class.getSimpleName());
	}
}
 

public void destroy() {
	synchronized (factoryLock) {
		isDestroyed = true;

		MemorySegment segment;
		while ((segment = availableMemorySegments.poll()) != null) {
			segment.free();
		}
	}
}
 

/**
 * Constructs {@link BufferConsumer} instance with content that can be changed by {@link BufferBuilder}.
 */
public BufferConsumer(
		MemorySegment memorySegment,
		BufferRecycler recycler,
		PositionMarker currentWriterPosition) {
	this(
		new NetworkBuffer(checkNotNull(memorySegment), checkNotNull(recycler), true),
		currentWriterPosition,
		0);
}
 

/** Same as above, but the number of bucket segments of the new table can be specified. */
private void rebuild(long newNumBucketSegments) throws IOException {
	// Get new bucket segments
	releaseBucketSegments();
	allocateBucketSegments((int)newNumBucketSegments);

	T record = buildSideSerializer.createInstance();
	try {
		EntryIterator iter = getEntryIterator();
		recordArea.resetAppendPosition();
		recordArea.setWritePosition(0);
		while ((record = iter.next(record)) != null && !closed) {
			final int hashCode = MathUtils.jenkinsHash(buildSideComparator.hash(record));
			final int bucket = hashCode & numBucketsMask;
			final int bucketSegmentIndex = bucket >>> numBucketsPerSegmentBits; // which segment contains the bucket
			final MemorySegment bucketSegment = bucketSegments[bucketSegmentIndex];
			final int bucketOffset = (bucket & numBucketsPerSegmentMask) << bucketSizeBits; // offset of the bucket in the segment
			final long firstPointer = bucketSegment.getLong(bucketOffset);

			long ptrToAppended = recordArea.noSeekAppendPointerAndRecord(firstPointer, record);
			bucketSegment.putLong(bucketOffset, ptrToAppended);
		}
		recordArea.freeSegmentsAfterAppendPosition();
		holes = 0;

	} catch (EOFException ex) {
		throw new RuntimeException("Bug in InPlaceMutableHashTable: we shouldn't get out of memory during a rebuild, " +
			"because we aren't allocating any new memory.");
	}
}
 

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

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

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

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

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