Flink内存管理源码解读之内存管理器

/** The default memory page size. Currently set to 32 KiBytes. */
public static final int DEFAULT_PAGE_SIZE = 32 * 1024;

/** The minimal memory page size. Currently set to 4 KiBytes. */
public static final int MIN_PAGE_SIZE = 4 * 1024;

/** The lock used on the shared structures. */
private final Object lock = new Object();

public MemoryManager(long memorySize, int numberOfSlots) {

public MemoryManager(long memorySize, int numberOfSlots, int pageSize, MemoryType memoryType, boolean preAllocateMemory) {

switch (memoryType) {
 case HEAP:
  this.memoryPool = new HeapMemoryPool(memToAllocate, pageSize);
  break;
 case OFF_HEAP:
  this.memoryPool = new HybridOffHeapMemoryPool(memToAllocate, pageSize);
  break;
 default:
  throw new IllegalArgumentException("unrecognized memory type: " + memoryType);
}

public HeapMemoryPool(int numInitialSegments, int segmentSize) {
 this.availableMemory = new ArrayDeque<byte[]>(numInitialSegments);
 this.segmentSize = segmentSize;
 
 for (int i = 0; i < numInitialSegments; i++) {
  this.availableMemory.add(new byte[segmentSize]);
 }
}

public HybridOffHeapMemoryPool(int numInitialSegments, int segmentSize) {
 this.availableMemory = new ArrayDeque<ByteBuffer>(numInitialSegments);
 this.segmentSize = segmentSize;

 for (int i = 0; i < numInitialSegments; i++) {
  this.availableMemory.add(ByteBuffer.allocateDirect(segmentSize));
 }
}

public static ByteBuffer allocateDirect(int capacity);

public void shutdown() {
 // -------------------- BEGIN CRITICAL SECTION -------------------
 synchronized (lock)
 {
  if (!isShutDown) {
   // mark as shutdown and release memory
   isShutDown = true;
   numNonAllocatedPages = 0;

   // go over all allocated segments and release them
   for (Set<MemorySegment> segments : allocatedSegments.values()) {
    for (MemorySegment seg : segments) {
     seg.free();
    }
   }
   
   memoryPool.clear();
  }
 }
 // -------------------- END CRITICAL SECTION -------------------
}

synchronized (lock)
{

 Set<MemorySegment> segmentsForOwner = allocatedSegments.get(owner);
 if (segmentsForOwner == null) {
  segmentsForOwner = new HashSet<MemorySegment>(numPages);
  allocatedSegments.put(owner, segmentsForOwner);
 }

 if (isPreAllocated) {
  for (int i = numPages; i > 0; i--) {
   MemorySegment segment = memoryPool.requestSegmentFromPool(owner);
   target.add(segment);
   segmentsForOwner.add(segment);
  }
 }
 else {
  for (int i = numPages; i > 0; i--) {
   MemorySegment segment = memoryPool.allocateNewSegment(owner);
   target.add(segment);
   segmentsForOwner.add(segment);
  }
  numNonAllocatedPages -= numPages;
 }
}

synchronized (lock)
  {
   // 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);
   }
  }

synchronized (lock)
{
 // 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 e) {
   // this may happen in the case where an asynchronous
   // call releases the memory. fall through the loop and try again
  }
 } while (!successfullyReleased);
}

public void releaseAll(Object owner) {

public void write(byte[] b, int off, int len) throws IOException {
 int remaining = this.segmentSize - this.positionInSegment;
 if (remaining >= len) {
  this.currentSegment.put(this.positionInSegment, b, off, len);
  this.positionInSegment += len;
 }
 else {
  if (remaining == 0) {
   advance();
   remaining = this.segmentSize - this.positionInSegment;
  }
  while (true) {
   int toPut = Math.min(remaining, len);
   this.currentSegment.put(this.positionInSegment, b, off, toPut);
   off += toPut;
   len -= toPut;
   
   if (len > 0) {
    this.positionInSegment = this.segmentSize;
    advance();
    remaining = this.segmentSize - this.positionInSegment; 
   }
   else {
    this.positionInSegment += toPut;
    break;
   }
  }
 }
}

remaining >= len

remaining == 0