Rust FFI 实践

背景

部门算法团队开始成长起来，开始有越来越多的尝试以及成果，但是现在工程方面严重的限制了（主要是做预测服务）他们的研究成果转化为实际输出的能力。去年下半年，我们就发现TF官方的Java binding 存在严重的内存泄露问题，而TF Java binding 因为封装包括训练和预测所需要的API,比较复杂，我们也难以更改。同时，使用TF serving，就需要提供标准的RPC调用来完成交互，而所有的数据处理等工作都是在Java端，这也对运维模式产生一定的压力，毕竟要维护serving集群，研发工程师要对接serving才能完成一个端到端的预测。

其实最简单的办法还是提供java binding,但是这个binding我们只提供Load model 以及tensor in tensor out 的predict,这样对我们的封装要求就可以小很多。而对于C/C++我个人并不是很熟悉，正好我最近开始学习Rust(主要还是看到Rust更符合我的编程习惯)，所以打算用Rust对TF进行二次封装，然后暴露出一个简单的C ABI供Java做绑定。 这里就会涉及到Rust FFI的使用，目前网络上资源也比较少，更多的是Example性质的，大家的文章大同小异，所以我这里就简单写下我这两天折腾的心得。代码还比较烂，欢迎大家指正，但勿喷。

Rust 和 C 交互的基础

语言之间能够交互的核心原因在于最终他们都会被编译为基于特定系统（如Linux）二进制文件，这种底层的共通性就为他们带来了直接交互的可能性。现阶段，如果要跨多语言，最好的方式是都遵循C ABI标准。

业务逻辑

这里业务逻辑比较简单，根据流程，我们只要提供五个核心函数即可，分别是：

//创建张量
CTensor *create_tensor(float *data, int32_t data_length, int32_t *shape, int32_t shape_length);

//加载模型
Predictor_t *load(char exported_dir[]);

//预测结果
OutputTensor *predict(Predictor_t *predictor, char *output_name[], char *input_names[], CTensorArray *input_values);

//获取结果
CTensor *get_ctensor(OutputTensor *outputTensor);

//释放结果张量
void *free_tensor(CTensor *ctensor);

因为我们是希望用Rust对TF做二次封装，这里我们直接使用crate rust-tf, 这样可以避免再次对TF原生的libtensorlfow做封装，毕竟工作量是不小的。

透明指针

前面的定义涉及到了 CTensor, OutputTensor, Predictor_t 等几个结构体，他们本质上对应的都是高级语言里的对象。因为Rust 支持和C一样的结构体布局，所以我们可以在两个语言之间直接传递结构体。下面是这几个结构体在C侧的定义：

typedef struct Predictor_S Predictor_t;
typedef struct OutputTensor OutputTensor;

typedef struct CTensor {
    const float *data;
    int32_t data_length;
    const int *shape;
    int32_t shape_length;
} CTensor;

typedef struct CTensorArray {
    CTensor *data;
    int32_t len;
} CTensorArray;

Predictor_t, OutputTensor 我们看只是定义了一个空结构体，因为这两个东西对应的是Rust返回的对象。在FFI里，我们可以使用一个空的struct 对象来代替一个实际的Rust对象，然后通过指针来进行应用。什么意思的呢？比如下面代码：

Predictor_t *predictor = load("model path");

这里的load其实对应的是rust里的函数，该函数返回了一个包装了TFSession的对象。我们通过Predictor_t来表示它，predictor虽然是指向Predictor_t的指针，但本质上是指向包装TFSession对象的指针。虽然在C里我们不能直接调用Predictor_t的方法，但是我们可以提供一些辅助方法将Predictor_t传递回Rust然后在Rust调用里调用他的方法，最后返回结果。比如，

Predictor_t *pre = load(path);
OutputTensor *wow = predict(pre, "y_hat", "x,y", tarray);

上面的load/predict 都是Rust提供的C方法，虽然我们没办法直接调用Predictor_t的predict方法，但是我们再提供一个C的predict方法，然后将Predictor_t传递回Rust里，调用该对象的实际predict方法。

下面是是predict方法的签名以及在Rust里的实际上实现：

OutputTensor *predict(Predictor_t *predictor, char *output_name[], char *input_names[], CTensorArray *input_values);

#[no_mangle]
pub extern "C" fn predict(predictor: *mut Predictor, output_name: FfiStr, input_names: FfiStr, input_values: *mut CTensorArray) -> *mut Tensor<f32> {
    let r_predictor = unsafe {
        assert!(!predictor.is_null());
        *Box::from_raw(predictor)
    };
    ......  
    let tensor = r_predictor.predict(r_output_name, r_input_names, r_input_values_with_mut_ref);

    Box::into_raw(Box::new(tensor))
}

