Ruby の拡張ライブラリでメモリリークを防ぐメモリ確保の方法

Ruby の拡張ライブラリを書く時には、Ruby の処理を呼び出すと例外が発生する可能性があることに気をつけないといけません。たとえば以下のように some_func という関数を呼び出す wrapper method を定義したとします。

extern int some_func(int len, int *ary);

VALUE
sample_func(VALUE self, VALUE arg)
{
    int *buf;
    int len, i, result;

    if (!RB_TYPE_P(arg, T_ARRAY))
        rb_raise(rb_eArgError, "arg must be an Array");

    len = RARRAY_LEN(arg);
    buf = ALLOC_N(int, len);
    for (i = 0; i < len; i++) {
        buf[i] = NUM2INT(RARRAY_PTR(arg)[i]);
    }
    result = some_func(len, buf);
    xfree(buf);
    return INT2NUM(result);
}

この実装にはメモリリークの可能性があります。 ALLOC_N() マクロで一時的に利用するメモリを確保していて、その後に NUM2INT() マクロを利用して Ruby のオブジェクト(VALUE 型の値)を int に変換しようとしています。 ここで arg の配列に格納されているオブジェクトが int に変換できないと例外が発生するため、メモリ解放処理(xfree) を実行しないままこの関数から抜けてしまいます。これを防ぐには rb_ensure() を使って ensure 節に相当する処理を登録して呼び出す方法、rb_protect() を使って例外など大域脱出を捕捉する方法*1などがあります。けど関数ポインタを渡すために、メインの処理を別の関数に切り出さないといけなくてちょっと面倒です。

そんな時に String オブジェクトを利用して、もし例外が起きても GC が回収してくれるから OK という発想の方法もあります。 rb_str_tmp_new() という API で String オブジェクトを生成して*2、そのバッファを利用するようにすると、仮に大域脱出が起きてもスタック上に作った VALUE の参照がなくなるため、String オブジェクトが回収と同時に解放されることが期待されます*3。 rb_str_tmp_new() をそのまま利用する時にはコンパイラの最適化による GC のマーク漏れを防ぐために RB_GC_GUARD() を適切に利用する必要もあるのですが、そのへんの面倒をうまく隠してくれる API が 1.9.3 から追加されています。

void *rb_alloc_tmp_buffer(volatile VALUE *store, long len);
void rb_free_tmp_buffer(volatile VALUE *store);

rb_alloc_tmp_buffer() で temporary buffer を確保して、その anchor である String オブジェクトを格納するポインタを store に渡します。これは通常マシンスタック上に置きたいので自動変数のポインタを渡します。 解放するところで rb_free_tmp_buffer() を呼べばバッファはすぐに解放されます。
この API を使うと先程の例はこのように書けます。([2011-10-16 追記] メモリ確保サイズが間違っていたので修正しました。なかださんありがとうございます)

extern int some_func(int len, int *ary);

VALUE
sample_func(VALUE self, VALUE arg)
{
    VALUE tmpstr;
    int *buf;
    int len, i, result;

    if (!RB_TYPE_P(arg, T_ARRAY))
        rb_raise(rb_eArgError, "arg must be an Array");

    len = RARRAY_LEN(arg);
    buf = rb_alloc_tmp_buffer(&tmpstr, (long)len * sizeof(int));
    for (i = 0; i < len; i++) {
        buf[i] = NUM2INT(RARRAY_PTR(arg)[i]);
    }
    result = some_func(len, buf);
    rb_free_tmp_buffer(&tmpstr);
    return INT2NUM(result);
}

しかし 1.9.2 にはこの API は存在しません。なので extconf.rb で関数の存在チェックを入れて、存在しなかったら自分で定義を追加するようにしておくと 1.9.2 でも使えるようになります。

extconf.rb に以下のようなチェックと create_header の呼び出しを(なければ)追加します。create_header を呼ぶと extconf.h というヘッダが生成されます。そこに HAVE_RB_ALLOC_TMP_BUFFER などのマクロが定義されます。

...
have_func("rb_alloc_tmp_buffer", "ruby.h")
have_func("rb_free_tmp_buffer", "ruby.h")

create_header

あとは missing.c などのソースファイルを用意して、そこに自前の関数定義を追加しましょう。 HAVE_RB_ALLOC_TMP_BUFFER などのマクロをチェックして条件コンパイルするようにすれば、関数が存在しない時だけ自前の実装を有効にできます。

#ifndef HAVE_RB_ALLOC_TMP_BUFFER
void *
rb_alloc_tmp_buffer(volatile VALUE *store, long len)
{
        VALUE s = rb_str_tmp_new(len);
        *store = s;
        return RSTRING_PTR(s);
}
#endif

#ifndef HAVE_RB_FREE_TMP_BUFFER
void
rb_free_tmp_buffer(volatile VALUE *store)
{
    VALUE s = *store;
    *store = 0;
        if (s) rb_str_clear(s);
}
#endif

それから利用するところでプロトタイプ宣言を追加すればよいです。ソースファイルが1つだけならそこに書けばいいですし、ヘッダがあるならそこに書くのがいいでしょう。

#include "extconf.h"
...
#ifndef HAVE_RB_ALLOC_TMP_BUFFER
extern void *rb_alloc_tmp_buffer(volatile VALUE *store, long len);
#endif

#ifndef HAVE_RB_FREE_TMP_BUFFER
extern void rb_free_tmp_buffer(volatile VALUE *store);
#endif

これで 1.9.2 でも rb_alloc_tmp_buffer() を利用してメモリリークを防ぎつつ手軽な一時メモリ確保ができますね。え、1.8? Ruby 1.8 has no future ですので 1.9 に移行しましょう。

*1:[追記] 捕捉した例外などは必要な後始末処理をしたあとで rb_jump_tag() で再発生させないといけません

*2:ちなみに rb_str_tmp_new() で生成した String オブジェクトはクラスを示す klass メンバが 0 に初期化されるため、ObjectSpace.each_object でも見えないように隠されています

*3:CRuby は保守的 GC を採用しているので、回収されないかもしれませんが