【Spark内核源码】内置的RPC框架，Spark的通信兵（二）

RPC管道处理TransportChannelHandler

RPC服务端处理RpcHandler

引导程序Bootstrap

RPC客户端TransportClient

总结

接着【Spark内核源码】内置的RPC框架，Spark的通信兵（一）接着分析

RPC管道处理TransportChannelHandler

TransportContext最后一个作用就是使用org.apache.spark.network.server.TransportChannelHandler设置Netty Channel pipelines（Netty的通信管道）。

在TransportClientFactory的createClient方法和TransportServer的init方法中都执行了初始化管道方法，也就是TransportContext中的initializePipeline方法。

TransportClientFactory中调用initializePipeline方法

initializePipeline方法的代码如下：

public TransportChannelHandler initializePipeline(
      SocketChannel channel,
      RpcHandler channelRpcHandler) {
    try {
      TransportChannelHandler channelHandler = createChannelHandler(channel, channelRpcHandler);
      /**
       * 管道设置
       * Request时按照顺序执行，TransportFrameDecoder-》MessageDecoder-》IdleStateHandler-》TransportChannelHandler
       * Response时按照逆序执行，IdleStateHandler-》MessageEncoder
       * */
      channel.pipeline()
        .addLast("encoder", encoder) //为pipeline设置encoder
        .addLast(TransportFrameDecoder.HANDLER_NAME, NettyUtils.createFrameDecoder()) //为pipeline设置frameDecoder
        .addLast("decoder", decoder) //为pipeline设置decoder
        .addLast("idleStateHandler", //为pipeline设置IdleStateHandler，Netty内置对象
                new IdleStateHandler(0, 0, conf.connectionTimeoutMs() / 1000))
        // NOTE: Chunks are currently guaranteed to be returned in the order of request, but this
        // would require more logic to guarantee if this were not part of the same event loop.
        .addLast("handler", channelHandler);
      return channelHandler;
    } catch (RuntimeException e) {
      logger.error("Error while initializing Netty pipeline", e);
      throw e;
    }
  }

首先是调用createChannelHandler方法创建TransportChannelHandler对象，createChannelHandler代码如下：

/**
   * TransportChannelHandler
   * 在服务端代理TransportRequestHandler处理请求消息
   * 在客户端代理TransportResponseHandler处理相应信息
   *
   * */
  private TransportChannelHandler createChannelHandler(Channel channel, RpcHandler rpcHandler) {
    // 创建TransportChannelHandler的同时，创建了TransportResponseHandler、TransportRequestHandler和TransportClient
    TransportResponseHandler responseHandler = new TransportResponseHandler(channel);
    // 真正意义上的创建TransportClient
    TransportClient client = new TransportClient(channel, responseHandler);
    TransportRequestHandler requestHandler = new TransportRequestHandler(channel, client,
      rpcHandler);
    return new TransportChannelHandler(client, responseHandler, requestHandler,
      conf.connectionTimeoutMs(), closeIdleConnections);
  }

创建TransportChannelHandler之前，先创建了TransportResponseHandler、TransportClient和TransportRequestHandler，在这里才是真正意义的创建TransportClient对象，与管道一一对应，保证所有用户使用channel时得到的是同一个TrasportClient对象。TransportChannelHandler在服务端代理TransportRequestHandler处理请求消息，在客户端代理TransportResponseHandler处理相应信息。创建了TransportChannelHandler之后，对管道pipeline进行设置，代码如下：

TransportFrameDecoder、MessageDecoder、TransportChannelHandler本质上都是继承了ChannelInboundHandler，MessageEncoder本质上都是继承了ChannelOutboundHandler，IdleStateHandler本质上都是继承了ChannelInboundHandler，ChannelOutboundHandler（继承和接口实现），根据Netty中handler的执行顺序，得出如下：

Request时按照顺序执行，TransportFrameDecoder-》MessageDecoder-》IdleStateHandler-》TransportChannelHandler，Response时按照逆序执行，IdleStateHandler-》MessageEncoder。结构如下：

RPC服务端处理RpcHandler

下面的代码是TransportRequestHandler中的代码，可以清楚的看到，TransportRequestHandler是将请求消息交给rpcHandler做进一步处理。

RpcHandler是一个抽象类，主要有以下几个方法：

receive方法，接收单一PRC消息，RpcResponseCallback用来处理结束后的回掉，无论成功与否，都会执行一次。有一个receive重载方法，默认执行OneWayRpcCallback回调，这个回调只负责打印成功和失败时的信息。

/**
   * 抽象方法用来接收单一RPC消息
   * RpcResponseCallback用来处理结束后的回掉，无论成功与否，都会执行一次
   * */
  public abstract void receive(
      TransportClient client,
      ByteBuffer message,
      RpcResponseCallback callback);
