Flink中基于Akka的RPC实现

Flink中基于Akka的RPC实现

版本说明：Flink: 1.10.1

1 前言

Flink中RPC是基于Akka实现的，在上一篇文章《使用Akka实现简单RPC框架》中，使用Akka的基本API加上Java动态代理实现了一个简单的RPC框架，对Akka不太熟悉的同学可以先参考那篇文章手写一下代码，然后再来阅读这篇文章会更好的理解Flink的RPC实现，基本原理都是一样的。我也是先看的Akka相关知识，然后再重新阅读的Flink代码，不得不承认，我上一篇文章的实现并没有Flink中实现的优雅，Flink里确实有不少值得借鉴的地方。本篇文章会先简单介绍一下Flink中RPC的实现原理，然后讲解几个比较关键的类实现，最后会使用Flink的RPC框架实现一个简单例子。

2 Flink中RPC的实现原理

Flink中RPC的实现原理与我之前文章里实现的一直，都是使用Akka actor实现通讯和序列化，使用动态代理和反射实现远程调用。如下图所示，用户需要自定义一个UDRpcGateway接口，用户自定义接口需要继承RpcGateway接口，当前接口需要定义远程调用的方法。服务端需要创建一个AkkaRpcService实例，继承RpcEndpoint和实现UDRpcGateway相关接口，AkkaRpcService会创建一个AkkaRpcActor，AkkaRpcActo监听指定的端口，当有请求到达之后会根据请求信息调用UDRpcEndpoint中对应方法。客户端同样需要创建一个AkkaRpcService服务，并调用connect方法连接远程AkkaRpcService，连接成功后会创建一个UDRpcGateway的代理，客户端可以直接调用UDRpcGateway中定义的方法，然后通过代理进入AkkaInvocationHandler中进行处理，AkkaInvocationHandler被调用会，会通过ActorRef与远程的actor进行远程调用。

3 关键类实现讲解

从上面原理中，我们可以看到几个比较重要的实现类及接口，RpcGateway、RpcEndpoint、AkkaRpcService、AkkaRpcActor、AkkaInvocationHandler，下面将分别介绍这几个类的作用及关键代码的解释。

3.1 RpcGateway

Rpc网关接口，用户RPC调用的方法，都需要在RpcGateway接口中定义，代码位置：org.apache.flink.runtime.rpc.RpcGateway

3.2 RpcEndpoint

Rpc端点的基类类，用户需要实现它，并实现用户RpcGateway中定义的接口，该类是RPC的入口，通过它启动AkkaRpcService以及相关的AkkaPrcActor。代码位置：org.apache.flink.runtime.rpc.RpcEndpoint

3.3 AkkaRpcService

RPC服务的核心实现，其中包括创建用以接收远程调用的Actor和连接远程Actor并获取代理的逻辑。diamante位置：org.apache.flink.runtime.rpc.akka.AkkaRpcService

其中有比较关键的几个实现方法:

