Go 调用 Rust 动态库：FFI 零拷贝传递 Arrow 数据块

最近小本本博主我接到项颇具挑战性的工作，具体是要把原来一个业务中立的 Rust 模块封装为动态库提供给其他团队使用，需求是能通过 Arrow 格式交换数据，同时还要提供 ETL 中常见的一些 T(ransformation) 操作。目前有 Go 服务需要接入，所以还需要由我提供 Go SDK。本篇记录一下我作为新人的一些探索学习过程。

刚接到需求有点懵，咱 Rust 仅仅才入门一周，Go 也才刚去做了几周的业务开发，这回可是要我做这俩的缝合怪啊😅

Snipaste_2023-08-07_01-23-35 我真的会谢

0. 何为 FFI

蒙蔽完了先理理思路吧。首先这俩都是静态编译的编程语言，根据我以前瞎捣鼓电脑得出的野鸡经验，任何语言编译出来的二进制产物，本质上跟 C 语言编译出来的没区别，这是因为现存的流行操作系统几乎都用 C 构造而来，所以内存管理、系统调用等等基本操作都是以 C 为基础。自然地，动态库之间的调用也能通过基础的 C 标准来做到。简单总结下就是我们把原来的 Rust 库包装做成一个 C 库，再由 Go 去调用这个 C 库，就完成了跨语言调用，也就是所谓的 FFI，外部函数接口，全称 Foreign Function Interface。

1. 内存管理

细节是魔鬼，本篇就从入门的内存管理说起。正常用单一种图灵完备语言编译构建出来的程序，比如 go 有个非常强大的 runtime，自带垃圾回收，自动管理内存，使用者基本无需操心。可一旦涉及 FFI 就不一样了，所有经过了 FFI 边界的指针，go 虚拟机无从确信它是否可以更改是否可以回收，甚至连数据结构可能都不确定，只能完全依赖用户在编译时给的头文件，事先声明好所有穿越 FFI 边界的行为。

嗯那也还好吧，然而接着内存所有权问题就来了，一块内存由谁申请、管理、回收呢？没有规划随意读写的话可是会 segfault 漫天飞哦。具体操作起来各种组合有七八种做法，不过基本原则是谁申请谁就有拥有权，就由谁负责回收，避免七手八脚写坏内容，关键是防止多次 free 或者 free 后再次使用之类的严重问题。常用的几种管理方式可以参考这篇博文，里面通过跨 FFI 字符串的例子展示了多种策略，包括提供回调，预先申请再写入等等。这里采用相对最简单的方式：我在 Rust 库同时提供对象创建和对象回收的方法，由 go 端的 caller 按需调用。

2. 举个简单栗子🌰

下面举个简单例子：根据 config 配置创建一个 reader，调用 reader 读取 arrow 数据，并通过 ffi 导出，使用完毕后回收 reader。

2.1 Rust 端

首先创建 Rust 端的 FFI 鞘层包，如果 rust 库本身够简单的话可以考虑下 cbindgen 之类的方案，帮你自动生成 C FFI 的鞘层和头文件。这里创建一个 libexample 动态库包，Cargo.toml：

[package]
name = "example"
version = "0.0.1"
edition = "2021"

[lib]
name = "example"
# 声明为外部动态库
crate-type = ["cdylib"]

[dependency]
# 这里把真正的 rust 功能模块引入进来
rustlib = "..."

创建 `lib.rs` 入口，点击展开代码和说明

use std::ffi::CStr;
use arrow::array::StructArray;
use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
use arrow::record_batch::RecordBatch;

use rustlib::prelude::new_reader;

/// `#[repr(C)]` 表示让编译器把 Reader 类型编译为兼容 C 内存对齐布局的 struct
#[repr(C)]
pub struct Reader {
    ...
}

/// `#[no_mangle]` 表示让 rustc 不要改写函数名，不然后续链接器找不到函数会无法调用
/// 这里为创建 reader 的函数，返回 reader 的指针
#[no_mangle]
pub extern "C" fn create_reader(config: *const libc::c_char) -> *const Reader {
    // 读取来自 c 的 char[] 字串
    let config_str: &str = unsafe { CStr::from_ptr(config as *mut _) }.to_str().unwarp();
    // 创建 reader，并通过 Box 将它导出为 ffi 指针，相当于把这个对象的内存故意泄漏出去，详情参考 rust 文档关于 `Box::into_raw` 的说明
    let reader: Reader = new_reader(config_str);
    Box::into_raw(Box::new(reader)) as *const Reader
}

/// 这里是回收 reader 的函数，需要调用方用完 reader 后再调这个来回收内存
#[no_mangle]
pub extern "C" fn free_reader(reader_ptr: *const Reader) {
    let _ = unsafe { Box::<Reader>::from_raw(reader_ptr as *mut _) };
}

/// 定义返回结果的 struct
#[repr(C)]
pub struct Batch {
    schema: *const FFI_ArrowSchema,
    array: *const FFI_ArrowArray,
}

