Flink源码分析之RPC通信

RpcEndpoint.java
  
    protected RpcEndpoint(final RpcService rpcService, final String endpointId) {
        this.rpcService = checkNotNull(rpcService, "rpcService");
        this.endpointId = checkNotNull(endpointId, "endpointId");
				
  			// 生成代理对象
        this.rpcServer = rpcService.startServer(this);

        this.mainThreadExecutor = new MainThreadExecutor(rpcServer, this::validateRunsInMainThread);
    }

AkkaRpcService.java
    startServer(C rpcEndpoint)
        registerAkkaRpcActor(rpcEndpoint)
   		{
     			 	// 创建Actor
             SupervisorActor.startAkkaRpcActor(
               supervisor.getActor(),
               actorTerminationFuture ->
               Props.create(
                 akkaRpcActorType,
                 rpcEndpoint,
                 actorTerminationFuture,
                 getVersion(),
                 configuration.getMaximumFramesize(),
                 flinkClassLoader),
               rpcEndpoint.getEndpointId());  
   
             // 在RpcService中保存ActorRef与RpcEndpoint引用关系
             actors.put(actorRegistration.getActorRef(), rpcEndpoint)    
   		}

SupervisorActor.java
     
   		// 1) 发送消息
       public static StartAkkaRpcActorResponse startAkkaRpcActor(
               ActorRef supervisor, StartAkkaRpcActor.PropsFactory propsFactory, String endpointId) {
     
       			// 以Ask方式发送消息并等待结果
     				// Ask在实现上实际上是会创建一个Actor等待响应结果，成功或者超时时，销毁Actor
           return Patterns.ask(
                           supervisor,
                           createStartAkkaRpcActorMessage(propsFactory, endpointId),
                           RpcUtils.INF_DURATION)
                   .toCompletableFuture()
                   .thenApply(SupervisorActor.StartAkkaRpcActorResponse.class::cast)
                   .join();
       } 
   
   		// 2) 处理消息
       private void createStartAkkaRpcActorMessage(StartAkkaRpcActor startAkkaRpcActor) {
           final String endpointId = startAkkaRpcActor.getEndpointId();
           final Props akkaRpcActorProps = ...
           ...
   
           try {
             	// 创建Actor
               final ActorRef actorRef = getContext().actorOf(akkaRpcActorProps, endpointId);
   
               registeredAkkaRpcActors.put(actorRef, akkaRpcActorRegistration);
   
             	// 回复消息
               getSender().tell(
                               	StartAkkaRpcActorResponse.success(...),
                               	getSelf()
               								);
           } catch (AkkaException akkaException) {
               getSender().tell(StartAkkaRpcActorResponse.failure(akkaException), getSelf());
           }
       }

AkkaRpcService.java
    startServer(C rpcEndpoint)
   	{
   		...
        // 服务端对象实现了RpcGateway接口
        Set<Class<?>> implementedRpcGateways =
        new HashSet<>(RpcUtils.extractImplementedRpcGateways(rpcEndpoint.getClass()));
            
            // 服务端对象是一个RpcServer
        implementedRpcGateways.add(RpcServer.class);
            // 服务端对象是Akka endpoint，可以获取到ActorRef引用
        implementedRpcGateways.add(AkkaBasedEndpoint.class); 
    
        if (rpcEndpoint instanceof FencedRpcEndpoint){
            implementedRpcGateways.add(FencedMainThreadExecutable.class);
        }
   	}

AkkaRpcService.java
   	startServer(C rpcEndpoint)
   	{
     	...
        RpcServer server =
                   (RpcServer)
                           Proxy.newProxyInstance(
                                   classLoader,
                                   implementedRpcGateways.toArray(
                                           new Class<?>[implementedRpcGateways.size()]),
                                   akkaInvocationHandler);
   
         return server;        
   	}