所以透明指针只是对一个对象的实际引用，方便我们在调用语言侧传递。同样的，我们也可以看到，两种语言直接的交互，都可以通过指针来完成。

如何在C/Rust之间传递指针

首先，Rust 的函数要返回一个指针，可以像下面那么做：

#[no_mangle]
pub extern "C" fn create_tensor(data: *const c_float,
                                data_length: c_int,
                                shape: *const c_int,
                                shape_length: c_int, ) -> *mut CTensor {
  let ctensor = CTensor {
        data,
        data_length,
        shape,
        shape_length,
    };
    Box::into_raw(Box::new(ctensor))
}

接着我们在C语言里申明函数头，就能使用Rust里的这个函数实现了：

CTensor *create_tensor(float *data, int32_t data_length, int32_t *shape, int32_t shape_length);

也就是任何复杂对象都可以通过Box::into_raw(Box::new(....))来完成。

那如何传递一个指针给Rust函数呢？ 我们看下面的代码：

#[no_mangle]
pub extern "C" fn get_ctensor(tensor: *mut OutputTensor) -> *mut CTensor {
    assert!(!tensor.is_null());
    let r_tensor = unsafe { *Box::from_raw(tensor) };
    let ctensor = r_tensor.to_ctensor();
    Box::into_raw(Box::new(ctensor))
}

获取一个C传递回来的指针后，需要对他进行处理之后才能在Rust里面使用：

let r_tensor = unsafe { *Box::from_raw(tensor) };

现在总结下，C传递来的指针，需要使用*Box::from_raw(...),而如果想把Rust对象作为指针传递出去，则需要做Box::into_raw(Box::new(....))。

如果我想传递数组怎么办？

数组使用太频繁了，那么C/Rust 应该如何传递数组呢？本质上我们是没办法直接传递数组的，除了普通的值类型，一切都是以指针进行交互的。在C里面，数组和指针具有很大的相关性，我们可以利用指针来模拟数组。我们以字符串为例子，因为对于字符串，不同语言的表示形态也是不一样的，但是都可以用char(u8)来表示，所以我们可以把字符串看成u8的数组。一个数组其实由两部分组成：

1. 一片连续的元素（内存）
2. 元素的个数

我们只要知道这片连续元素的起始地址，以及元素的个数，就能描述这个数组，所以通过下面的struct 就能描述数组：

typedef struct cstring_t {
    const char *data;
    int32_t data_length;    
} cstring_t;

这样，data是指向char的指针，同时我们也可以认为是数组第一个元素的指针，然后我们提供了该指针指向数组的长度。

接着我们在Rust定义各一个相同的结构体,并且提供一个函数供C调动。

#[repr(C)]
pub struct cstring_t {
    data: *const u8,
    data_length: c_int,
}
imp

#[no_mangle]
pub extern "C" fn pass_str(cst: *const cstring_t) {
    let r_cst = unsafe { *Box::from_raw(cst as *mut cstring_t) };
    let s = unsafe {
        slice::from_raw_parts(r_cst.data, r_cst.data_length as usize)
    };
    println!("{:?}", std::str::from_utf8(s))
}

现在，提供一个C的函数签名：

void *pass_str(cstring_t *csr);

现在就可以在C侧调用了：

char *ye = "abc";
    int len = 3;

    cstring_t a;

    cstring_t *a_p = malloc(sizeof(a));
    a_p->data = ye;
    a_p->data_length = len;
    pass_str(a_p);

然后大家就看到了Rust侧打印了如下内容：

Ok("abc")

知道了字符串是怎么处理的，那么我们想传递一个张量该怎们办呢？张量本质上由两部分组成：

1. 一个存储实际数据的数组（一维）
2. 描述形状的数组 （一维）

所以其实是两个数组，前面我们知道，描述一个数组只要一个指针和一个长度就可以了，所以我们描述一个张量可以这么做：

typedef struct CTensor {
    const float *data;
    int32_t data_length;
    const int *shape;
    int32_t shape_length;
} CTensor;

对应rust的结构为：

#[repr(C)]
pub struct CTensor {
    data: *const c_float,
    data_length: c_int,
    shape: *const c_int,
    shape_length: c_int,
}

这样就实现了在C/Rust之间实现了张量的交换。

更复杂的数组传递

前面我们看到，数组里面还都是一些基本类型，那如果数组里面是个对象怎么办？比如，我希望提供一个张量数组，其实没有什么差别，申明大概是这样的：

#[repr(C)]
pub struct CTensorArray {
    data: *const *const CTensor,
    len: c_int,
}

#[no_mangle]
pub extern "C" fn create_tensor_array(data: *const *const CTensor,
                                      len: c_int) -> *mut CTensorArray {
    assert!(!data.is_null());
    let tensor_array = CTensorArray {
        data,
        len,
    };
    Box::into_raw(Box::new(tensor_array))
}

