PySpark源码分析之Driver端基于Py4j的通信详解

概述

接上文 PySpark源码分析之AM端运行流程（Driver） 的最后部分可知，PySpark是通过Py4j来实现与Scala端JVM通信交互的（注：Py4j只用于Driver端Python调用JVM中方法；Executor端是直接通过socket通信的。），PySpark主要提供一层Python API的封装。其过程主要是JVM会开启一个Socket端口提供RPC服务，Python需要调用Spark API时，它会作为客户端将调用指令序列化成字节流发送到Socket服务端口，JVM接受字节流后解包成对应的指令，然后找到目标对象和代码进行执行，然后将执行结果序列化成字节流通过Socket返回给客户端，客户端收到字节流后再解码成Python对象，于是Python客户端就成功拿到了远程调用的结果（注：Py4j同时支持【Python -> JVM】和【JVM -> Python】，本文仅仅分析前者，后者在Spark Streaming中有适用场景，以后有时间再分析）。

官方文档：https://www.py4j.org/index.html
github：https://github.com/bartdag/py4j

Py4j内存模型：

Every time a Java object is sent to the Python side, a reference to the object is kept on the Java side (in the Gateway class). Once the object is garbage collected on the Python VM (reference count == 0), the reference is removed on the Java VM: if this was the last reference, the object will likely be garbage collected too. When a gateway is shut down, the remaining references are also removed on the Java VM.
Because Java objects on the Python side are involved in a circular reference (JavaObject and JavaMember reference each other), these objects are not immediately garbage collected once the last reference to the object is removed (but they are guaranteed to be eventually collected if the Python garbage collector runs before the Python program exits).
In doubt, users can always call the detach function on the Python gateway to explicitly delete a reference on the Java side. A call to gc.collect() also usually works.
每次将Java对象发送到Python端时，对该对象的引用都保存在Java端(在Gateway类中)。一旦对象在Python VM上被垃圾收集(引用计数== 0)，该引用在Java的VM上也会被删除：如果这是最后一个引用，那么该对象也很可能被垃圾收集。当gateway关闭时，剩余的引用也会在Java VM上删除。
因为Python端的Java对象间循环引用（JavaObject和JavaMember相互引用），所以在删除最后一个对象引用后并不会立即进行垃圾回收（但能够保证在Python程序最终退出之前，进行Python的垃圾回收)。
毫无疑问，用户也可通过在Python端调用detach函数的来显式的删除Java端上的引用。也可调用gc.collect()实现。
参考：https://www.py4j.org/advanced_topics.html#py4j-memory-model

服务端启动

首先我们先看Py4j的原理图，如下图：

如上图，Spark的Driver运行PythonRunner启动Py4j服务端，当服务启动后GatewayServer开始循环监听端口接受socket连接，针对每个socket连接开启一个子线程，创建对应的GatewayConnect，并根据客户端传递的消息指令选择命令类进行执行，处理后的结果同样通过封装socket返回给Python客户端。
命令类包含：

• ArrayCommand：处理Array的相关操作。
• CallCommand：在已实例化类上反射调用函数。
• ConstructorCommand：通过反射实例化类。
• FieldCommand：通过反射设置或获取实例化类变量值。
• HelpPageCommand：为Java对象生成帮助文档。
• ListCommand：处理List的相关操作。
• MemoryCommand：用于垃圾回收过程中删除服务端对象。
• ReflectionCommand：用于反射package、class和静态成员。
• ShutdownGatewayServerCommand：接受处理关闭GatewayServer的指令。
• JVMViewCommand：管理JVMView的操作，包含：创建新的JVMView（提供了默认的defaultJVMView，但也支持创建新的）、添加imports和搜索拼接类完全限定名。
• ExceptionCommand：处理异常类的操作。
• DirCommand：用于获取类的相关属性，包含：public字段名、public方法名、public static属性（包含：字段名、方法名和类名）和JVMView中的import
• StreamCommand：同CallCommand，区别是CallCommand直接返回结果，StreamCommand以流的方式返回结果。

