Flink 内存管理

TaskManager的内存布局

Flink内部并非直接将对象存储在堆上，而是将对象序列化到一个个预先分配的MemorySegment中。MemorySegment是一段固定长度的内存(默认32KB)；也是Flink中最小的内存分配单元。MemorySegment提供了高效的读写方法，它的底层可以是堆上的byte[]，也可以是堆外(off-heap)ByteBuffer。可以把MemorySegment看作Java NIO中的ByteBuffer，Flink还实现了Java的java.io.DataOutput和java.io.DataInput接口，分别是AbstractPagedInputView和AbstractPagedOutputView，其可以通过一种逻辑视图的方式来操作连续的多块MemorySegment。

在Flink中，TaskManager负责任务的实际运行，通常一个TaskManager对应一个JVM进程(非MiniCluster模式)。抛开JVM内存模型，单从TaskManager内存的主要使用方式来看，TaskManager的内存主要分为三个部分：

Network Buffers：一定数量的MemorySegment，主要用于网络传输。在TaskManager启动时分配，通过NetworkEnvironment和NetworkBufferPool进行管理
Managed Memory：由MemoryManager管理的一组MemorySegment集合，主要用于Batch模式下的sorting，hashing和cache等。
Remaining JVM heap：余下的堆内存留给TaskManager的数据结构以及用户代码处理数据时使用。TaskManager自身的数据结构并不会占用太多内存，因而主要都是供用户代码使用，用户代码创建的对象通常生命周期都较短

需要注意的是，上面所说的三部分的内存并非都是JVM堆上的内存，因为MemorySegment底层的内存可以在堆上，也可以在堆外(不由JVM管理)。对于Network Buffers，这一部分内存就是在堆外(off-heap)进行分配的；对于Managed Memory，这一部分内存可以配置在堆上，也可以配置在堆外。另外还需要注意的一点是，Managed Memory主要是在Batch模式下使用，在Streaming模式下这一部分内存并不会预分配，因而空闲出来的内存其实都是可以给用户自定义函数使用的。

通过二进制数据管理对象

Flink是通过MemorySegment来管理数据对象的，因而对象首先需要被序列化保存到MemorySegment中。在Java的生态系统中，已经存在很多现有的序列化框架了，如Java自带的序列化机制、Kryo、Avro、Thrift、Protobuf等，但Flink也实现了一套自己的序列化框架。这主要是出于以下考虑：首先，比较和操作二进制数据需要准确了解序列化的布局，针对二进制数据的操作来配置序列化的布局可以显著提升性能；其次，对于Flink应用而言，它所处理的数据对象类型通常是完全已知的，由于数据集对象的类型固定，对于数据集可以只保存一份对象Schema信息，可以进一步节省存储空间。

Flink可以处理任意的Java或Scala对象，而不必实现特定的接口。对于Java实现的Flink程序，Flink会通过反射框架获取用户自定义函数返回的类型；而对于Scala实现的Flink程序，则通过Scala Compiler分析用户自定义函数返回的类型。每一种数据类型都对应一个TypeInfomation。其基本的数据类型如下：

Flink在其内部构建了一套自己的类型系统，Flink现阶段支持的类型分类如图所示，从图中可以看到Flink类型可以分为基础类型（Basic）、数组（Arrays）、复合类型（Composite）、辅助类型（Auxiliary）、泛型和其它类型（Generic）。Flink支持任意的Java或是Scala类型。不需要像Hadoop一样去实现一个特定的接口（org.apache.hadoop.io.Writable），Flink能够自动识别数据类型。

通过TypeInfomation可以获取到对应数据类型的序列化器TypeSerializer。对于BasicTypeInfo，Flink提供了对应的序列化器；对于WritableTypeInfo，Flink会将序列化和反序列化操作委托给Hadoop Writable接口的write() and readFields()；对于GenericTypeInfo，Flink默认使用Kyro进行序列化；而TupleTypeInfo、CaseClassTypeInfo和PojoTypeInfo是一种组合类型，序列化时分别委托给成员的序列化器进行序列化即可。

Flink自带了很多TypeSerializer子类，大多数情况下各种自定义类型都是常用类型的排列组合，因而可以直接复用，如果内建的数据类型和序列化方式不能满足你的需求，Flink的类型信息系统也支持用户拓展。若用户有一些特殊的需求，只需要实现TypeInformation、TypeSerializer和TypeComparator即可定制自己类型的序列化和比较大小方式，来提升数据类型在序列化和比较时的性能。基本的序列化示例如下：