startServer(C rpcEndpoint)

	/**
	 * 启动PRC服务
	 *
	 * @param rpcEndpoint Rpc protocol to dispatch the rpcs to
	 * @param <C>         用户定义的RPCEndpoint
	 * @return RPCEndpoint代理
	 */
	@Override
	public <C extends RpcEndpoint & RpcGateway> RpcServer startServer(C rpcEndpoint) {
		checkNotNull(rpcEndpoint, "rpc endpoint");

		CompletableFuture<Void> terminationFuture = new CompletableFuture<>();
		final Props akkaRpcActorProps;

		// 根据RpcEndpoint的类型来创建对应的Actor，目前支持两种Actor的创建
		// 1. AkkaRpcActor
		// 2. FencedAkkaRpcActor 对AkkaRpcActor进行扩展，能够过滤到与指定token无关的消息
		if (rpcEndpoint instanceof FencedRpcEndpoint) {
			akkaRpcActorProps = Props.create(
				FencedAkkaRpcActor.class,
				rpcEndpoint,
				terminationFuture,
				getVersion(),
				configuration.getMaximumFramesize());
		} else {
			akkaRpcActorProps = Props.create(
				AkkaRpcActor.class,
				rpcEndpoint,
				terminationFuture,
				getVersion(),
				configuration.getMaximumFramesize());
		}

		ActorRef actorRef;

		synchronized (lock) {
			checkState(!stopped, "RpcService is stopped");
			actorRef = actorSystem.actorOf(akkaRpcActorProps, rpcEndpoint.getEndpointId());
			actors.put(actorRef, rpcEndpoint);
		}

		LOG.info("Starting RPC endpoint for {} at {} .", rpcEndpoint.getClass().getName(), actorRef.path());

		final String akkaAddress = AkkaUtils.getAkkaURL(actorSystem, actorRef);
		final String hostname;
		Option<String> host = actorRef.path().address().host();
		if (host.isEmpty()) {
			hostname = "localhost";
		} else {
			hostname = host.get();
		}

		// 提取集成RpcEndpoint的所有子类
		Set<Class<?>> implementedRpcGateways = new HashSet<>(RpcUtils.extractImplementedRpcGateways(rpcEndpoint.getClass()));

		implementedRpcGateways.add(RpcServer.class);
		implementedRpcGateways.add(AkkaBasedEndpoint.class);

		// 对上述指定的类集合进行代理
		final InvocationHandler akkaInvocationHandler;

		if (rpcEndpoint instanceof FencedRpcEndpoint) {
			// a FencedRpcEndpoint needs a FencedAkkaInvocationHandler
			akkaInvocationHandler = new FencedAkkaInvocationHandler<>(
				akkaAddress,
				hostname,
				actorRef,
				configuration.getTimeout(),
				configuration.getMaximumFramesize(),
				terminationFuture,
				((FencedRpcEndpoint<?>) rpcEndpoint)::getFencingToken);

			implementedRpcGateways.add(FencedMainThreadExecutable.class);
		} else {
			akkaInvocationHandler = new AkkaInvocationHandler(
				akkaAddress,
				hostname,
				actorRef,
				configuration.getTimeout(),
				configuration.getMaximumFramesize(),
				terminationFuture);
		}

		// Rather than using the System ClassLoader directly, we derive the ClassLoader
		// from this class . That works better in cases where Flink runs embedded and all Flink
		// code is loaded dynamically (for example from an OSGI bundle) through a custom ClassLoader
		ClassLoader classLoader = getClass().getClassLoader();

		@SuppressWarnings("unchecked")
		RpcServer server = (RpcServer) Proxy.newProxyInstance(
			classLoader,
			implementedRpcGateways.toArray(new Class<?>[implementedRpcGateways.size()]),
			akkaInvocationHandler);

		return server;
	}

