Rust RDMA 编程「二、接口绑定」

2024.05.24
SF-Zhou

1. libibverbs

libibverbs 是一个 API 库,使应用程序可以在用户态访问 InfiniBand 硬件。它的源码可以在 rdma-core 里找到,对外提供的是 C 接口。Linux 系统一般不需要下载源码自行编译,例如 Debian 系可以通过 sudo apt install libibverbs-dev 直接安装,安装完成后头文件默认在 /usr/include/infiniband/verbs.h

Rust 无法直接使用 C 的头文件,所以需要对 C 暴露出来的结构体和函数进行绑定,一般可以使用 bindgen 帮助生成绑定接口。目前已经有 rdma-sysrust-ibverbs 做了类似的工作,但为了更好的实现我们的需求,这里仍然自行封装一个库实现绑定。

2. 绑定

生成绑定的 build.rs 内容如下:

use std::collections::HashSet;
use std::env;
use std::path::PathBuf;

fn main() {
    // 0. 查找 libibverbs,定位头文件的位置
    let lib = pkg_config::Config::new()
        .statik(false)
        .probe("libibverbs")
        .unwrap_or_else(|_| panic!("please install libibverbs-dev"));

    let mut include_paths = lib.include_paths.into_iter().collect::<HashSet<_>>();
    include_paths.insert(PathBuf::from("/usr/include"));

    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
    let gen_path = out_path.join("verbs_inline.c");
    let obj_path = out_path.join("verbs_inline.o");
    let lib_path = out_path.join("libverbs_inline.a");

    // 1. 根据头文件生成绑定代码
    let mut builder = bindgen::Builder::default()
        .clang_args(include_paths.iter().map(|p| format!("-I{:?}", p)))
        .header_contents("header.h", "#include <infiniband/verbs.h>")
        .derive_copy(true)
        .derive_debug(true)
        .derive_default(true)
        .generate_comments(false)
        .generate_inline_functions(true)
        .wrap_static_fns(true) // 生成 static 函数的绑定
        .wrap_static_fns_path(&gen_path) // 修改生成的位置
        .prepend_enum_name(false)
        .formatter(bindgen::Formatter::Rustfmt)
        .size_t_is_usize(true)
        .translate_enum_integer_types(true)
        .layout_tests(true)
        .default_enum_style(bindgen::EnumVariation::Rust {
            non_exhaustive: false,
        })
        .opaque_type("pthread_.*")
        .blocklist_type("timespec")
        .allowlist_function("ibv_.*")
        .allowlist_type("ibv_.*")
        .allowlist_type("ib_uverbs_access_flags")
        .bitfield_enum("ib_uverbs_access_flags")
        .bitfield_enum("ibv_.*_bits")
        .bitfield_enum("ibv_.*_caps")
        .bitfield_enum("ibv_.*_flags")
        .bitfield_enum("ibv_.*_mask")
        .bitfield_enum("ibv_pci_atomic_op_size")
        .bitfield_enum("ibv_port_cap_flags2")
        .bitfield_enum("ibv_rx_hash_fields");

    // 这几个 struct 中有 mutex,无法启用 copy 和 debug
    for name in [
        "ibv_srq",
        "ibv_wq",
        "ibv_qp",
        "ibv_cq",
        "ibv_cq_ex",
        "ibv_context",
    ] {
        builder = builder.no_copy(name).no_debug(name)
    }

    let bindings = builder.generate().expect("Unable to generate bindings");

    // 2. 使用 clang 编译 static 函数生成的 C wrapper 文件
    let clang_output = std::process::Command::new("clang")
        .arg("-O2")
        .arg("-c")
        .arg("-o")
        .arg(&obj_path)
        .arg(&gen_path)
        .args(["-include", "infiniband/verbs.h"])
        .output()
        .unwrap();
    if !clang_output.status.success() {
        panic!(
            "Could not compile object file: {}",
            String::from_utf8_lossy(&clang_output.stderr)
        );
    }

    // 3. 打包 .o 文件成 .a
    let lib_output = std::process::Command::new("ar")
        .arg("rcs")
        .arg(&lib_path)
        .arg(&obj_path)
        .output()
        .unwrap();
    if !lib_output.status.success() {
        panic!(
            "Could not emit library file: {}",
            String::from_utf8_lossy(&lib_output.stderr)
        );
    }

    println!("cargo:rustc-link-lib=static=verbs_inline");
    println!("cargo:rustc-link-search=native={}", out_path.display());

    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}

上方代码有很大一部分是为了处理头文件中的 static 函数。在启用 bindgen experimental feature 的情况下,可以生成 static 函数 wrapper 的 C 文件,内容大概如下:

// Static wrappers

void ibv_wr_atomic_cmp_swp__extern(struct ibv_qp_ex *qp, uint32_t rkey, uint64_t remote_addr, uint64_t compare, uint64_t swap) { ibv_wr_atomic_cmp_swp(qp, rkey, remote_addr, compare, swap); }
void ibv_wr_atomic_fetch_add__extern(struct ibv_qp_ex *qp, uint32_t rkey, uint64_t remote_addr, uint64_t add) { ibv_wr_atomic_fetch_add(qp, rkey, remote_addr, add); }
void ibv_wr_bind_mw__extern(struct ibv_qp_ex *qp, struct ibv_mw *mw, uint32_t rkey, const struct ibv_mw_bind_info *bind_info) { ibv_wr_bind_mw(qp, mw, rkey, bind_info); }
void ibv_wr_local_inv__extern(struct ibv_qp_ex *qp, uint32_t invalidate_rkey) { ibv_wr_local_inv(qp, invalidate_rkey); }
void ibv_wr_rdma_read__extern(struct ibv_qp_ex *qp, uint32_t rkey, uint64_t remote_addr) { ibv_wr_rdma_read(qp, rkey, remote_addr); }
void ibv_wr_rdma_write__extern(struct ibv_qp_ex *qp, uint32_t rkey, uint64_t remote_addr) { ibv_wr_rdma_write(qp, rkey, remote_addr); }
void ibv_wr_flush__extern(struct ibv_qp_ex *qp, uint32_t rkey, uint64_t remote_addr, size_t len, uint8_t type, uint8_t level) { ibv_wr_flush(qp, rkey, remote_addr, len, type, level); }
void ibv_wr_rdma_write_imm__extern(struct ibv_qp_ex *qp, uint32_t rkey, uint64_t remote_addr, __be32 imm_data) { ibv_wr_rdma_write_imm(qp, rkey, remote_addr, imm_data); }
// ...

实际上就是生成一批新的函数,内部调用这些 static 的函数,绑定到原先的函数名。生成的 bindings.rs 中大概是这样的:

extern "C" {
    #[link_name = "ibv_wr_atomic_cmp_swp__extern"]
    pub fn ibv_wr_atomic_cmp_swp(
        qp: *mut ibv_qp_ex,
        rkey: u32,
        remote_addr: u64,
        compare: u64,
        swap: u64,
    );
}

显然,static 函数这样处理后性能肯定会下降,毕竟多一次函数调用。rdma-sys 中对 static 函数和宏通过手工重写成 Rust 代码的方式实现替换,是更好的解决方案。计划以后有时间的时候将所有使用到的 static 函数重写为 Rust 代码。

本文中提到的代码可以在 r2dma-sys 中找到。