如图所示，Tuple3这个对象的序列化过程如下：Tuple3包含三个层面，一是int类型，一是double类型，还有一个是Person。Person包含两个字段，一是int型的ID，另一个是String类型的name，它在序列化操作时，会委托相应具体序列化的序列化器进行相应的序列化操作。从图中可以看到Tuple3会把int类型通过IntSerializer进行序列化操作，此时int只需要占用四个字节就可以了。根据int占用四个字节，这个能够体现出Flink可序列化过程中的一个优势，即在知道数据类型的前提下，可以更好的进行相应的序列化与反序列化操作。相反，如果采用Java的序列化，虽然能够存储更多的属性信息，但一次占据的存储空间会受到一定的损耗。

对于可以用作key的数据类型，TypeInfomation还可以生成TypeComparator，用来直接在序列化后的二进制数据上进行compare、hash等操作。

在批处理的场景下，诸如group、sort和join等操作都需要访问大量的数据。借助于MemorySegment并直接操作二进制数据，Flink可以高效地完成这些操作，避免了频繁地序列化/反序列化，并且这些操作是缓存友好的。这种基于MemorySegment和二进制数据直接管理数据对象的方式可以带来如下好处：

保证内存安全：由于分配的MemorySegment的数量是固定的，因而可以准确地追踪MemorySegment的使用情况。在Batch模式下，如果MemorySegment资源不足，会将一批MemorySegment写入磁盘，需要时再重新读取。这样有效地减少了OOM的情况。
减少了GC的压力：因为分配的MemorySegment是长生命周期的对象，数据都以二进制形式存放，且MemorySegment可以回收重用，所以MemorySegment会一直保留在老年代不会被GC；而由用户代码生成的对象基本都是短生命周期的，MinorGC可以快速回收这部分对象，尽可能减少MajorGC的频率。此外，MemorySegment还可以配置为使用堆外内存，进而避免GC。
节省内存空间：数据对象序列化后以二进制形式保存在MemorySegment中，减少了对象存储的开销。
高效的二进制操作和缓存友好的计算：可以直接基于二进制数据进行比较等操作，避免了反复进行序列化于反序列；另外，二进制形式可以把相关的值，以及hash值，键值和指针等相邻地放进内存中，这使得数据结构可以对高速缓存更友好。

MemorySegment

MemorySegment是一段固定长度的内存，也是Flink中最小的内存分配单元。在早期版本的实现中，MemorySegment使用的都是堆上的内存。尽管Flink的内存管理机制已经做了很多优化，但是Flink团队仍然加入了对堆外内存的支持。主要是考虑到以下几个方面：

启动很大堆内存（100s of GBytes heap memory）的JVM需要很长时间，GC停留时间也会很长（分钟级）。使用堆外内存的话，JVM只需要分配较少的堆内存（只需要分配RemainingHeap那一块）。
堆外内存在写磁盘或网络传输时是可以利用zero-copy特性，I/O和网络传输的效率更高。
堆外内存是进程间共享的，也就是说，即使JVM进程崩溃也不会丢失数据。这可以用来做故障恢复。Flink暂时没有利用起这个，不过未来有可能会利用这个特性。

但是使用堆外内存同样存在一些潜在的问题：

堆内存可以很方便地进行监控和分析，相较而言堆外内存则更加难以控制；
Flink有时可能需要短生命周期的MemorySegment，在堆上申请开销会更小；
一些操作在堆内存上会更快一些

Flink将原来的MemorySegment变成了抽象类，并提供了两个具体的子类：HeapMemorySegment和HybridMemorySegment。前者是用于分配堆内存，后者用来分配堆外内存和堆内存的。

在早期版本中，由于MemorySegment是只基于堆内存的，因而只需要提供一种类型的MemorySegment实现即可；而在引入对堆外内存的支持后，按一般的思路是应该在新增一个基于堆外内存的实现即可。但是，这里涉及到一个JIT优化的性能问题。在只有一种类型的MemorySegment的情况下，通过ClassHierarchyAnalysis(CHA)，JIT编译器能够确定方法调用的具体实现，因而方法调用可以通过去虚化（de-virtualized）和内联（inlined）来提升性能。而一旦有了两种类型的实现，在同时使用两种类型的MemorySegment的情况下，JIT编译器就无法进行优化，这大概会导致2.7倍的性能差异。因而Flink做了这两种优化：1）确保只有一种MemorySegment的实现被加载；2）提供一种能同时处理管理堆内存和堆外内存的MemorySegment实现，从而保证频繁调用的MemorySegment能够被JIT优化。其MemorySegment的源码实现如下：

