Java Code Examples for org.apache.flink.core.memory.MemorySegment#getOwner()

/**
 * Tries to release the memory for the specified segment.
 *
 * <p>If the segment has already been released, it is only freed. If it is null or has no owner, the request is simply ignored.
 * The segment is only freed and made eligible for reclamation by the GC. The segment will be returned to
 * the memory pool, increasing its available limit for the later allocations.
 *
 * @param segment The segment to be released.
 */
public void release(MemorySegment segment) {
	Preconditions.checkState(!isShutDown, "Memory manager has been shut down.");

	// check if segment is null or has already been freed
	if (segment == null || segment.getOwner() == null) {
		return;
	}

	// remove the reference in the map for the owner
	try {
		allocatedSegments.computeIfPresent(segment.getOwner(), (o, segsForOwner) -> {
			segment.free();
			segsForOwner.remove(segment);
			return segsForOwner.isEmpty() ? null : segsForOwner;
		});
	}
	catch (Throwable t) {
		throw new RuntimeException("Error removing book-keeping reference to allocated memory segment.", t);
	}
}
 

private void returnMemory(
		LazyMemorySegmentPool pool, MemorySegment[] buckets, MemorySegment[] overflowSegments) {
	pool.returnAll(Arrays.asList(buckets));
	for (MemorySegment segment : overflowSegments) {
		if (segment != null &&
				// except stealing from heap.
				segment.getOwner() != this) {
			pool.returnPage(segment);
		}
	}
}
 

private MemorySegment releaseSegmentsForOwnerUntilNextOwner(
		MemorySegment firstSeg,
		Iterator<MemorySegment> segmentsIterator) {
	AtomicReference<MemorySegment> nextOwnerMemorySegment = new AtomicReference<>();
	Object owner = firstSeg.getOwner();
	allocatedSegments.compute(owner, (o, segsForOwner) -> {
		freeSegment(firstSeg, segsForOwner);
		while (segmentsIterator.hasNext()) {
			MemorySegment segment = segmentsIterator.next();
			try {
				if (segment == null || segment.isFreed()) {
					continue;
				}
				Object nextOwner = segment.getOwner();
				if (nextOwner != owner) {
					nextOwnerMemorySegment.set(segment);
					break;
				}
				freeSegment(segment, segsForOwner);
			} catch (Throwable t) {
				throw new RuntimeException(
					"Error removing book-keeping reference to allocated memory segment.", t);
			}
		}
		return segsForOwner == null || segsForOwner.isEmpty() ? null : segsForOwner;
	});
	return nextOwnerMemorySegment.get();
}
 

/**
 * Tries to release the memory for the specified segment. If the segment has already been released or
 * is null, the request is simply ignored.
 *
 * <p>If the memory manager manages pre-allocated memory, the memory segment goes back to the memory pool.
 * Otherwise, the segment is only freed and made eligible for reclamation by the GC.
 *
 * @param segment The segment to be released.
 * @throws IllegalArgumentException Thrown, if the given segment is of an incompatible type.
 */
public void release(MemorySegment segment) {
	// check if segment is null or has already been freed
	if (segment == null || segment.getOwner() == null) {
		return;
	}

	final Object owner = segment.getOwner();

	// -------------------- BEGIN CRITICAL SECTION -------------------
	synchronized (lock) {
		// prevent double return to this memory manager
		if (segment.isFreed()) {
			return;
		}
		if (isShutDown) {
			throw new IllegalStateException("Memory manager has been shut down.");
		}

		// remove the reference in the map for the owner
		try {
			Set<MemorySegment> segsForOwner = this.allocatedSegments.get(owner);

			if (segsForOwner != null) {
				segsForOwner.remove(segment);
				if (segsForOwner.isEmpty()) {
					this.allocatedSegments.remove(owner);
				}
			}

			if (isPreAllocated) {
				// release the memory in any case
				memoryPool.returnSegmentToPool(segment);
			}
			else {
				segment.free();
				numNonAllocatedPages++;
			}
		}
		catch (Throwable t) {
			throw new RuntimeException("Error removing book-keeping reference to allocated memory segment.", t);
		}
	}
	// -------------------- END CRITICAL SECTION -------------------
}
 

/**
 * Tries to release many memory segments together.
 *
 * <p>If the memory manager manages pre-allocated memory, the memory segment goes back to the memory pool.
 * Otherwise, the segment is only freed and made eligible for reclamation by the GC.
 *
 * @param segments The segments to be released.
 * @throws NullPointerException Thrown, if the given collection is null.
 * @throws IllegalArgumentException Thrown, id the segments are of an incompatible type.
 */