connect

		/***
	 * 连接远程RPC Server
	 * @param address Address of the remote rpc server
	 * @param clazz Class of the rpc gateway to return
	 * @param <C> 用于自定义的RpcGateway
	 * @return 用于自定义的RpcGateway代理
	 */
	// this method does not mutate state and is thus thread-safe
	@Override
	public <C extends RpcGateway> CompletableFuture<C> connect(
		final String address,
		final Class<C> clazz) {

		return connectInternal(
			address,
			clazz,
			(ActorRef actorRef) -> {
				Tuple2<String, String> addressHostname = extractAddressHostname(actorRef);

				return new AkkaInvocationHandler(
					addressHostname.f0,
					addressHostname.f1,
					actorRef,
					configuration.getTimeout(),
					configuration.getMaximumFramesize(),
					null);
			});
	}

  /**
	 * 连接远程RPC
	 * @param address 远程RPC服务的地址
	 * @param clazz 用户自定义RpcGateway.class
	 * @param invocationHandlerFactory 代理处理器工厂实例
	 * @param <C> 用户自定义RpcGateway
	 * @return 用户自定义RpcGateway实例
	 */
	private <C extends RpcGateway> CompletableFuture<C> connectInternal(
		final String address,
		final Class<C> clazz,
		Function<ActorRef, InvocationHandler> invocationHandlerFactory) {
		checkState(!stopped, "RpcService is stopped");

		LOG.debug("Try to connect to remote RPC endpoint with address {}. Returning a {} gateway.",
			address, clazz.getName());

		// 根据Akka Actor地址获取ActorRef
		final ActorSelection actorSel = actorSystem.actorSelection(address);

		final Future<ActorIdentity> identify = Patterns
			.ask(actorSel, new Identify(42), configuration.getTimeout().toMilliseconds())
			.<ActorIdentity>mapTo(ClassTag$.MODULE$.<ActorIdentity>apply(ActorIdentity.class));

		final CompletableFuture<ActorIdentity> identifyFuture = FutureUtils.toJava(identify);

		final CompletableFuture<ActorRef> actorRefFuture = identifyFuture.thenApply(
			(ActorIdentity actorIdentity) -> {
				if (actorIdentity.getRef() == null) {
					throw new CompletionException(new RpcConnectionException("Could not connect to rpc endpoint under address " + address + '.'));
				} else {
					return actorIdentity.getRef();
				}
			});

		// 发送一个握手成功的消息给远程Actor
		final CompletableFuture<HandshakeSuccessMessage> handshakeFuture = actorRefFuture.thenCompose(
			(ActorRef actorRef) -> FutureUtils.toJava(
				Patterns
					.ask(actorRef, new RemoteHandshakeMessage(clazz, getVersion()), configuration.getTimeout().toMilliseconds())
					.<HandshakeSuccessMessage>mapTo(ClassTag$.MODULE$.<HandshakeSuccessMessage>apply(HandshakeSuccessMessage.class))));

		// 创建动态代理，并返回
		return actorRefFuture.thenCombineAsync(
			handshakeFuture,
			(ActorRef actorRef, HandshakeSuccessMessage ignored) -> {
				InvocationHandler invocationHandler = invocationHandlerFactory.apply(actorRef);

				// Rather than using the System ClassLoader directly, we derive the ClassLoader
				// from this class . That works better in cases where Flink runs embedded and all Flink
				// code is loaded dynamically (for example from an OSGI bundle) through a custom ClassLoader
				ClassLoader classLoader = getClass().getClassLoader();

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

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

3.4 AkkaRpcActor

就是一个Akka的Actor实现，主要看createReceive()、handleHandshakeMessage(RemoteHandshakeMessage handshakeMessage)、handleControlMessage(ControlMessages controlMessage)、handleMessage(final Object message)几个方法即可。代码位置：org.apache.flink.runtime.rpc.akka.AkkaRpcActor

createReceive()

	/**
	 * 创建消息接收器
	 * @return Receive
	 */
	@Override
	public Receive createReceive() {
		return ReceiveBuilder.create()
			// 处理握手消息
			.match(RemoteHandshakeMessage.class, this::handleHandshakeMessage)
			// 处理控制消息，如启动、停止、中断，START、STOP、TERMINATE
			.match(ControlMessages.class, this::handleControlMessage)
			// 处理通用消息
			.matchAny(this::handleMessage)
			.build();
	}

handleHandshakeMessage

	/**
	 * 处理握手消息
	 * @param handshakeMessage 握手消息
	 */
	private void handleHandshakeMessage(RemoteHandshakeMessage handshakeMessage) {
		// 判断消息是否是兼容版本
		if (!isCompatibleVersion(handshakeMessage.getVersion())) {
      		// 发送失败消息
			sendErrorIfSender(new AkkaHandshakeException(
				String.format(
					"Version mismatch between source (%s) and target (%s) rpc component. Please verify that all components have the same version.",
					handshakeMessage.getVersion(),
					getVersion())));
			// 判断RpcGateway是否继承自RpcGateway
		} else if (!isGatewaySupported(handshakeMessage.getRpcGateway())) {
      		// 发送失败消息
			sendErrorIfSender(new AkkaHandshakeException(
				String.format(
					"The rpc endpoint does not support the gateway %s.",
					handshakeMessage.getRpcGateway().getSimpleName())));
		} else {
			// 发送握手成功消息
			getSender().tell(new Status.Success(HandshakeSuccessMessage.INSTANCE), getSelf());
		}
	}

handleControlMessage

	/**
	 * 处理控制消息，根据对象的消息状态，调用对应的方法
	 * @param controlMessage 控制消息
	 */
	private void handleControlMessage(ControlMessages controlMessage) {
		try {
			switch (controlMessage) {
				case START:
					state = state.start(this);
					break;
				case STOP:
					state = state.stop();
					break;
				case TERMINATE:
					state = state.terminate(this);
					break;
				default:
					handleUnknownControlMessage(controlMessage);
			}
		} catch (Exception e) {
			this.rpcEndpointTerminationResult = RpcEndpointTerminationResult.failure(e);
			throw e;
		}
	}

handleMessage

	private void handleMessage(final Object message) {
		if (state.isRunning()) {
			mainThreadValidator.enterMainThread();

			try {
				handleRpcMessage(message);
			} finally {
				mainThreadValidator.exitMainThread();
			}
		} else {
			log.info("The rpc endpoint {} has not been started yet. Discarding message {} until processing is started.",
				rpcEndpoint.getClass().getName(),
				message.getClass().getName());

			sendErrorIfSender(new AkkaRpcException(
				String.format("Discard message, because the rpc endpoint %s has not been started yet.", rpcEndpoint.getAddress())));
		}
	}
		protected void handleRpcMessage(Object message) {

		if (message instanceof RunAsync) {
			// 处理异步执行Runnable消息
			handleRunAsync((RunAsync) message);
		} else if (message instanceof CallAsync) {
			// 处理异步执行Callable消息
			handleCallAsync((CallAsync) message);
		} else if (message instanceof RpcInvocation) {
			// 处理Rpc调用消息
			handleRpcInvocation((RpcInvocation) message);
		} else {
			log.warn(
				"Received message of unknown type {} with value {}. Dropping this message!",
				message.getClass().getName(),
				message);

			sendErrorIfSender(new AkkaUnknownMessageException("Received unknown message " + message +
				" of type " + message.getClass().getSimpleName() + '.'));
		}
	}

3.5 AkkaInvocationHandler

Akka调用代理的处理器，熟悉Java动态代理的同学对这个实现肯定并不陌生。当前类继承自InvocationHandler，并实现其invoke方法，代码位置：org.apache.flink.runtime.rpc.akka.AkkaInvocationHandler

	@Override
	public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
		Class<?> declaringClass = method.getDeclaringClass();

		Object result;

		// 判断方法的类是否为指定的类，符合如下特定的类，执行本地调用，否则实行远程调用
		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)) {
			throw new UnsupportedOperationException("AkkaInvocationHandler does not support the call FencedRpcGateway#" +
				method.getName() + ". This indicates that you retrieved a FencedRpcGateway without specifying a " +
				"fencing token. Please use RpcService#connect(RpcService, F, Time) with F being the fencing token to " +
				"retrieve a properly FencedRpcGateway.");
		} else {
			// 执行远程调用
			result = invokeRpc(method, args);
		}

		return result;
	}
	
		/**
	 * Invokes a RPC method by sending the RPC invocation details to the rpc endpoint.
	 *
	 * @param method to call
	 * @param args of the method call
	 * @return result of the RPC
	 * @throws Exception if the RPC invocation fails
	 */
	private Object invokeRpc(Method method, Object[] args) throws Exception {
		String methodName = method.getName();
		Class<?>[] parameterTypes = method.getParameterTypes();
		Annotation[][] parameterAnnotations = method.getParameterAnnotations();
		Time futureTimeout = extractRpcTimeout(parameterAnnotations, args, timeout);

		final RpcInvocation rpcInvocation = createRpcInvocationMessage(methodName, parameterTypes, args);

		Class<?> returnType = method.getReturnType();

		final Object result;

		if (Objects.equals(returnType, Void.TYPE)) {
			// 判断是否存在返回值，如果不存在返回值，直接调用
			tell(rpcInvocation);

			result = null;
		} else {
			// execute an asynchronous call
			CompletableFuture<?> resultFuture = ask(rpcInvocation, futureTimeout);

			CompletableFuture<?> completableFuture = resultFuture.thenApply((Object o) -> {
				if (o instanceof SerializedValue) {
					try {
						return  ((SerializedValue<?>) o).deserializeValue(getClass().getClassLoader());
					} catch (IOException | ClassNotFoundException e) {
						throw new CompletionException(
							new RpcException("Could not deserialize the serialized payload of RPC method : "
								+ methodName, e));
					}
				} else {
					return o;
				}
			});

			if (Objects.equals(returnType, CompletableFuture.class)) {
				result = completableFuture;
			} else {
				try {
					result = completableFuture.get(futureTimeout.getSize(), futureTimeout.getUnit());
				} catch (ExecutionException ee) {
					throw new RpcException("Failure while obtaining synchronous RPC result.", ExceptionUtils.stripExecutionException(ee));
				}
			}
		}

		return result;
	}

