Spark Shuffle内存分析

分布式系统里的Shuffle 阶段往往是非常复杂的，而且分支条件也多，我只能按着我关注的线去描述。肯定会有不少谬误之处，我会根据自己理解的深入，不断更新这篇文章。

前言

用Spark写程序，一个比较蛋疼的地方就是OOM，或者GC严重，导致程序响应缓慢，一般这种情况都会出现在Shuffle阶段。Shuffle 是一个很复杂的过程，任何一个环节都足够写一篇文章。所以这里，我尝试换个方式，从实用的角度出发，让读者有三方面的收获：

构建出Shuffle的一个概览图剖析哪些环节，哪些代码可能会让内存产生问题控制相关内存的参数

有时候，我们宁可程序慢点，也不要OOM，至少要抛弃来跑步起来，希望这篇文章能够让你达成这个目标。

同时我们会提及一些类名，这些类方便你自己想更深入了解时，可以方便的找到他们，自己去探个究竟。
Shuffle 概览

Spark 的Shuffle 分为 Write,Read 两阶段。我们预先建立三个概念：

Write 对应的是ShuffleMapTask,具体的写操作ExternalSorter来负责Read 阶段由ShuffleRDD里的HashShuffleReader来完成。如果拉来的数据如果过大，需要落地，则也由ExternalSorter来完成的所有Write 写完后，才会执行Read。 他们被分成了两个不同的Stage阶段。

也就是说，Shuffle Write ,Shuffle Read 两阶段都可能需要落磁盘，并且通过Disk Merge 来完成最后的Sort归并排序。

我们看下面的图片:

图片来源于该文章： https://github.com/JerryLead/SparkInternals/blob/master/markdown/4-shuffleDetails.md

实际上，这是一个概览图，实际的要比这复杂的多，spill过程中的临时文件也很多。图片的几个点注意下：

图中的Bucket 部分以及Shuffle File 部分，都是由 ExternalSorter 来完成的。Shuffle Read 部分， 读取数据由HashShuffleReader来完成，并且透过ExternalSorter来完成Disk 排序文件持有数按Core来进行计算,假设Core数为C,假设Reduce数为R, 那么一个Executor 的持有文件数为: C*R

Shuffle Write 内存消耗分析

Shuffle Write 的入口链路为：

org.apache.spark.scheduler.ShuffleMapTask—> org.apache.spark.shuffle.sort.SortShuffleWriter
—> org.apache.spark.util.collection.ExternalSorter

会产生内存瓶颈的其实就是 org.apache.spark.util.collection.ExternalSorter。我们看看这个复杂的ExternalSorter都有哪些地方在占用内存：

第一个地：

private var map = new PartitionedAppendOnlyMap[K, C]

我们知道，数据都是先写内存，内存不够了，才写磁盘。这里的map就是那个放数据的内存了。

这个PartitionedAppendOnlyMap内部维持了一个数组，是这样的：

private var data = new Array[AnyRef](2 * capacity)

也就是他消耗的并不是Storage的内存，所谓Storage内存，指的是由blockManager管理起来的内存。

PartitionedAppendOnlyMap 放不下，要落地，那么不能硬生生的写磁盘，所以需要个buffer,然后把buffer再一次性写入内存。这个buffer是由参数

spark.shuffle.file.buffer=32k

控制的。数据获取的过程中，序列化反序列化，也是需要空间的，所以Spark 对数量做了限制，通过如下参数控制：

spark.shuffle.spill.batchSize=10000

假设一个Executor的可使用的Core为 C个，那么对应需要的内存消耗为：

C * 32k + C * 10000个Record + C * PartitionedAppendOnlyMap

这么看来，写文件的buffer不是问题，而序列化的batchSize也不是问题，几万或者十几万个Record 而已。那C * PartitionedAppendOnlyMap 到底会有多大呢？我先给个结论:

C * PartitionedAppendOnlyMap < ExecutorHeapMemeory * 0.2 * 0.8

怎么得到上面的结论呢？核心店就是要判定PartitionedAppendOnlyMap 需要占用多少内存，而它到底能占用内存，则由触发写磁盘动作决定，因为一旦写磁盘，PartitionedAppendOnlyMap所占有的内存就会被释放。下面是判断是否写磁盘的逻辑代码：

estimatedSize = map.estimateSize() if (maybeSpill(map, estimatedSize)) {
map = new PartitionedAppendOnlyMap[K, C]
}

每放一条记录，就会做一次内存的检查，看PartitionedAppendOnlyMap 到底占用了多少内存。如果真是这样，假设检查一次内存1ms, 1kw 就不得了的时间了。所以肯定是不行的，所以 estimateSize其实是使用采样算法来做的。

第二个，我们也不希望mayBeSpill太耗时,所以 maybeSpill 方法里就搞了很多东西，减少耗时。我们看看都设置了哪些防线

首先会判定要不要执行内部逻辑：

elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold

每隔32次会进行一次检查，并且要当前PartitionedAppendOnlyMap currentMemory > myMemoryThreshold 才会进一步判定是不是要spill.

其中 myMemoryThreshold可通过如下配置获得初始值

spark.shuffle.spill.initialMemoryThreshold = 5 * 1024 * 1024

接着会向 shuffleMemoryManager 要 2 * currentMemory - myMemoryThreshold 的内存，shuffleMemoryManager 是被Executor 所有正在运行的Task(Core) 共享的，能够分配出去的内存是：

ExecutorHeapMemeory * 0.2 * 0.8

上面的数字可通过下面两个配置来更改：

spark.shuffle.memoryFraction=0.2spark.shuffle.safetyFraction=0.8

如果无法获取到足够的内存，就会出发真的spill操作了。

看到这里，上面的结论就显而易见了。

然而，这里我们忽略了一个很大的问题，就是

estimatedSize = map.estimateSize()

为什么说它是大问题，前面我们说了，estimateSize 是近似估计，所以有可能估的不准，也就是实际内存会远远超过预期。

具体的大家尅看看 org.apache.spark.util.collection.SizeTracker

我这里给出一个结论：

如果你内存开的比较大，其实反倒风险更高，因为estimateSize 并不是每次都去真实的算缓存是通过通过采样数据来完成的，而采样的周期不是固定的，而是指数增长的，比如第一次采样完后，PartitionedAppendOnlyMap 要经过1.1次之后才进行第二次采样，然后经过1.1*.1.1次之后进行第三次采样，以此递推，假设你内存开的大，那PartitionedAppendOnlyMap可能要经过几十万次更新之后之后才会进行一次采样，然后才能计算出新的大小，这个时候几十万次更新带来的新的内存压力，可能已经让你的GC不抗重负了。

当然，这是一种折中，因为确实不能频繁采样。

如果你不想出现这种问题，要么自己替换实现这个类，要么将

spark.shuffle.safetyFraction=0.8

设置的更小一些。
Shuffle Write 内存消耗分析

Shuffle Write 的入口链路为：

org.apache.spark.rdd.ShuffledRDD—> org.apache.spark.shuffle.sort.HashShuffleReader
—> org.apache.spark.util.collection.ExternalSorter

这里内存占用点和ShuffleWrite 其实是差不多的。
另外就是需要考虑下拉数据的问题，就是去读取Shuffle Write的数据。目前提供了 NIO 和 Netty的方式。这里理论上不会出现太大内存占用，除非比如触发了Netty的Bug等。