Spark闭包清理类ClosureCleaner简析

Spark闭包清理类ClosureCleaner简析

从6月初开始因为一些工作上的事情，已经好久没有写博客了，这次把之前Spark源码阅读中深入了解的Spark闭包清理类ClosureCleaner简单介绍下，将知识留个档以便以后忘记了还有个地方来还原下思路。

Scala闭包机制回顾

在之前文章Spark闭包清理类ClosureCleaner简析中已经简单介绍了Scala的闭包实现方式，即用$outer字段来从闭包中引用外部的变量。

另外在另一篇文章慎用Scala中的return表达式中，介绍了在lambda表达式中，return语句的含义是NonLocalReturn而不是仅仅退出该lambda表达式，因此在lambda表达式中，return语句是很危险的，并且随时可能引起严重的后果。

Spark算子

这里并不会详细讲解Spark的RDD算子，我们仅仅从最简单的一个算子map入手，来看下你在rdd.map(func)中填入的func函数是如何运行在各个执行机之上的。

我们先从RDD类的map源码看起：

def map[U: ClassTag](f: T => U): RDD[U] = withScope {
    val cleanF = sc.clean(f)
    new MapPartitionsRDD[U, T](this, (context, pid, iter) => iter.map(cleanF))
}

def mapPartitions[U: ClassTag](f: Iterator[T] => Iterator[U], preservesPartitioning: Boolean = false): RDD[U] = withScope {
    val cleanedF = sc.clean(f)
    new MapPartitionsRDD(this, (context, index, iter) => cleanedF(iter), preservesPartitioning)
}

在这里多列出了下mapPartition的源码，原因是：当初我还是个Spark小白的时候，阅读各种博客，在文章中经常会看到一个spark效率优化点，就是使用mapPartition而不是map可以提高运行效率，然而从源码来看，两者的效率在大多数情况下应该是一致的。除非某些特殊情况：你自己写了Partitioner做分区，并且需要对整个分区的元素集合做某种操作，这种情况下必须用mapPartition。

我们回到map方法中来，它一共只有两句：

1、第一句对我们输出的func方法做了clean，我们下一节会详细介绍这个方法。
2、构造MapPartitionsRDD类来存放当前的RDD与计算函数。

从这里可以看出RDD的原理其实就是一个类似链表的存储结构，链表中存储的就是该链节点应该对上个节点的结果所做的操作函数。

在我们的func函数被存在了RDD里之后，它就会最终被Action算子触发，进入自己该在的Stage里，并被Spark序列化taskBinary = sc.broadcast(taskBinaryBytes)后存在Task类中下发给各个Executor节点运行。

这里有个重要的一点，Stage会被序列化成二进制，这说明其所有的成员都应该能被序列化，否则不可能在Executor中将其还原出来。因此，我们在写Spark代码的时候经常会碰到的org.apache.spark.SparkException: Task not serializable问题，就是因为我们在func闭包里引用了不可序列化的对象导致的。如果我们确实在闭包里使用了这些对象，那没话说，就是你的锅，自己改成可序列化得类去。但是还有些我们根本没用到的类，由于Scala的闭包机制，自动给我们引进来了，导致类的不可序列化，这种情况可就让我们为难了，举个栗子：
 

object TestObject {
  def run(): Int = {
    var nonSer = new NonSerializable
    val x = 5
    withSpark(new SparkContext("local", "test")) { sc =>
      val nums = sc.parallelize(Array(1, 2, 3, 4))
      nums.map(_ + x).reduce(_ + _)
    }
  }
}

以上我引用了Spark测试代码中的一个类来进行解释。我们传入map算子的闭包中，引用了外部的变量x，x为val整数，从之前的文章中分析，它会被直接存进闭包里，因此在闭包里不会存在$outer引用外部，因此并没啥问题。所以这是假栗子，仅仅帮助大家回顾下闭包而已。

再看个真正的栗子：
 

class TestClass extends Serializable {
  var nonSer = new NonSerializable
  var x = 5
  def getX: Int = x

  def run(): Int = {

    withSpark(new SparkContext("local", "test")) { sc =>
      val nums = sc.parallelize(Array(1, 2, 3, 4))
      nums.map(_ + getX).reduce(_ + _)
    }
  }
}

在传入map的闭包里，我们引用了外部TestClass类的getX方法，因此在闭包里会存在外部类的引用，这时如果我们直接对这个闭包做序列化，那由于闭包->$outer->nonSer间接包含了不可序列化对象，因此肯定会造成序列化失败异常。

Spark闭包清理类ClosureCleaner

为了解决以上问题，spark引入了闭包清理类，即在map的第一行调用了val cleanF = sc.clean(f)方法来实现的。下面我们深入该方法来一探究竟。

闭包清理最终调用了ClosureCleaner.clean(f, checkSerializable)方法来实现对f函数的闭包清理。

其内部代码涉及到的点很多，本文就不详细进行介绍了，以下介绍几个关键的技术点供大家来理解：

1、前一节提到return表达式是非常危险的，因此代码中对return做了判断，一旦发现return出现就会抛出异常快速失败：
 

if (op == NEW && tp.contains("scala/runtime/NonLocalReturnControl")) {
            throw new ReturnStatementInClosureException
          }

2、既然Scala的闭包是使用成员变量$outer来引用的，并且一个成员变量是否被用到必然要从字节码层面来进行判断，因此Spark闭包清理类中大量使用了org.apache.xbean.asm5包来遍历所有用到的方法的字节码，通过GETFIELD、INVOKEVIRTUAL、INVOKESPECIALJVM操作码来找到所引用的对象、方法、内部使用的类。

3、以传入clean的方法func为基准，向内查找找出所有内部使用到的对象所对应的成员变量和方法。

4、向外$outer查找所有外部所引用到的类，并将所有这些类克隆一份，但保持所有成员变量为初始值。

5、使用在内部查找到的所有引用变量值填充外部的克隆类成员变量。并将外部的$outer引用关系按照原始形势链接起来。

6、最终，使用反射将func的$outer成员换成对应的克隆后的克隆$outer对象

经过以上操作，就完成了对外部闭包引用的清理。简而言之就是，克隆个新的，用到的就填上，没用到的对象就是null了，这样就避免了不可序列化对象影响到了Spark的任务序列化功能。

不过这个闭包清理类还并不是那么完善，如SPARK-22328中的以下对话就隐藏了一个如果外部引用对象存在父类时的问题，当时我对这种场景以及延伸场景进行了验证，确实无法正确处理。
还有就是用方法获取了一个外部类，但是只用到了类中的一个可序列化对象，而该类存在其他不可序列化对象，这种情况下，序列化也是失败的。因为1、确实使用到了这个对象，因此对象会被赋值。2、无法简单从字节码中识别出改引用对象，因此代码中也没对该场景做特殊处理。

cloud-fan on 25 Oct 2017 Contributor
Assume we have class A and B having the same parent class P. P has 2 fields a and b. The closure accessed A.a and B.b, so when we clone A object, we should only set field a, when we clone B object, we should only set field b. However here seems we set field a and b for A and B object, which is sub-optimal.

cloud-fan on 25 Oct 2017 Contributor
Seems this is also a issue for the outerClasses, maybe I missed something…

viirya on 25 Oct 2017 Contributor
Seems that is true. For a closure that only accessed A.a, we clone the whole A object which contains both a and b fields. This is the fact in existing ClosureCleaner.

viirya on 25 Oct 2017 Contributor
As this is not a regression, IIUC, will it block this change?