4 基于Flink的RPC框架编写RPC示例

上面了解了Flink的RPC框架的实现原理，接下来我们就使用Flink的RPC框架实现一个简单RPC调用来巩固一下掌握的知识。包括如下步骤：

1. 定义Gateway
2. 实现Endpoint
3. 编写服务端示例
4. 编写客户端示例
5. 运行测试

4.1 定义Gateway

集成RpcGateway自定义接口

package com.hollysys.flink.src.rpc;

import org.apache.flink.runtime.rpc.RpcGateway;

/**
 * @author shirukai
 */
public interface DemoGateway extends RpcGateway {
    String sayHello(String name);
    String sayGoodbye(String name);
}

4.2 实现Endpoint

继承RpcEndpoint并实现DemoGateway接口

package com.hollysys.flink.src.rpc;

import org.apache.flink.runtime.rpc.RpcEndpoint;
import org.apache.flink.runtime.rpc.RpcService;

/**
 * @author shirukai
 */
public class DemoEndpoint extends RpcEndpoint implements DemoGateway {
    public DemoEndpoint(RpcService rpcService) {
        super(rpcService);
    }

    @Override
    public String sayHello(String name) {
        return "Hello," + name + ".";
    }

    @Override
    public String sayGoodbye(String name) {
        return "Goodbye," + name + ".";
    }
}

