Rubyのバッチを直してメモリ1/1000、速度4倍にした話

弊社、GoやScala、Node.js、Pythonに比べてRubyの地位があまり高くありません。そんな中メモリを20GBも使い、2時間もかかるRubyのバッチがあって、やっぱRubyダメじゃんと思われていたりします。

いやいやそれRubyのせいじゃないでしょう、と。私もRubyプログラマの端くれとして、こんな事でRubyの評判を下げられるのも許せないので修正したのですが、消費メモリ1/1000、処理時間1/4と大幅なパフォーマンスアップを達成。

とは言っても凄い魔法を使ったと言うわけではなく、元々のアレな部分をセオリー通りのコードに書き直しただけ、誰でもできる改善です。

バッチの内容

バッチはAuroraのデータをselectしてtsvに変換してs3にアップすると言う単純なものです。

1. DBからテーブル毎にスキーマ情報とデータを取得
2. それをTSVファイルとして出力
3. gzip圧縮
4. s3にアップロード

なんでこんなのが20GBも使うの？という感じですが、Rubyに慣れてない人が書いたというならだいたい予想できますし、予想通りでした。

コマンドを活用

Rubyスクリプトの修正の前にバッチそのものの設計を軽く直しました。

UNIX哲学にはこう記されています。

小さいものは美しい。各プログラムが一つのことをうまくやるようにせよ。

つまり１つのスクリプトで全てをやる必要なんてないのです。

元のバッチのgzip圧縮処理は、一度出力したtsvを再度オープンしてGzipWriterで書き出すという効率の悪い処理をしていました。最初にファイルに書き出す時点で圧縮すべきです。しかし、そんな修正をせずとも出力したTSVファイルのパスを標準出力に書き出して

ruby create_tsv.rb | xargs gzip 

としてしまった方が圧倒的に簡単で、しかも速いのです。

awsへのアップロードも再度ファイルを開いてput_objectすると言う非効率なものでしたが、TSVファイル出力先ディレクトリを aws s3 sync コマンドでアップロードするようにしました。

ruby create_tsv.rb --output-dir=./tmp --db-config=./db/${env}.yaml < tables.txt | xargs gzip
aws s3 sync ./tmp/ s3://${bucket_name}/${env}/

バッチ本体はこのようなシェルスクリプトになりました。この他に前処理としてロック獲得だのgemライブラリのbundle installだの色々やっております。

Rubyスクリプトの修正

MySQLライブラリの置き換え

まずruby-mysqlをrailsでもおなじみのmysql2に変更。

streamingの使用

元のコードはselect結果をクライアントに全て持ってきて処理していたため、メモリを大量に消費していました。mysql2はstreamingというオプションを渡すと１行ずつデータを取得するようになり、メモリ消費が抑えられます。

db.query(sql, :streaming => true).results

ただしstreamingオプションはクライアントのリソース消費が減る分、DBサーバのリソースを食うのでその点は注意が必要。

処理の効率化

元々のコードはこうなっていました

tsv = db.query(sql).map { |row| row_to_tsv_record(row) }
CSV.open { |f|
    tsv.each { |tsv|
        f << tsv
    }
}

これではせっかくstreamingを使っても、テーブルの全データが入った超巨大なArrayを作ってしまうためメモリを大量消費してします。また、2回ループするので時間もかかります。これを素直に書き直しました。

CSV.open { |f|
    db.query(sql, :streaming => true).results(:as => :array).each {
        |row| f << row_to_tsv_record(row)
    }
}

結果

ちまちまと修正した結果、最終的には処理時間は1/4、メモリはRSSで20GB→20MB、VSZで18GB→200MB程度まで減り、バッチを実行していたインスタンスがr4.xlargeからt2.mediumになりました。

DBのタイムアウトはサーバ側をクライアント側より長くする

多くのデータベースにはコネクションのタイムアウトという設定が存在します。指定した期間使われていないコネクションを解放するというもので、クライアント、サーバ双方で設定可能になっています。

タイムアウトはサーバ側の寿命がクライアント側よりも長くなるように設定すべきです。逆にしてしまうとサーバ側で切断されたコネクションに対してクライアントが読み書きを行おうとしてエラーになってしまうことがあります。

MySQLとnode.jsで再現すると以下のようになります。

$ node index.js
[ RowDataPacket { Variable_name: 'wait_timeout', Value: '28800' } ]
wait 11 seconds
[ RowDataPacket { '1': 1 } ]
end
$
$ mysql -u root -e "set global wait_timeout=10;"
$
$ node index.js
[ RowDataPacket { Variable_name: 'wait_timeout', Value: '10' } ]
wait 11 seconds
{ Error: Cannot enqueue Query after fatal error.
    at Protocol._validateEnqueue (/Users/nullpon/mysql-node/node_modules/mysql/lib/protocol/Protocol.js:201:16)
    at Protocol._enqueue (/Users/nullpon/mysql-node/node_modules/mysql/lib/protocol/Protocol.js:139:13)
    at PoolConnection.query (/Users/nullpon/mysql-node/node_modules/mysql/lib/Connection.js:208:25)
    at PoolConnection.query (internal/util.js:230:26)
    at __dirname (/Users/nullpon/mysql-node/index.js:21:59)
    at <anonymous> code: 'PROTOCOL_ENQUEUE_AFTER_FATAL_ERROR', fatal: false }

index.jsのコードは以下になります

'use strict';
const {promisify} = require('util')
const mysql = require('mysql');

(async () => {
    var pool  = mysql.createPool({
        host     : 'localhost',
        user     : 'root',
        password : '',
        database : 'test',
        acquireTimeout: 60000 // timeout 60秒
    });
    const conn = await promisify(pool.getConnection).bind(pool)();

    const result1 = await promisify(conn.query).bind(conn)('show global variables like \'wait_timeout\'');
    console.log(result1);

    console.log('wait 11 seconds');
    await new Promise((r) => setTimeout(r, 11000))

    const result2 = await promisify(conn.query).bind(conn)('SELECT 1');
    console.log(result2);

    conn.release();

    await promisify(pool.end).bind(pool)();

})().then(() => {
    console.log('end');
}).catch((e) => {
    console.log(e);
});

Redisのようにデフォルトでサーバ側のタイムアウトが無限大となっているDBもあります。特にAWSのElastiCacheでありがちなのですが、アプリがOOMなどで終了したり、正常終了時でもコネクションを解放しない雑な実装だったりすると、サーバ側のコネクションは消えずに残ってしまいます。タイムアウトを無限大にしておくと、このようなゴミが永遠に解放されず少しずつ溜まっていき、いずれサーバのコネクション数の上限に達して障害の原因となります。

こういう事態を防ぐためにもサーバ側でもタイムアウトを設定した方が無難なのですが、クライアントよりも短い値を設定しないように注意してください。