public void release(Collection<MemorySegment> segments) {
	if (segments == null) {
		return;
	}

	// -------------------- BEGIN CRITICAL SECTION -------------------
	synchronized (lock) {
		if (isShutDown) {
			throw new IllegalStateException("Memory manager has been shut down.");
		}

		// since concurrent modifications to the collection
		// can disturb the release, we need to try potentially multiple times
		boolean successfullyReleased = false;
		do {
			final Iterator<MemorySegment> segmentsIterator = segments.iterator();

			Object lastOwner = null;
			Set<MemorySegment> segsForOwner = null;

			try {
				// go over all segments
				while (segmentsIterator.hasNext()) {

					final MemorySegment seg = segmentsIterator.next();
					if (seg == null || seg.isFreed()) {
						continue;
					}

					final Object owner = seg.getOwner();

					try {
						// get the list of segments by this owner only if it is a different owner than for
						// the previous one (or it is the first segment)
						if (lastOwner != owner) {
							lastOwner = owner;
							segsForOwner = this.allocatedSegments.get(owner);
						}

						// remove the segment from the list
						if (segsForOwner != null) {
							segsForOwner.remove(seg);
							if (segsForOwner.isEmpty()) {
								this.allocatedSegments.remove(owner);
							}
						}

						if (isPreAllocated) {
							memoryPool.returnSegmentToPool(seg);
						}
						else {
							seg.free();
							numNonAllocatedPages++;
						}
					}
					catch (Throwable t) {
						throw new RuntimeException(
								"Error removing book-keeping reference to allocated memory segment.", t);
					}
				}

				segments.clear();

				// the only way to exit the loop
				successfullyReleased = true;
			}
			catch (ConcurrentModificationException | NoSuchElementException e) {
				// this may happen in the case where an asynchronous
				// call releases the memory. fall through the loop and try again
			}
		} while (!successfullyReleased);
	}
	// -------------------- END CRITICAL SECTION -------------------
}
 

/**
 * Tries to release the memory for the specified segment. If the segment has already been released or
 * is null, the request is simply ignored.
 *
 * <p>If the memory manager manages pre-allocated memory, the memory segment goes back to the memory pool.
 * Otherwise, the segment is only freed and made eligible for reclamation by the GC.
 *
 * @param segment The segment to be released.
 * @throws IllegalArgumentException Thrown, if the given segment is of an incompatible type.
 */
public void release(MemorySegment segment) {
	// check if segment is null or has already been freed
	if (segment == null || segment.getOwner() == null) {
		return;
	}

	final Object owner = segment.getOwner();

	// -------------------- BEGIN CRITICAL SECTION -------------------
	synchronized (lock) {
		// prevent double return to this memory manager
		if (segment.isFreed()) {
			return;
		}
		if (isShutDown) {
			throw new IllegalStateException("Memory manager has been shut down.");
		}

		// remove the reference in the map for the owner
		try {
			Set<MemorySegment> segsForOwner = this.allocatedSegments.get(owner);

			if (segsForOwner != null) {
				segsForOwner.remove(segment);
				if (segsForOwner.isEmpty()) {
					this.allocatedSegments.remove(owner);
				}
			}

			if (isPreAllocated) {
				// release the memory in any case
				memoryPool.returnSegmentToPool(segment);
			}
			else {
				segment.free();
				numNonAllocatedPages++;
			}
		}
		catch (Throwable t) {
			throw new RuntimeException("Error removing book-keeping reference to allocated memory segment.", t);
		}
	}
	// -------------------- END CRITICAL SECTION -------------------
}
 

/**
 * Tries to release many memory segments together.
 *
 * <p>If the memory manager manages pre-allocated memory, the memory segment goes back to the memory pool.
 * Otherwise, the segment is only freed and made eligible for reclamation by the GC.
 *
 * @param segments The segments to be released.
 * @throws NullPointerException Thrown, if the given collection is null.
 * @throws IllegalArgumentException Thrown, id the segments are of an incompatible type.
 */
public void release(Collection<MemorySegment> segments) {
	if (segments == null) {
		return;
	}

	// -------------------- BEGIN CRITICAL SECTION -------------------
	synchronized (lock) {
		if (isShutDown) {
			throw new IllegalStateException("Memory manager has been shut down.");
		}

		// since concurrent modifications to the collection
		// can disturb the release, we need to try potentially multiple times
		boolean successfullyReleased = false;
		do {
			final Iterator<MemorySegment> segmentsIterator = segments.iterator();

			Object lastOwner = null;
			Set<MemorySegment> segsForOwner = null;

			try {
				// go over all segments
				while (segmentsIterator.hasNext()) {

					final MemorySegment seg = segmentsIterator.next();
					if (seg == null || seg.isFreed()) {
						continue;
					}

					final Object owner = seg.getOwner();

					try {
						// get the list of segments by this owner only if it is a different owner than for
						// the previous one (or it is the first segment)
						if (lastOwner != owner) {
							lastOwner = owner;
							segsForOwner = this.allocatedSegments.get(owner);
						}

						// remove the segment from the list
						if (segsForOwner != null) {
							segsForOwner.remove(seg);
							if (segsForOwner.isEmpty()) {
								this.allocatedSegments.remove(owner);
							}
						}

						if (isPreAllocated) {
							memoryPool.returnSegmentToPool(seg);
						}
						else {
							seg.free();
							numNonAllocatedPages++;
						}
					}
					catch (Throwable t) {
						throw new RuntimeException(
								"Error removing book-keeping reference to allocated memory segment.", t);
					}
				}

				segments.clear();

				// the only way to exit the loop
				successfullyReleased = true;
			}
			catch (ConcurrentModificationException | NoSuchElementException e) {
				// this may happen in the case where an asynchronous
				// call releases the memory. fall through the loop and try again
			}
		} while (!successfullyReleased);
	}
	// -------------------- END CRITICAL SECTION -------------------
}