/**
   * 重载receive方法，默认执行ONE_WAY_CALLBACK回调
   * */
  public void receive(TransportClient client, ByteBuffer message) {
    receive(client, message, ONE_WAY_CALLBACK);
  }

getStreamManager方法，抽象方法，用于获取getStreamManager

/**
   * 抽象方法获取StreamManager
   * */
  public abstract StreamManager getStreamManager();

channelActive、channelInactive、exceptionCaught方法，分别与客户端相关联的channel处于活动/非活动/异常状态时调用

  /**
   * 当与客户端相关联的channel处于活动状态时调用
   * */
  public void channelActive(TransportClient client) { }

  /**
   * 当与客户端相关联的channel处于非活动状态时调用
   * */
  public void channelInactive(TransportClient client) { }

  /**
   * 当与客户端相关联的channel产生异常时调用
   * */
  public void exceptionCaught(Throwable cause, TransportClient client) { }

了解了RpcHandler的结构后，再看一下TransportRequestHandler的handle(RequestMessage request)方法

public void handle(RequestMessage request) {
    if (request instanceof ChunkFetchRequest) {
      processFetchRequest((ChunkFetchRequest) request);
    } else if (request instanceof RpcRequest) {
      processRpcRequest((RpcRequest) request);
    } else if (request instanceof OneWayMessage) {
      processOneWayMessage((OneWayMessage) request);
    } else if (request instanceof StreamRequest) {
      processStreamRequest((StreamRequest) request);
    } else {
      throw new IllegalArgumentException("Unknown request type: " + request);
    }
  }

TransportRequestHandler处理4中RequestMessage：

1、处理块获取请求

private void processFetchRequest(final ChunkFetchRequest req) {
    if (logger.isTraceEnabled()) {
      logger.trace("Received req from {} to fetch block {}", getRemoteAddress(channel),
        req.streamChunkId);
    }

    ManagedBuffer buf;
    try {
      // this.streamManager = rpcHandler.getStreamManager();
      // 校验客户端是否有权限从流中读取消息
      streamManager.checkAuthorization(reverseClient, req.streamChunkId.streamId);
      // 将一个流与一个客户端的TCP链接关联起来，单个流只会有一个客户端读取
      streamManager.registerChannel(channel, req.streamChunkId.streamId);
      // 获取块
      buf = streamManager.getChunk(req.streamChunkId.streamId, req.streamChunkId.chunkIndex);
    } catch (Exception e) {
      logger.error(String.format("Error opening block %s for request from %s",
        req.streamChunkId, getRemoteAddress(channel)), e);
      respond(new ChunkFetchFailure(req.streamChunkId, Throwables.getStackTraceAsString(e)));
      return;
    }
    // 将ManagedBuffer和流的块ID封装到ChunkFetchSuccess中，调用respond方法返回给客户端
    respond(new ChunkFetchSuccess(req.streamChunkId, buf));
  }

processFetchRequest做了以下几件事：

校验客户端是否有权限从流中读取消息
将一个流与一个客户端的TCP链接关联起来，单个流只会有一个客户端读取
获取块
将ManagedBuffer和流的块ID封装到ChunkFetchSuccess中，调用respond方法返回给客户端

2、处理RPC请求

代码如下：

private void processRpcRequest(final RpcRequest req) {
    try {
      /**
       * 将发送消息的客户端、RpcRequest消息的内容和RpcResponseCallback回调类作为参数传递给RpcHandler的receive方法
       * 所以说真正处理消息的是RpcHandler，而不是TrnsportRequestHandler
       * */
      rpcHandler.receive(reverseClient, req.body().nioByteBuffer(), new RpcResponseCallback() {
        @Override
        public void onSuccess(ByteBuffer response) {
          respond(new RpcResponse(req.requestId, new NioManagedBuffer(response)));
        }

        @Override
        public void onFailure(Throwable e) {
          respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
        }
      });
    } catch (Exception e) {
      logger.error("Error while invoking RpcHandler#receive() on RPC id " + req.requestId, e);
      respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
    } finally {
      req.body().release();
    }
  }

将发送消息的客户端、RpcRequest消息的内容和RpcResponseCallback回调类作为参数传递给RpcHandler的receive方法，所以说真正处理消息的是RpcHandler，而不是TrnsportRequestHandler。

3、处理无需回复的RPC请求

处理无需回复的RPC请求，回调类是OneWayRpcCallback，处理完RPC请求后不会给客户端作出响应。

private void processOneWayMessage(OneWayMessage req) {
    try {
      /**
       * 处理无需回复的RPC请求，回调类是OneWayRpcCallback，处理完RPC请求后不会给客户端作出响应
       * */
      rpcHandler.receive(reverseClient, req.body().nioByteBuffer());
    } catch (Exception e) {
      logger.error("Error while invoking RpcHandler#receive() for one-way message.", e);
    } finally {
      req.body().release();
    }
  }

