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

huacnlee · 2022年04月22日 · 最后由 jinxin238357 回复于 2022年04月30日 · 1292 次阅读

最近我发现 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

🎉 性能提升 6 - 10 倍

参考内容

目前 Rust 社区里面似乎还没有太多 Rust 集成的例子,如果也你想这么做,可以阅读以下 auto-correct 这个 Gem 这个 PR 的变更内容。

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

这 gem 发布出去,是不是需要使用方有 cargo 环境?

之前我给 gem 通过 FFI 写过 C 的扩展,性能也不赖,而且不用依赖 rust 编译环境 现在有一个疑问,这样写,那安装 gem 的人是不是也要安装 rust 编译环境才能安装 gem?

bluecoda 回复

可以发布静态库,类似 Nokogiri 那样

很不错啊

现在 gem 对 binary 的支持好像不如 python whl

看了 rust 代码,缺少抽象。。。 不想每个模块都去写 Init_xxx 这样,需要宏来解决。。。

比如 ruby_module!(xxx){ }

如果没有抽象,只是单纯把 c 代码转化为 rust,还引入不必要的依赖,坑。

attributes macro 好像可以更好的实现这套 DSL

such a ugly language

需要 登录 后方可回复, 如果你还没有账号请 注册新账号