/// 使用 reader 读取文件并导出为 arrow 数组数据，返回结果数据的指针
#[no_mangle]
pub extern "C" fn read_as_arrow(reader_ptr: *const Reader, path: *const libc::c_char) -> *const Batch {
    // 只获取 reader 的引用而不回收 reader，调用方可能还没用完
    let reader: &Reader = unsafe { &*(reader_ptr as *mut _) };
    let path_str = unsafe { CStr::from_ptr(path as *mut _) }.to_str().unwrap();

    // 读取并导出数据
    let batch: &RecordBatch = reader.arrow_from_path(path_str);
    let table_struct: StructArray = batch.into();
    let schema = FFI_ArrowSchema::try_from(table_struct.data_type()).unwrap();
    let array = FFI_ArrowArray::new(table_struct.data());
    Box::into_raw(Box::new(Batch {
        schema: Box::into_raw(Box::new(schema)),
        array: Box::into_raw(Box::new(array)),
    })) as *const Batch
}

这里有个特殊的地方，返回的 Batch 结果我们无须再提供对应的回收方法。rust arrow-rs 库中两个类 FFI_ArrowSchema 和 FFI_ArrowArray 是专门用来穿越 FFI 界面的（这俩也是 arrow 标准里定义要求的内容），这两种对象传送出 FFI 界面被调用端导入后，就由调用方取得所有权，由调用方的 arrow 库自行回收，不需要我们在创建端操心回收了。

2.2 header 头文件

在 C 头文件上声明刚刚导出的类型和函数，创建 example.h：

#include<stddef.h>
typedef struct Reader {
...
} Reader;

Reader* create_reader(char[] config);
void free_reader(Reader* reader_ptr);

#define FFI_ArrowSchema_ptr uintptr_t
#define FFI_ArrowArray_ptr uintptr_t
typedef struct Batch {
    FFI_ArrowSchema_ptr schema;
    FFI_ArrowArray_ptr array;
} Batch;

Batch* read_as_arrow(Reader* reader_ptr, char[] path);

2.3 Go 端

接着在 Go 端我们也需要一个 C 库的鞘层，命名为 libexample-go。go.mod：

module github.com/examplellc/libexample-go

go 1.21.0

require (...)

接着新建 reader.go，创建一个 Reader 包装类型。

新建 `reader.go` 如下，点击查看代码

package main

/*
#cgo LDFLAGS: -lexample -L.
#include <stdlib.h>
#include <stddef.h>
#include "example.h"
*/
import "C"
import (
	"github.com/apache/arrow/go/v12/arrow"
	"github.com/apache/arrow/go/v12/arrow/cdata"
	"unsafe"
)

type Reader struct {
	readerPtr *C.struct_Reader
}

func NewReader(config string) Reader {
	return Reader{
		readerPtr: C.create_reader(C.CString(config)),
	}
}

func (r *Reader) Close() {
	if r.readerPtr != nil {
		C.free_reader(r.readerPtr)
	}
	r.readerPtr = nil
}

func (r *Reader) ReadArrow(path string) arrow.Record {
	batchPtr := C.read_as_arrow(r.readerPtr, C.CString(path))
	return cdata.ImportCRecordBatch(
		(*cdata.CArrowArray)(unsafe.Pointer(batchPtr.array)),
		(*cdata.CArrowSchema)(unsafe.Pointer(batchPtr.schema)),
	)
}

有一些值得注意的地方，

首先这个包是个 cgo 包，需要调用对应平台的编译器和连接器，开发环境上需要额外安装；
然后上面源码中引用了一个 go 内置的 "C" 假包，紧挨着它 import 前面的注释有特殊作用，要用 C 的语法写需要导入的头，里面用 #cgo 开头的行可以定义编译参数。比如这里定义了编译后需要链接的库名称，以及搜索路径：
```
#cgo LDFLAGS: -lexample -L.
```
需要引用头文件内 struct 类型时，可以用 C.struct_ 开头的格式引用，注意匿名 struct 不能这样引用，必须像 typedef struct {} xxx; 这样导出后才可用 C.struct_xxx；
要引用头文件内声明的函数，使用 C.read_as_arrow 这样的格式即可；
假包 "C" 还包含了一些工具函数，具体可参考 cgo 官方文档，像这里用到的 C.CString 能把 go string 转成 C 中以 null 结尾的 char[] 字串；

至此所有包装层的包都制作完毕了。

2.4 使用方法

接下来在正常 go 项目中如何调用上面这个库呢？

首先 cargo build --release 编译出 c 库本体 libexample.so/libexample.dylib/libexample.dll；
把以上编译产物拷贝到包的根目录，因为我们刚指定了链接库的搜索路径 -L.，当然也可以放到常用的 /usr/local/lib 等目录下；
最后 go 编译时打开 cgo 开关，CGO_ENABLED=1 go build . -o build/main 像这样用环境变量 CGO_ENABLED 即可；

不出意外，build/main 已经可以正常执行了。

3. Debug

上面开个玩笑哈哈，第一次运行不出意外那是不可能的😝程序起来铁定少不了 segfaults，还有各种奇奇怪怪的数据错乱，建议多打断点多调试，善用 gdb，还有 go delv 和 rust lldb，也是不错的官方调试工具。