public abstract class MemorySegment {protected final byte[] heapMemory; // 堆内存引用protected long address; // 堆外内存地址// 基于堆内存创建MemorySegmentMemorySegment(byte[] buffer, Object owner) {if (buffer == null) {throw new NullPointerException("buffer");}this.heapMemory = buffer;this.address = BYTE_ARRAY_BASE_OFFSET;this.size = buffer.length;this.addressLimit = this.address + this.size;this.owner = owner;}// 基于堆外内存创建MemorySegmentMemorySegment(long offHeapAddress, int size, Object owner) {if (offHeapAddress <= 0) {throw new IllegalArgumentException("negative pointer or size");}if (offHeapAddress >= Long.MAX_VALUE - Integer.MAX_VALUE) {// this is necessary to make sure the collapsed checks are safe against numeric overflowsthrow new IllegalArgumentException("Segment initialized with too large address: " + offHeapAddress+ " ; Max allowed address is " + (Long.MAX_VALUE - Integer.MAX_VALUE - 1));}this.heapMemory = null;this.address = offHeapAddress;this.addressLimit = this.address + size;this.size = size;this.owner = owner;}public boolean isOffHeap() {return heapMemory == null;}public final long getLong(int index) {final long pos = address + index;if (index >= 0 && pos <= addressLimit - 8) {return UNSAFE.getLong(heapMemory, pos); // 这是能够在一个实现中同时操作堆内存和堆外内存的关键}else if (address > addressLimit) {throw new IllegalStateException("segment has been freed");}else {// index is in fact invalidthrow new IndexOutOfBoundsException();}}//.........
}public final class HybridMemorySegment extends MemorySegment {private final ByteBuffer offHeapBuffer; // 堆外内存private final Runnable cleaner; // The cleaner is called to free the underlying native memory.//堆外内存初始化HybridMemorySegment(@Nonnull ByteBuffer buffer, @Nullable Object owner, @Nullable Runnable cleaner) {super(checkBufferAndGetAddress(buffer), buffer.capacity(), owner);this.offHeapBuffer = buffer;this.cleaner = cleaner;}//堆内内存初始化HybridMemorySegment(byte[] buffer, Object owner) {super(buffer, owner);this.offHeapBuffer = null;this.cleaner = null;}//.........
}public final class HeapMemorySegment extends MemorySegment {private byte[] memory;HeapMemorySegment(byte[] memory, Object owner) {super(Objects.requireNonNull(memory), owner);this.memory = memory;}//......
}

之所以能够使用同一份代码实现既能够处理堆内存又能够处理堆外内存的效果，其关键点在于sun.misc.Unsafe的一些方法会根据对象引用表现出不同的行为，列如sun.misc.Unsafe.getLong(Object reference, long offset)；在reference不为null的情况下，则会取该对象的地址，加上后面的offset，从相对地址处取出8字节；而在reference为null的情况下，则offset就是要操作的绝对地址。所以，通过控制对象引用的值，就可以灵活地管理堆外内存和堆内存。

既然HybridMemorySegment可以同时管理堆内存和堆外内存，为什么还需要HeapMemorySegment呢？这是因为假如所有的MemorySegment都是在堆上分配的，使用HeapMemorySegment相比于HybridMemorySegment会有更好的性能。但实际上，由于Flink中Network buffer使用的MemorySegment一定是在堆外分配的，HeapMemorySegment在Flink中已经不会再使用了，具体可以参考FLINK-7310 always use the HybridMemorySegment。MemorySegment通常不直接构造，而是通过MemorySegmentFactory来创建如下：

