Datenlord | Rust实现RDMA异步编程（二）：async Rust 封装 UCX 通信库

UCX 是一个高性能网络通信库，它作为 MPI 所依赖的通信模块之一在高性能计算领域得到广泛的使用。UCX 使用 C 语言编写，为了在 Rust 项目中使用它，我们需要将它的 C 接口包装成 Rust 库。在这个过程中我们充分利用了 Rust 的杀手级特性—— async-await 协程来包装异步 IO 接口，从而极大降低了应用的编程复杂度。

去年我们用 Rust 实现的高性能分布式文件系统 MadFS，底层就使用了我们自己包装过的 UCX 作为通信模块，它在大规模 RDMA 网络上展现出了良好的性能。UCX 官方在得知这一消息后也很开心地宣传了我们这个项目 : )

本文首先会介绍一下 UCX 通信库的功能和编程模型，然后介绍我们用 async Rust 封装 UCX 的过程，具体代码可以参考 GitHub 仓库：async-ucx。值得注意的是，这里介绍的 IO 模型和封装异步 IO 的方法是通用的，可以适用到其它 IO 库当中。

UCX 通信接口简介

UCX 的全称是 Unified Communication X。正如它名字所展示的，UCX 旨在提供一个统一的抽象通信接口，能够适配任何通信设备，并支持各种应用的需求。

下图是 UCX 官方提供的架构图：

可以看到，UCX 整体分为两层：上层的 UCP 接口和底层的 UCT 接口。

底层的 UCT 适配了各种通信设备：从单机的共享内存，到常用的 TCP Socket，以及数据中心常用的 RDMA 协议，甚至新兴的 GPU 上的通信，都有很好的支持。

上层的 UCP 则是在 UCT 不同设备的基础上，封装了更抽象的通信接口，以方便应用使用。具体来说有以下几类：

Active Message：最底层的接口，提供类似 RPC 的语义。每条 Active Message 会触发接收端进行一些操作。
RMA / Atomic：是对远程直接内存访问（RDMA）的抽象。通信双方可以直接读写远端的内存，但是需要有额外的内存注册过程。
Tag Matching：常用于高性能计算 MPI 程序中。每条消息都会附带一个 64 位整数作为 tag，接收方每次可以指定接收哪种 tag 的消息。
Stream：对字节流（TCP）的抽象。

一般来说，和底层通信设备模型最匹配的接口具有最高的性能，其它不匹配的接口都会有一次软件转换过程。另一方面，同一种 UCP 接口发送不同大小的消息可能也会使用不同的 UCT 方法。例如在 RDMA 网络中，由于内存注册也有不小的开销，因此对于小消息来说，拷贝到预注册好的缓冲区再发送的性能更高。这些策略默认是由 UCX 自己决定的，用户也可以通过设置环境变量的方式手动修改。

在我们的系统中，使用了 UCP Tag 接口并基于此实现了轻量级的 RPC。在 RPC 场景下，Tag 可以用于区分不同上下文的消息：每个链接双方首先随机生成一个 tag 作为请求的标识，对于每次请求再随机生成一个 tag 作为回复的标识。此外 Tag 接口还支持 IO Vector，即将不连续的多个内存段合并成一个消息发送。这个特性可以用来将用户提供的数据缓冲区和 RPC 请求打包在一起，一定程度上避免数据拷贝。

UCX 编程模型简介

UCX 采用了以异步 IO 为核心的编程模型。其中 UCP 层定义的核心对象有以下四种：

Context：全局资源的上下文，管理所有通信设备。一般每个进程创建一个即可。
Worker：任务的管理调度中心，以轮询方式执行任务。一般每个线程创建一个，会映射为网卡上的一个队列。
Listener：类似 TCP Listener，用来在 worker 之间创建连接。
Endpoint：表示一个已经建立的连接。在此之上提供了各种类型的通信接口。

它们之间的所属关系如下图所示：

建立连接

UCX 中双方首先要建立连接，拿到一个 Endpoint 之后才能进行通信。建立连接一般要通过 Listener，过程和 TCP 比较类似：

通信双方 A/B 首先建立各自的 Context 和 Worker，其中一方 A 在 Worker 上创建 Listener 监听连接请求，Listener 的地址会绑定到本机的一个端口上。用户需要通过某种方法将这个地址传递给另一方 B。B 拿到地址后在 Worker 上发起 connect 操作，此时 A 会收到新连接请求，它可以选择接受或拒绝。如果接受则需要在 Worker 上 accept 这个请求，将其转换为 Endpoint。之后 B 会收到 A 的回复，connect 操作完成，返回一个 Endpoint。此后双方就可以通过这对 Endpoint 进行通信了。

内存注册

对于常规的通信接口，用户可以直接在 Endpoint 上发起请求。但对于 RMA（远程内存访问）操作，需要被访问的一方首先在自己的 Context 上注册内存，同时指定访问权限，获得一个 Mem handle。然后将这个本地 handle 转化为其他节点可以访问的一串 token，称为 remote key（rkey）。最后想办法把 rkey 传给远端。远端拿着这个 rkey 进行远程内存访问操作。

异步任务处理（重点）

为了发挥最高的性能，整个 UCX 通信接口是全异步的。所谓异步指的是 IO 操作的执行不会阻塞当前线程，一次操作的发起和完成是独立的两个步骤。如此一来 CPU 就可以同时发起很多 IO 请求，并且在它们执行的过程中可以做别的事情。

不过接下来问题来了：程序如何知道一个异步任务是否完成了？常见的有两种做法：主动轮询，被动通知。前者还是需要占用 CPU 资源，所以一般都采用通知机制。在 C 这种传统过程式语言中，异步完成的通知一般通过 回调函数（callback）实现：每次发起异步操作时，用户都需要传入一个函数指针作为参数。当任务完成时，后台的运行时框架会调用这个函