Rust Ruby on Rust - 用 Rust 编写 RubyGem Extension 提升 600% 的性能

huacnlee · 2022年04月22日 · 最后由 ericguo 回复于 2022年10月16日 · 1951 次阅读

最近我发现 RubyGems 以及合并了 Cargo builder 支持的 PR，也就意味着，我们除了可以用 C 来写 RubyGem 扩展之外，也可以用 Rust 了。

https://github.com/rubygems/rubygems/pull/5175

为了吃螃蟹，我用 autocorrect 的项目测试了一下，已经完成集成上 CI Result 以及跑同，集成好以后，我一次性删除了大量之前的 Rust 实现的代码。

目前在折腾 Cross Compile 打包发布的事情。

详细内容，可以作为参考：

https://github.com/huacnlee/auto-correct/pull/17

如何编写 Rust 扩展

在 RubyGems 的 PR 里面给到的例子实现方法比较初级。我实际走下来看，这样还不够。我们需要原来 Gem 的 extconf.rb 那套流程来自动化发布。

创建类似这样的目录结构：

├── Gemfile
├── Rakefile
├── auto-correct.gemspec
├── ext
│   └── autocorrect
│       ├── Cargo.lock
│       ├── Cargo.toml
│       ├── extconf.rb
│       └── src
│           └── lib.rs
├── lib
│   ├── auto-correct
│   │   └── version.rb
│   └── auto-correct.rb

Cargo.toml 需要引入 rb-sys，我们需要用它来定义 Ruby Module, Function, Const 之类的东西，这个类似 C Extension 里面相同的方法。

[package]
edition = "2021"
name = "autocorrect"
version = "1.5.11"

[lib]
crate-type = ["cdylib"]

[dependencies]
rb-sys = {version = "0.8.0", features = ["link-ruby", "ruby-static"]}
autocorrect = "1.5.11"

src/lib.rs

extern crate rb_sys;

use rb_sys::{
    rb_define_module, rb_define_module_function, rb_string_value_cstr, rb_utf8_str_new, VALUE,
};
use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_long};

#[inline]
unsafe fn cstr_to_string(str: *const c_char) -> String {
    CStr::from_ptr(str).to_string_lossy().into_owned()
}

#[no_mangle]
unsafe extern "C" fn pub_format(_klass: VALUE, mut input: VALUE) -> VALUE {
    // 需要做类型转换，将 CString 转换成 Rust String
    let ruby_string = cstr_to_string(rb_string_value_cstr(&mut input));

    // 调用 autocorrect (Rust 那个 crate) 的 format 方法（名字有点混淆注意区分）
    // https://rubygems.org/gems/autocorrect
    let result = autocorrect::format(&ruby_string);
    let size = result.len() as c_long;

    // 再把 Rust String 类型转换成 CString
    let result_cstring = CString::new(result).unwrap();

    rb_utf8_str_new(result_cstring.as_ptr(), size)
}

// 这个是 Ruby Extension 的标准流程，你需要定义一个 `Init_xxxx` 的初始化方法。
#[allow(non_snake_case)]
#[no_mangle]
pub extern "C" fn Init_autocorrect() {
    let name = CString::new("AutoCorrect").unwrap();

    let klass = unsafe { rb_define_module(name.as_ptr()) };
    let fn_format_name = CString::new("format").unwrap();

    let format_callback = unsafe {
        std::mem::transmute::<
            unsafe extern "C" fn(VALUE, VALUE) -> VALUE,
            unsafe extern "C" fn() -> VALUE,
        >(pub_format)
    };

    unsafe {
        // 定义 AutoCorrect.format 方法
        rb_define_module_function(klass, fn_format_name.as_ptr(), Some(format_callback), 1);
    }
}

定义 extconf.rb，本身 Ruby 有提供 mkmf，但我们额外需要 rb_sys 提供的改进，所以项目需要增加依赖：

s.add_runtime_dependency "rb_sys"

NOTE: 截止目前 rb_sys v0.1.0 编译还有 Bug，我已经提 PR 改进了，等发布新版本。所以我那个 PR 里面用的是我自己临时发布的 rb_sys1。

require "mkmf"
require "rb_sys/mkmf"

create_rust_makefile("auto-correct/autocorrect")

Rakefile 改动

require "rubygems"
require "bundler/setup"
require "bundler/gem_tasks"
require "rake/testtask"
require "rake/extensiontask"
require "bundler"

Rake::ExtensionTask.new("autocorrect") do |ext|
  ext.lib_dir = "lib/auto-correct"
  ext.source_pattern = "*.{rs,toml}"
end

这样你就有 rake compile 命令了，执行它，会将 Rust 代码编译成一个 autocorrect.bundle 的文件，放到 lib/auto-correct 目录，于是我们就可以在 Ruby 里面调用 Rust 函数了。

最后效果演示和意义

之前 Ruby 实现的相同逻辑与 Rust 实现的版本基本上是完全一致的（正则、规则、处理流程基本完全），甚至在 Rust 实现的 AutoCorrect 版本里面后面我又增加了更多的细节。

Rust 实现的性能结果

Warming up --------------------------------------
     format 50 chars     1.886k i/100ms
    format 100 chars     1.060k i/100ms
    format 400 chars   342.000  i/100ms
         format_html    85.000  i/100ms
Calculating -------------------------------------
     format 50 chars     18.842k (± 1.5%) i/s -     94.300k in   5.005815s
    format 100 chars     10.357k (± 1.8%) i/s -     51.940k in   5.016770s
    format 400 chars      3.336k (± 2.2%) i/s -     16.758k in   5.026230s
         format_html    839.761  (± 2.1%) i/s -      4.250k in   5.063225s

用上面方式集成 Rust 实现以后

Warming up --------------------------------------
     format 50 chars    12.376k i/100ms
    format 100 chars     6.918k i/100ms
    format 400 chars     1.944k i/100ms
         format_html   617.000  i/100ms
Calculating -------------------------------------
     format 50 chars    120.884k (± 4.8%) i/s -    606.424k in   5.031500s
    format 100 chars     67.903k (± 3.9%) i/s -    345.900k in   5.102803s
    format 400 chars     19.024k (± 2.9%) i/s -     95.256k in   5.011560s
         format_html      6.004k (± 2.3%) i/s -     30.233k in   5.038348s