AkkaRpcService.java
  connect(String address, Class<C> clazz)
		{
  			...
        // 1) 使用ActorSystem.actorSelection(address).resolveOne的方式来获取Actor的引用ActorRef(ActorRef可以用来向服务端Actor发送消息)
        final CompletableFuture<ActorRef> actorRefFuture = resolveActorAddress(address);
  
  			// 2) ActorRef创建完成后，使用ask的方式向服务端发送一条握手消息(用来验证Client和Server彼此版本一致)
        final CompletableFuture<HandshakeSuccessMessage> handshakeFuture =
                actorRefFuture.thenCompose(
                        (ActorRef actorRef) ->
                                AkkaFutureUtils.toJava(
                                        Patterns.ask(
                                                        actorRef,
                                                        new RemoteHandshakeMessage(
                                                                clazz, getVersion()),
                                                        configuration.getTimeout().toMilliseconds())
                                                .<HandshakeSuccessMessage>mapTo(
                                                        ClassTag$.MODULE$
                                                                .<HandshakeSuccessMessage>apply(
                                                                        HandshakeSuccessMessage
                                                                                .class))));  
  
  			// 3) 以上2个事都做完后，异步创建代理对象并返回
        final CompletableFuture<C> gatewayFuture =
                actorRefFuture.thenCombineAsync(
                        handshakeFuture,
                        (ActorRef actorRef, HandshakeSuccessMessage ignored) -> {
                          	// invocationHandlerFactory.apply(actorRef) = new AkkaInvocationHandler 或 new 或 FencedAkkaInvocationHandler
                            InvocationHandler invocationHandler =
                                    invocationHandlerFactory.apply(actorRef);

                            ClassLoader classLoader = getClass().getClassLoader();

                            C proxy =
                                    (C)
                                            Proxy.newProxyInstance(
                                                    classLoader,
                                                    new Class<?>[] {clazz},
                                                    invocationHandler);

                            return proxy;
                        },
                        actorSystem.dispatcher());

        return guardCompletionWithContextClassLoader(gatewayFuture, flinkClassLoader);  	
		}

class AkkaInvocationHandler implements InvocationHandler, AkkaBasedEndpoint, RpcServer{
  
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        Class<?> declaringClass = method.getDeclaringClass();

        Object result;

      	// 非Rpc方法，直接本地执行。这个是服务端通过自己的代理对象RpcServer调用自己非Rpc方法时走的逻辑
        if (declaringClass.equals(AkkaBasedEndpoint.class)
                || declaringClass.equals(Object.class)
                || declaringClass.equals(RpcGateway.class)
                || declaringClass.equals(StartStoppable.class)
                || declaringClass.equals(MainThreadExecutable.class)
                || declaringClass.equals(RpcServer.class)) {
            result = method.invoke(this, args);
        } else if (declaringClass.equals(FencedRpcGateway.class)) { // 支持HA的见FencedAkkaInvocationHandler
            throw new UnsupportedOperationException(...);
        } else {
          	// RPC方法，指RpcGateway子接口中定义的方法
          	// 接口：ResourceManagerGateway、DispatcherGateway、JobMasterGateway、MetricQueryServiceGateway、TaskExecutorGateway
            result = invokeRpc(method, args);
        }

        return result;
    }
  
    private Object invokeRpc(Method method, Object[] args) throws Exception {
				...

        // 1) 封装消息
        final RpcInvocation rpcInvocation =
                createRpcInvocationMessage(methodName, parameterTypes, args);

        // 2) 借助akka发送消息，进行RPC调用
        Class<?> returnType = method.getReturnType();

        final Object result;

        if (Objects.equals(returnType, Void.TYPE)) {
          	// 无返回值，用akka tell模式
            tell(rpcInvocation);

            result = null;
        } else {
						...
            // 有返回值，用akka ask模式
            final CompletableFuture<?> resultFuture = ask(rpcInvocation, futureTimeout);
						...
        }

        return result;
    }
  
}