4.3 编写服务端示例

服务端有三步操作：

1. 创建RPC服务
2. 创建RpcEndpoint实例
3. 启动Endpoint

package com.hollysys.flink.src.rpc;

import akka.actor.ActorSystem;
import org.apache.flink.runtime.akka.AkkaUtils;
import org.apache.flink.runtime.rpc.akka.AkkaRpcService;
import org.apache.flink.runtime.rpc.akka.AkkaRpcServiceConfiguration;

/**
 * @author shirukai
 */
public class FlinkRpcServerExample {
    public static void main(String[] args) {
        // 1. 创建RPC服务
        ActorSystem defaultActorSystem = AkkaUtils.createDefaultActorSystem();
        AkkaRpcService akkaRpcService = new AkkaRpcService(defaultActorSystem,
                AkkaRpcServiceConfiguration.defaultConfiguration());

        // 2. 创建RpcEndpoint实例
        DemoEndpoint endpoint = new DemoEndpoint(akkaRpcService);
        System.out.println("Address: "+endpoint.getAddress());
        // 3. 启动Endpoint
        endpoint.start();
    }
}

4.4 编写客户端示例

客户端有三步操作：

1. 创建RPC服务
2. 连接远程RPC服务
3. 远程调用

package com.hollysys.flink.src.rpc;

import akka.actor.ActorSystem;
import org.apache.flink.runtime.akka.AkkaUtils;
import org.apache.flink.runtime.rpc.akka.AkkaRpcService;
import org.apache.flink.runtime.rpc.akka.AkkaRpcServiceConfiguration;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

/**
 * @author shirukai
 */
public class FlinkRpcClientExample {
    public static void main(String[] args) throws ExecutionException, InterruptedException {
        // 1. 创建RPC服务
        ActorSystem defaultActorSystem = AkkaUtils.createDefaultActorSystem();
        AkkaRpcService akkaRpcService = new AkkaRpcService(defaultActorSystem,
                AkkaRpcServiceConfiguration.defaultConfiguration());
        // 2. 连接远程RPC服务，注意：连接地址是服务端程序打印的地址
        CompletableFuture<DemoGateway> gatewayFuture = akkaRpcService
                .connect("akka.tcp://flink@169.254.1.71:63289/user/1dd321ce-2d48-4ecd-95a5-f1853d3d452b", DemoGateway.class);
        // 3. 远程调用
        DemoGateway gateway = gatewayFuture.get();
        System.out.println(gateway.sayHello("flink-rpc"));
        System.out.println(gateway.sayGoodbye("flink-rpc"));
    }
}

4.5 运行测试

首选运行服务端示例程序：

然后根据服务端输出的地址，修改客户端的连接地址，启动客户端示例程序：

5 总结

Flink的RPC实现的核心逻辑差不多就这么多吧，我也是一遍阅读一遍整理的文章，其中有一些代码并没有进行详细的阅读，如果文章有错误的地方，麻烦给位同学多多指正。