4、处理流请求

使用streamManager.openStream方法将流数据封装为ManagedBuffer。

private void processStreamRequest(final StreamRequest req) {
    ManagedBuffer buf;
    try {
      // 将获取的流数据封装为ManagedBuffer
      buf = streamManager.openStream(req.streamId);
    } catch (Exception e) {
      logger.error(String.format(
        "Error opening stream %s for request from %s", req.streamId, getRemoteAddress(channel)), e);
      respond(new StreamFailure(req.streamId, Throwables.getStackTraceAsString(e)));
      return;
    }
    // 无论成功还是失败，都要响应客户端
    if (buf != null) {
      respond(new StreamResponse(req.streamId, buf.size(), buf));
    } else {
      respond(new StreamFailure(req.streamId, String.format(
        "Stream '%s' was not found.", req.streamId)));
    }
  }

上面这四种处理请求的方法，除了processOneWayMessage不需要调用respond方法外，其他三个都需要调用respond方法，用来响应客户端。respond方法中实际是调用channel.writeAndFlush来响应客户端的。

private void respond(final Encodable result) {
    final SocketAddress remoteAddress = channel.remoteAddress();
    //响应客户端
    channel.writeAndFlush(result).addListener(
      new ChannelFutureListener() {
        @Override
        public void operationComplete(ChannelFuture future) throws Exception {
          if (future.isSuccess()) {
            logger.trace("Sent result {} to client {}", result, remoteAddress);
          } else {
            logger.error(String.format("Error sending result %s to %s; closing connection",
              result, remoteAddress), future.cause());
            channel.close();
          }
        }
      }
    );
  }

引导程序Bootstrap

在TransportServer中有一个成员变量List<TransportServerBootstrap>是TransportServer引导程序列表，在初始化管道时，调用了每一个引导程序的doBootstrap方法。

TransportServerBootstrap定义了服务端引导程序的规范，当客户端与服务端建立了连接，在服务端持有的客户端管道上执行引导程序。

TransportServerBootstrap接口定义如下：

public interface TransportServerBootstrap {
  
  RpcHandler doBootstrap(Channel channel, RpcHandler rpcHandler);
}

TransportServerBootstrap有两个实现类，一个是SaslServerBootstrap，另一个是EncryptionCheckerBootstrap，以SaslServerBootstrap为例说明引导程序的作用。

直接看SaslServerBootstrap的doBootstrap方法：

public RpcHandler doBootstrap(Channel channel, RpcHandler rpcHandler) {
    return new SaslRpcHandler(conf, channel, rpcHandler, secretKeyHolder);
  }

doBootstrap方法直接创建了一个RpcHandler的具体实现类SaslRpcHandler。SaslRpcHandler负责对管道进行SASL加密，它集成了RpcHandler，所以核心代码就在receive中：

public void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback) {
    if (isComplete) {
      // Authentication complete, delegate to base handler.
      // 将处理好的消息传递给下游的RpcHandler
      delegate.receive(client, message, callback);
      return;
    }

    ByteBuf nettyBuf = Unpooled.wrappedBuffer(message);
    SaslMessage saslMessage;
    try {
      // 进行SASL加密
      saslMessage = SaslMessage.decode(nettyBuf);
    } finally {
      nettyBuf.release();
    }

    if (saslServer == null) {
      // First message in the handshake, setup the necessary state.
      client.setClientId(saslMessage.appId);
      // 如果saslServer为空，创建SparkSaslServer
      saslServer = new SparkSaslServer(saslMessage.appId, secretKeyHolder,
        conf.saslServerAlwaysEncrypt());
    }

    byte[] response;
    try {
      // saslServer处理已经解密的信息
      response = saslServer.response(JavaUtils.bufferToArray(
        saslMessage.body().nioByteBuffer()));
    } catch (IOException ioe) {
      throw new RuntimeException(ioe);
    }
    callback.onSuccess(ByteBuffer.wrap(response));

    // Setup encryption after the SASL response is sent, otherwise the client can't parse the
    // response. It's ok to change the channel pipeline here since we are processing an incoming
    // message, so the pipeline is busy and no new incoming messages will be fed to it before this
    // method returns. This assumes that the code ensures, through other means, that no outbound
    // messages are being written to the channel while negotiation is still going on.
    if (saslServer.isComplete()) {
      logger.debug("SASL authentication successful for channel {}", client);
      isComplete = true; // 处理完成
      if (SparkSaslServer.QOP_AUTH_CONF.equals(saslServer.getNegotiatedProperty(Sasl.QOP))) {
        logger.debug("Enabling encryption for channel {}", client);
        // 进行管道加密
        SaslEncryption.addToChannel(channel, saslServer, conf.maxSaslEncryptedBlockSize());
        saslServer = null;
      } else {
        saslServer.dispose();
        saslServer = null;
      }
    }
  }

