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哲学にはこう記されています。

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

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

元のバッチの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-mysqlrailsでもおなじみのmysql2に変更。

streamingの使用

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

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になりました。