只不过除了基础类型以外，一切都是要以指针传递，所以这里的数据data是一个指针，这个指针指向CTensor的指针。使用起来大概是这样的：

CTensor *xTensor = create_tensor(xP, 1, shape_x_p, 1);
CTensor *yTensor = create_tensor(yP, 1, shape_y_p, 1);
CTensor *xy[] = {xTensor, yTensor};
CTensor **xy_p;
xy_p = xy;

CTensorArray *tarray = create_tensor_array(xy_p, 2);

由此可见，你是可以传递任意复杂的东西的，不过代价也比较高。

所有权在Rust/C之间的转移

我们知道Rust是一门内存安全的问题，响应的有所有权和申明周期的问题。所以在做跨语言交互的过程会遇到一些相关的问题。

首先，一个对象如果传递给了调用者，那么所有权会转移到调用者，这个是由

Box::into_raw(Box::new(....))

自动完成的。

其次，借出的对象一旦重新返回Rust,那么所有权就转移回Rust了，这个也是由

*Box::from_raw(...)

自动完成的。

所以，下面代码大家发现什么问题了么？

//load predict  都是rust实现的方法
    Predictor_t *pre = load(path);
    OutputTensor *wow = predict(pre, "y_hat", "x,y", tarray);
    OutputTensor *wow2 = predict(pre, "y_hat", "x,y", tarray);

load 会将rust里的Predictor_t所有权转移给C,这个时候pre持有所有权。接着第一次调用predict,在predict方法里面我们的调用了

*Box::from_raw(pre)

重新获取了所有权，那么这个时候调用者(也就是前面的 Predictor_t *pre)指向的pre 就成了野指针了。接着在第二次使用的时候，就会出现错误。同样的tarray也会自动被释放，无法使用两次。

其实我们希望pre能够完全由调用者来决定是否释放，有解决办法么？其实本质在于from_raw会获取所有权，所以我们只要不使用他就行，使用如下方式：

&*tensor

这里面我们只是简单的解应用pre然后获取地址，避免去获取所有权。 不过所有权虽然带来麻烦，但是同时也能简化内存释放的问题。

所有权导致的另外一个空指针问题

让我们按一段代码：

#[no_mangle]
pub extern "C" fn to_tensor(data: *const Tensor<f32>) -> *mut CTensor {
    let tensor = unsafe {
        *Box::from_raw(data as *mut Tensor<f32>)
    };
    let tensor_ref = &tensor;
    let ctensor = CTensor::from(tensor_ref);
    Box::into_raw(Box::new(ctensor))
}

这段代码接受了tensor,返回ctensor。 在前面，我们获取了tensor的使用权，接着我们够着CTensor的时候使用了tensor的引用，然后我们返回了ctensor. 但是返回之后，tensor就被释放掉了，导致ctensor对tensor的引用成了野指针。

很多场景下，我们确实需要一个包装对象，那怎么解决这个问题呢？我一开始想到的是不释放tensor就可以了：

// here tensor ownership have be moved and
    // we can not use the pointer to get tensor again.
    std::mem::forget(tensor);

但这样会导致内存泄露，因为我们没有其他地方可以调用tensor并且进行释放。而且forget tensor,其实是将tensor的所有权转给了ManuallyDrop 对象。

其实比较好的办法是不获取tensor的所有权，

let tensor = unsafe {
        &* （data data as *mut Tensor<f32>）
    };

缺点也就来了，调用方需要去释放tensor。

还有第三个办法就是提供一个对象，该对象有一个to_ctensor方法，我们在to_tensor里调用这个对象的to_ctensor方法：

#[repr(C)]
pub struct OutputTensor {
    tensor: Tensor<f32>,
}
impl OutputTensor {
    fn to_ctensor(&self) -> CTensor {
        println!("to_tensor dims:{:?}  data:{:?}", self.tensor.dims(), self.tensor.to_vec());
        let ctensor = CTensor::from(&self.tensor);
        ctensor
    }
}

#[no_mangle]
pub extern "C" fn get_ctensor(tensor: *mut OutputTensor) -> *mut CTensor {
    assert!(!tensor.is_null());
    let r_tensor = unsafe { *Box::from_raw(tensor) };
    let ctensor = r_tensor.to_ctensor();
    Box::into_raw(Box::new(ctensor))
}

这样在C端可以这么使用了：

OutputTensor *wow = predict(pre, "y_hat", "x,y", tarray);
CTensor *res = get_ctensor(wow);

其原理也很简单，把tensor的所有权绑定到了OutputTensor身上。

总结

跨语言交互本身是比较难的，尤其是指针问题，这也是为什么C/C++更容易写出不安全的代码。我们应该尽量使用Rust Safe部分来完成我们的逻辑。