RSpecでオブジェクト割り当てをテストする——パフォーマンス最適化のリグレッションを防ぐ

最近、Rubyのパフォーマンスが大きな話題になっています。実は、コードにほんの少しの修正を加えるだけで、パフォーマンスを最大99.9%も向上させることが可能です。

最適化はなぜ「元に戻ってしまう」のか

コードの最適化方法を解説する記事は数多く存在しますが、その最適化が維持され続けていることをどうやって保証すればよいのでしょうか。

頻繁に呼び出されるメソッド内で、フリーズした定数ではなく文字列リテラルを埋め込んだ場合の影響を、常に意識しているとは限りません。将来のコード保守の中で、せっかくの最適化による恩恵をうっかり失ってしまうことは、とても起こりやすいのです。

私自身、HoneybadgerのRuby gemのコードを2回目(あるいは3回目)の最適化をしていたときに、ふとこんなことを考えました。「これらの最適化がリグレッション(退行)しないことを保証する仕組みがあれば素晴らしいのに」と。

リグレッションとは何か

リグレッションは、名前こそ知らなくても、ソフトウェア開発に携わるほとんどの人が経験したことのある現象です。過去に解決済みだったバグや問題が、同じコードへの後の変更によって再発することを指します。同じ作業を二度やりたくないのは誰でも同じでしょう。リグレッションとは、掃き掃除をした直後に床へ土を持ち込んでしまうようなものです。

幸い、私たちには秘密兵器があります。そう、テストです。厳格なTDDを実践しているかどうかにかかわらず、テストはバグ修正において非常に強力です。問題とその解決策をプログラムとして明示できるからです。テストがあれば、変更を加えてもリグレッションは起こらないという自信が得られます。

どこかで聞いた話だと思いませんか?私もそう感じました。そこでこう考えたのです。「パフォーマンス最適化にもリグレッションがあり得るなら、それもテストで検出できるはずではないか?」と。

プロファイリングツールとallocation_stats

Rubyには、オブジェクト割り当て、メモリ、CPU、ガベージコレクションなど、さまざまなパフォーマンスの側面をプロファイリングする優れたツールが多数あります。代表例としては、ruby-prof、stackprof、allocation_tracerなどが挙げられます。

私は最近、オブジェクト割り当てのプロファイリングにallocation_statsを使っています。割り当てを減らすのは比較的簡単な作業でありながら、メモリ消費量や速度のチューニングにおいて、数多くの「手の届きやすい成果」をもたらしてくれるのです。

例として、デフォルト値が'foo'である文字列を5つ格納するArrayを持つ、基本的なRubyクラスを見てみましょう。

class MyClass
  def initialize
    @values = Array.new(5)
    5.times { @values << 'foo' }
  end
end

AllocationStatsのAPIはシンプルです。プロファイリングしたいブロックを渡すだけで、どこで最も多くのオブジェクトが割り当てられているかを出力してくれます。

$ ruby -r allocation_stats -r ./lib/my_class
stats = AllocationStats.trace { MyClass.new } 
puts stats.allocations(alias_paths: true).group_by(:sourcefile, :sourceline, :class).to_text
^D
     sourcefile        sourceline   class   count
---------------------  ----------  -------  -----
/lib/my_class.rb           4       String       5
/lib/my_class.rb           3       Array        1
-                          1       MyClass      1

#to_textメソッド(割り当てのグループに対して呼び出します)は、指定した基準でグループ化された、人間が読みやすい整形済みの表を出力します。

手動でのプロファイリングにはこの出力で十分ですが、私の目標は、RSpecで書かれた通常のユニットテストスイートと一緒に実行できるテストを作ることでした。my_class.rbの4行目で5つの文字列が割り当てられていますが、すべて同じ値だと分かっている以上、これは不要に思えます。私が望んでいたのは、「MyClassの初期化時には6個未満のオブジェクトを割り当てる」というような、読みやすいシナリオ記述です。RSpecでは次のように書けます。

describe MyClass do
  context "when initializing" do
    specify { expect { MyClass.new }.to allocate_under(6).objects }
  end
end

カスタムRSpecマッチャーの作成

この構文を使えば、expectブロック内のコードについて、オブジェクト割り当て数が指定値未満であることを、カスタムRSpecマッチャーでテストできます。