public final class MemorySegmentFactory { // A factory for (hybrid) memory segments ({@link HybridMemorySegment}).
// 申请分配堆内存public static MemorySegment wrap(byte[] buffer) { // 创建堆内存 heap memory regionreturn new HybridMemorySegment(buffer, null);}public static MemorySegment allocateUnpooledSegment(int size) { // 按照指定大小 分配堆内存 new byte[size]  实现如下：return allocateUnpooledSegment(size, null);}public static MemorySegment allocateUnpooledSegment(int size, Object owner) {return new HybridMemorySegment(new byte[size], owner);}
// 申请分配堆外内存public static MemorySegment allocateUnpooledOffHeapMemory(int size) { //  按照指定大小 分配堆外内存ByteBuffer.allocateDirect(size)  实现如下：return allocateUnpooledOffHeapMemory(size, null);}public static MemorySegment allocateUnpooledOffHeapMemory(int size, Object owner) {ByteBuffer memory = ByteBuffer.allocateDirect(size);return new HybridMemorySegment(memory, owner, null);}public static MemorySegment allocateOffHeapUnsafeMemory(int size, Object owner) {long address = MemoryUtils.allocateUnsafe(size);ByteBuffer offHeapBuffer = MemoryUtils.wrapUnsafeMemoryWithByteBuffer(address, size);return new HybridMemorySegment(offHeapBuffer, owner, MemoryUtils.createMemoryGcCleaner(offHeapBuffer, address));}public static MemorySegment wrapOffHeapMemory(ByteBuffer memory) {return new HybridMemorySegment(memory, null, null);}
}

MemorySegment的管理

在TaskManager的内存布局中我们说过，TaskManager的内存主要分为三个部分：其中Network Buffers和Managed Memory都是一组MemorySegment的集合；下面就分别介绍下这两块内存是如何管理的。

Buffer和Network Buffer Pool

Buffer接口是对池化的MemorySegment的包装，带有引用计数，类似于Netty的ByteBuf。Buffer也使用两个指针分别表示写入的位置和读取的位置。Buffer的具体实现实现类NetworkBuffer继承自Netty的AbstractReferenceCountedByteBuf，这使得它很容易地集成了引用计数和读写指针的功能。同时，在非Netty场景下使用时，Buffer也提供了java.nio.ByteBuffer的包装，但需要手动设置读写指针的位置。ReadOnlySlicedNetworkBuffer则提供了只读模式的buffer的包装。

BufferBuilder和BufferConsumer构成了写入和消费buffer的通用模式：通过BufferBuilder向底层的MemorySegment写入数据，再通过BufferConsumer生成只读的Buffer，读取BufferBuilder写入的数据。这两个类都不是线程安全的，但可以实现一个线程写入，另一个线程读取的效果。

BufferPool接口继承了BufferProvider和BufferRecycler接口，提供了申请以及回收Buffer的功能。LocalBufferPool是BufferPool的具体实现，LocalBufferPool中Buffer的数量是可以动态调整的。

BufferPoolFactory接口是BufferPool的工厂，用于创建及销毁BufferPool。NetworkBufferPool是BufferPoolFactory的具体实现类。所以按照BufferPoolFactory->BufferPool->Buffer这样的结构进行组织。NetworkBufferPool在初始化的时候创建一组MemorySegment，这些MemorySegment会在所有的LocalBufferPool之间进行均匀分配。

public class NetworkBufferPool implements BufferPoolFactory, MemorySegmentProvider, AvailabilityProvider {private final int totalNumberOfMemorySegments;private final int memorySegmentSize;// 所有可用的MemorySegment，阻塞队列private final ArrayBlockingQueue<MemorySegment> availableMemorySegments;// ---- Managed buffer pools ----------------------------------------------private final Object factoryLock = new Object();private final Set<LocalBufferPool> allBufferPools = new HashSet<>();private int numTotalRequiredBuffers;private final int numberOfSegmentsToRequest;private final Duration requestSegmentsTimeout;/*** Allocates all {@link MemorySegment} instances managed by this pool.*/public NetworkBufferPool(int numberOfSegmentsToAllocate,int segmentSize,int numberOfSegmentsToRequest,Duration requestSegmentsTimeout) {this.totalNumberOfMemorySegments = numberOfSegmentsToAllocate;this.memorySegmentSize = segmentSize;checkArgument(numberOfSegmentsToRequest > 0, "The number of required buffers should be larger than 0.");this.numberOfSegmentsToRequest = numberOfSegmentsToRequest;Preconditions.checkNotNull(requestSegmentsTimeout);checkArgument(requestSegmentsTimeout.toMillis() > 0, "The timeout for requesting exclusive buffers should be positive.");this.requestSegmentsTimeout = requestSegmentsTimeout;final long sizeInLong = (long) segmentSize;try {this.availableMemorySegments = new ArrayDeque<>(numberOfSegmentsToAllocate); // 所有可用的MemorySegment，阻塞队列} catch (OutOfMemoryError err) {throw new OutOfMemoryError("Could not allocate buffer queue of length " + numberOfSegmentsToAllocate + " - " + err.getMessage());}try {for (int i = 0; i < numberOfSegmentsToAllocate; i++) { // NetworkBufferPool使用的MemorySegment全是堆外内存availableMemorySegments.add(MemorySegmentFactory.allocateUnpooledOffHeapMemory(segmentSize, null));}}catch (OutOfMemoryError err) {// ............}
}