AkkaRpcActor.java
  
    private void handleRpcInvocation(RpcInvocation rpcInvocation) {
        Method rpcMethod = null;

        try {
            String methodName = rpcInvocation.getMethodName();
            Class<?>[] parameterTypes = rpcInvocation.getParameterTypes();

            rpcMethod = lookupRpcMethod(methodName, parameterTypes);
        } catch (Exception e) {
          	...
            getSender().tell(new Status.Failure(rpcException), getSelf());
        }

        if (rpcMethod != null) {
            try {
                rpcMethod.setAccessible(true);
                final Method capturedRpcMethod = rpcMethod;
              
              	// 1) 无返回值
                if (rpcMethod.getReturnType().equals(Void.TYPE)) {
                    runWithContextClassLoader(
                            () -> capturedRpcMethod.invoke(rpcEndpoint, rpcInvocation.getArgs()),
                            flinkClassLoader);
                } else {
                  	// 2) 有返回值
                    final Object result;
                    try {
                        result =
                                runWithContextClassLoader(
                                        () ->
                                                capturedRpcMethod.invoke(
                                                        rpcEndpoint, rpcInvocation.getArgs()),
                                        flinkClassLoader);
                    } catch (InvocationTargetException e) {
												...
                        getSender().tell(new Status.Failure(e.getTargetException()), getSelf());
                        return;
                    }

                    final String methodName = rpcMethod.getName();

                    if (result instanceof CompletableFuture) {
                        final CompletableFuture<?> responseFuture = (CompletableFuture<?>) result;
                        sendAsyncResponse(responseFuture, methodName);
                    } else {
                        sendSyncResponse(result, methodName);
                    }
                }
            } catch (Throwable e) {
              	...
                getSender().tell(new Status.Failure(e), getSelf());
            }
        }
    }  

JobMaster.java
     reconnectToResourceManager
     	tryConnectToResourceManager
     		connectToResourceManager
   			{
     			// 这里用resourceManagerAddress地址重新构建一个连接
     			// RM地址发生切换时，resourceManagerAddress值也会得到更新
   				resourceManagerConnection = new ResourceManagerConnection(..., resourceManagerAddress, ...)  
                resourceManagerConnection.start	
   			}
     
RegisteredRpcConnection.java
    start
    	newRegistration.startRegistration
     
RetryingRegistration.java
    startRegistration
     	// 1). 建立与RM的连接
     	rpcService.connect	
     	// 2). 向RM注册自己
     	register
     		JobMaster.ResourceManagerConnection.invokeRegistration
     			gateway.registerJobManager	
   				{
     				// 3). RM方法内部回连JM
     				getRpcService().connect(jobManagerAddress, jobMasterId, JobMasterGateway.class)
   				}

JobMaster.java
   
  // LeaderRetrievalService用于发现leader，启用HA时start方法会创建一个具体的LeaderRetrievalDriver
  // driver上的leader切换事件最后会通知到LeaderRetrievalListener的notifyLeaderAddress上
  // ResourceManagerLeaderListener实现了LeaderRetrievalListener接口
  LeaderRetrievalService resourceManagerLeaderRetriever = highAvailabilityServices.getResourceManagerLeaderRetriever()
  resourceManagerLeaderRetriever.start(new ResourceManagerLeaderListener())
  ResourceManagerLeaderListener.notifyLeaderAddress(leaderAddress, leaderSessionID)  

JobMaster.java
    ResourceManagerHeartbeatListener.notifyHeartbeatTimeout
        reconnectToResourceManager

JobMaster.java
    ResourceManagerLeaderListener.notifyLeaderAddress
        notifyOfNewResourceManagerLeader
        	createResourceManagerAddress
        	reconnectToResourceManager

ResourceManager.java
    JobManagerHeartbeatListener.notifyHeartbeatTimeout
        closeJobManagerConnection
        	jobMasterGateway.disconnectResourceManager(getFencingToken(), cause)
      		{
        		// JobMaster disconnectResourceManager方法内部
        		if (isConnectingToResourceManager(resourceManagerId)) {
                  	reconnectToResourceManager(cause);
              	}
      		}
      
      
    private boolean isConnectingToResourceManager(ResourceManagerId resourceManagerId) {
        return resourceManagerAddress != null
            && resourceManagerAddress.getResourceManagerId().equals(resourceManagerId);
    }