receive做了以下几件事：

如果认证已经完成（isComplete=true）,将消息传递给下游RpcHandler
如果认证未经完成（isComplete=false）对客户端发送的消息进行加密
如果saslServer=null，创建SparkSaslServer，SaslRpcHandler接收客户端第一条消息时执行此操作
使用saslServer处理已解密的消息，并执行回调返回给客户端
如果认证已经完成，改变isComplete=true
对管道进行Sasl加密

可以看到引导程序主要起到了引导、包装、传递、代理的作用，类似的还有TransportClientBootstrap。

RPC客户端TransportClient

看完RPC服务端利用RpcHandler处理消息后，这里也看看RPC客户端如何处理消息的，在TransportContext的createChannelHandler中创建TransportClient。TransportClient一共有5个方法用于发送请求：

fetchChunk：从远端协商好的流中请求单个块
stream：使用流的ID，从远端获取流数据
sendRpc：向服务端发送RPC请求，通过at least once delivery原则保证请求不丢失
sendRpcSync：向服务端发送异步RPC请求
send：想服务端发送RPC请求，但并期望获取响应，不能保证可靠性

这里重点分析一下sendRpc方法，代码如下：

public long sendRpc(ByteBuffer message, final RpcResponseCallback callback) {
    final long startTime = System.currentTimeMillis();
    if (logger.isTraceEnabled()) {
      logger.trace("Sending RPC to {}", getRemoteAddress(channel));
    }

    // UUID生成requestId
    final long requestId = Math.abs(UUID.randomUUID().getLeastSignificantBits());
    /**
     * 这里的handler是TransportResponseHandler
     * 更新最后一次请求时间
     * addRpcRequest中利用Map设置requestId和回调类的关系，requestId为key，callback为value
     * */
    handler.addRpcRequest(requestId, callback);

    /**
     * channel.writeAndFlush发送请求，无论成功还是失败都会回调ChannelFutureListener的operationComplete方法
     * 成功的话打印日志信息
     * 失败的话不仅要打印日志，还要执行handler.removeRpcRequest(requestId)，移除此次请求
     * */
    channel.writeAndFlush(new RpcRequest(requestId, new NioManagedBuffer(message))).addListener(
      new ChannelFutureListener() {
        @Override
        public void operationComplete(ChannelFuture future) throws Exception {
          if (future.isSuccess()) {
            long timeTaken = System.currentTimeMillis() - startTime;
            if (logger.isTraceEnabled()) {
              logger.trace("Sending request {} to {} took {} ms", requestId,
                getRemoteAddress(channel), timeTaken);
            }
          } else {
            String errorMsg = String.format("Failed to send RPC %s to %s: %s", requestId,
              getRemoteAddress(channel), future.cause());
            logger.error(errorMsg, future.cause());
            handler.removeRpcRequest(requestId);
            channel.close();
            try {
              callback.onFailure(new IOException(errorMsg, future.cause()));
            } catch (Exception e) {
              logger.error("Uncaught exception in RPC response callback handler!", e);
            }
          }
        }
      });

    return requestId;
  }

从上面的代码可以看出，sendRpc做了如下事情：

利用UUID生成requestId
更新最后一次请求时间，并利用Map设置了requestId和回调类的对应关系
channel.writeAndFlush发送请求，无论成功还是失败都会回调ChannelFutureListener的operationComplete方法，成功的话打印日志信息，失败的话不仅要打印日志，还要执行handler.removeRpcRequest(requestId)，移除此次请求
返回requestId

请求发送成功后，客户端将会等待接收服务端响应，返回的消息会传会给TransportChannelHandler的channelRead方法。

接着进入responseHandler.handle((ResponseMessage) request);方法，其中有6中类型的判断，RPC对应的是RpcResponse和RpcFailure。

RpcResponse对应的处理如下：

RpcFailure对应的处理如下：

总结

根据上面对Spark RPC组件的分析可以得到RPC客户端服务端的请求响应流程，如下图所示：

客户端发送请求是在TransportClient的对应方法执行了channel.writeAndFlush方法，并设置了成功和失败监听；服务端响应请求是TransportRequestHandler的respond方法中执行了channel.writeAndFlush方法。无论是服务端得到请求还是客户端接收响应都是通过TransportChannelHandler的channelRead方法判断Message类型，服务端得到请求交由TransportRequestHandler处理，客户端接收响应交由TransportResponseHandler处理。引导程序则贯穿于各个步骤当中。

最后总结出Spark RPC框架结构如下图所示：