MemoryManager

MemoryManager是管理Managed Memory的类，这部分主要是在Batch模式下使用，在Streaming模式下这一块内存不会分配。MemoryManager主要通过内部接口MemoryPool来管理所有的MemorySegment。Managed Memory和管理相比于Network Buffers的管理更为简单，因为不需要Buffer的那一层封装。其源代码如下：

public class MemoryManager {// 管理所有的MemorySegmentprivate final MemoryPool memoryPool; // The memory pool from which we draw memory segments. Specific to on-heap or off-heap memoryprivate final HashMap<Object, Set<MemorySegment>> allocatedSegments; // Memory segments allocated per memory owner.public MemoryManager(long memorySize, int numberOfSlots, int pageSize,MemoryType memoryType, boolean preAllocateMemory) {// sanity checks......this.memoryType = memoryType;this.memorySize = memorySize;this.numberOfSlots = numberOfSlots;// assign page size and bit utilitiesthis.pageSize = pageSize;this.roundingMask = ~((long) (pageSize - 1));final long numPagesLong = memorySize / pageSize;if (numPagesLong > Integer.MAX_VALUE) {throw new IllegalArgumentException("The given number of memory bytes (" + memorySize+ ") corresponds to more than MAX_INT pages.");}this.totalNumPages = (int) numPagesLong; // 所有可用的MemorySegment数量if (this.totalNumPages < 1) {throw new IllegalArgumentException("The given amount of memory amounted to less than one page.");}this.allocatedSegments = new HashMap<Object, Set<MemorySegment>>();this.isPreAllocated = preAllocateMemory;this.numNonAllocatedPages = preAllocateMemory ? 0 : this.totalNumPages;final int memToAllocate = preAllocateMemory ? this.totalNumPages : 0; // 是否需要预分配内存，Streaming不会预分配switch (memoryType) {case HEAP: // 堆上内存this.memoryPool = new HybridHeapMemoryPool(memToAllocate, pageSize);break;case OFF_HEAP: // 堆外内存if (!preAllocateMemory) {LOG.warn("It is advisable to set 'taskmanager.memory.preallocate' to true when" +" the memory type 'taskmanager.memory.off-heap' is set to true.");}this.memoryPool = new HybridOffHeapMemoryPool(memToAllocate, pageSize);break;default:throw new IllegalArgumentException("unrecognized memory type: " + memoryType);}// ......}abstract static class MemoryPool {abstract int getNumberOfAvailableMemorySegments();abstract MemorySegment allocateNewSegment(Object owner);abstract MemorySegment requestSegmentFromPool(Object owner);abstract void returnSegmentToPool(MemorySegment segment);abstract void clear();}static final class HybridHeapMemoryPool extends MemoryPool {private final ArrayDeque<byte[]> availableMemory; // The collection of available memory segments. private final int segmentSize;HybridHeapMemoryPool(int numInitialSegments, int segmentSize) {this.availableMemory = new ArrayDeque<>(numInitialSegments);this.segmentSize = segmentSize;for (int i = 0; i < numInitialSegments; i++) {this.availableMemory.add(new byte[segmentSize]); // 堆上直接使用byte数组}}}static final class HybridOffHeapMemoryPool extends MemoryPool {private final ArrayDeque<ByteBuffer> availableMemory; // The collection of available memory segments. private final int segmentSize;HybridOffHeapMemoryPool(int numInitialSegments, int segmentSize) {this.availableMemory = new ArrayDeque<>(numInitialSegments);this.segmentSize = segmentSize;for (int i = 0; i < numInitialSegments; i++) {this.availableMemory.add(ByteBuffer.allocateDirect(segmentSize)); // 堆外使用DirectByteBuffer}}}
}