トレース結果の出力に加えて、AllocationStatsはRubyから割り当て情報へアクセスするためのいくつかのメソッド(#allocationsや#new_allocationsなど)を提供しています。私はこれらを活用してマッチャーを構築しました。

begin
  require 'allocation_stats'
rescue LoadError
  puts 'Skipping AllocationStats.'
end

RSpec::Matchers.define :allocate_under do |expected|
  match do |actual|
    return skip('AllocationStats is not available: skipping.') unless defined?(AllocationStats)
    @trace = actual.is_a?(Proc) ? AllocationStats.trace(&actual) : actual
    @trace.new_allocations.size < expected
  end

  def objects
    self
  end

  def supports_block_expectations?
    true
  end

  def output_trace_info(trace)
    trace.allocations(alias_paths: true).group_by(:sourcefile, :sourceline, :class).to_text
  end

  failure_message do |actual|
    "expected under #{ expected } objects to be allocated; got #{ @trace.new_allocations.size }:\n\n" << output_trace_info(@trace)
  end

  description do
    "allocates under #{ expected } objects"
  end
end

最初のrequire文でLoadErrorをrescueしているのは、すべてのテスト実行でAllocationStatsを含めたくない場合があるためです(テスト全体が遅くなる傾向があります)。続いて:allocate_underマッチャーを定義し、matchブロック内でトレースを実行します。failure_messageブロックも重要なポイントです。AllocationStatsトレースのto_text出力を、失敗メッセージの中に直接組み込んでいるからです!マッチャーの残りの部分は、ほぼ標準的なRSpecの設定です。

テストの実行——まずは失敗を確認する

マッチャーを読み込めば、先ほどのシナリオを実行して、失敗する様子を確認できます。

$ rspec spec/my_class_spec.rb 

MyClass
  when initializing
    should allocates under 6 objects (FAILED - 1)

Failures:

  1) MyClass when initializing should allocates under 6 objects
     Failure/Error: expect { MyClass.new }.to allocate_under(6).objects
       expected under 6 objects to be allocated; got 7:

               sourcefile           sourceline   class   count
       ---------------------------  ----------  -------  -----
       <PWD>/spec/my_class_spec.rb           6  MyClass      1
       <PWD>/lib/my_class.rb                 3  Array        1
       <PWD>/lib/my_class.rb                 4  String       5
     # ./spec/my_class_spec.rb:6:in `block (3 levels) in <top (required)>'

Finished in 0.15352 seconds (files took 0.22293 seconds to load)
1 example, 1 failure

Failed examples:

rspec ./spec/my_class_spec.rb:5 # MyClass when initializing should allocates under 6 objects

これで、パフォーマンス上の問題——MyClassが同じ値の余分な文字列オブジェクトを割り当てていること——をプログラム的に実証できました。次は、これらの値をフリーズした定数に移して問題を修正しましょう。

class MyClass
  DEFAULT = 'foo'.freeze

  def initialize
    @values = Array.new(5)
    5.times { @values << DEFAULT }
  end
end

修正後——テストが通ることを確認する

問題を修正したので、もう一度テストを実行すると、今度は成功します。

$ rspec spec/my_class_spec.rb

MyClass
  when initializing
    should allocates under 6 objects

Finished in 0.14952 seconds (files took 0.22056 seconds to load)
1 example, 0 failures

今後MyClass#initializeメソッドを変更しても、過剰なオブジェクトを割り当てていないことに自信が持てます。

パフォーマンステストを必要なときだけ実行する

割り当てのプロファイリングは比較的処理が重いため、常時ではなくオンデマンドで実行するのが理想的です。allocation_statsが存在しない場合も適切にハンドリングしているため、Bundlerで複数のGemfileを作成し、BUNDLE_GEMFILE環境変数で使用するGemfileを切り替えられます。

$ BUNDLE_GEMFILE=with_performance.gemfile bundle exec rspec spec/
$ BUNDLE_GEMFILE=without_performance.gemfile bundle exec rspec spec/

もうひとつの選択肢は、appraisal gemのようなライブラリを利用することです。これは同じアプローチを採りつつ、Bundler特有の落とし穴を解決してくれます。Jason Clarkが2015年3月のRuby on Alesで、この手法に関する優れたプレゼンテーションを行っています。詳細は彼のスライドを参照してください。

また、こうした種類のテストは通常のユニットテストとは別に管理するのが良いと考えています。新しい「performance」ディレクトリを作成し、ユニットテストスイートはspec/unit/に、パフォーマンススイートはspec/performance/に配置します。

spec/
|-- spec_helper.rb
|-- unit/
|-- features/
|-- performance/

まとめ

私はまだRubyコードのパフォーマンスプロファイリングへのアプローチを磨いている途中ですが、パフォーマンステストスイートを維持することで、今この瞬間のコード速度向上、将来にわたる高速性の維持、そして自分自身や他の開発者のための生きたドキュメント作成につながると期待しています。