然后我们看核心类 Gateway （分析见注释，只截取部分关键代码）：

/**
 *
 * 

 * A Gateway manages various states: entryPoint, references to objects returned
 * to a Python program, etc.
 * 
 *
 * 

 * This class is not intended to be directly accessed by users.
 * 
 *
 * @author Barthelemy Dagenais
 *
 */
public class Gateway {

    /** Server创建的任意对象都会保存在bindings变量 */
    private final Map bindings = new ConcurrentHashMap();
    private final AtomicInteger objCounter = new AtomicInteger();
    private final AtomicInteger argCounter = new AtomicInteger();
    private final static String OBJECT_NAME_PREFIX = "o";
    private final Object entryPoint;
    /** 反射工具类 */
    private final ReflectionEngine rEngine = new ReflectionEngine();
    private Py4JPythonClient cbClient;
    /** 默认JVM视图，主要保存import package */
    private final JVMView defaultJVMView;

    private final Logger logger = Logger.getLogger(Gateway.class.getName());

    private boolean isStarted = false;

    public Gateway(Object entryPoint) {
        this(entryPoint, null);
    }

    public Gateway(Object entryPoint, Py4JPythonClient cbClient) {
        this.entryPoint = entryPoint;
        this.cbClient = cbClient;
        this.defaultJVMView = new JVMView("default", Protocol.DEFAULT_JVM_OBJECT_ID);
    }
    public void deleteObject(String objectId) {
        bindings.remove(objectId);
    }

    public JVMView getDefaultJVMView() {
        return this.defaultJVMView;
    }
    protected String getNextObjectId() {
        return OBJECT_NAME_PREFIX + objCounter.getAndIncrement();
    }
    ...

    /**
     * 所有Server创建的对象返回给client前都会调用该方法，主要功能包含：
     * 1、新创建的对象都通过putNewObject函数生成objectId，并把对象添加到bindings变量（Map类型），注：基础数据类型、BigDecimal和无返回的对象无需添加到bindings变量
     * 2、根据创建的对象类型组织返回给Client的对象ReturnObject，注：一般最终调用Protocol.getOutputCommand(returnObject)拼接成返回的字符串
     */
    @SuppressWarnings("rawtypes")
    public ReturnObject getReturnObject(Object object) {
        ReturnObject returnObject;
        if (object != null) {
            if (isPrimitiveObject(object)) {
                returnObject = ReturnObject.getPrimitiveReturnObject(object);
            } else if (object == ReflectionEngine.RETURN_VOID) {
                returnObject = ReturnObject.getVoidReturnObject();
            } else if (isDecimalObject(object)) {
                returnObject = ReturnObject.getDecimalReturnObject(object);
            } else if (isList(object)) {
                String objectId = putNewObject(object);
                returnObject = ReturnObject.getListReturnObject(objectId, ((List) object).size());
            } else if (isMap(object)) {
                String objectId = putNewObject(object);
                returnObject = ReturnObject.getMapReturnObject(objectId, ((Map) object).size());
            } else if (isArray(object)) {
                String objectId = putNewObject(object);
                returnObject = ReturnObject.getArrayReturnObject(objectId, Array.getLength(object));
            } else if (isSet(object)) {
                String objectId = putNewObject(object);
                returnObject = ReturnObject.getSetReturnObject(objectId, ((Set) object).size());
            } else if (isIterator(object)) {
                String objectId = putNewObject(object);
                returnObject = ReturnObject.getIteratorReturnObject(objectId);
            } else {
                String objectId = putNewObject(object);
                returnObject = ReturnObject.getReferenceReturnObject(objectId);
            }
        } else {
            returnObject = ReturnObject.getNullReturnObject();
        }
        return returnObject;
    }

    /**
     * 

     * Invokes a constructor and returned the constructed object.
     * 根据类的全限定名（例如：package1.package2.className）和构造参数反射实例化类
     * 
     *
     * @param fqn
     *            The fully qualified name of the class.
     * @param args
     * @return
     */
    public ReturnObject invoke(String fqn, List