ActiveRecordのように振る舞うオブジェクトでJSON生成を高速化

こんにちは、freee ソフトウェアエンジニア @terashi58 です。
この記事は freee Engineers Advent Calendar 2015 7日目です。

ActiveRecord のパフォーマンス

Rails の ActiveRecord 便利ですよね。
私はまだ Rails 一年生ですが、日々お世話になっています。
個人的には has_many とかの Relation と where に始まる QueryMethods が特にいい感じです。

しかしこの ActiveRecord、便利機能満載のためか非常に遅いことがしばしば。
プロファイルを取ると、たいていはコンストラクタの呼び出しが重いです。
特に index 処理では大量の AR オブジェクトを生成するためボトルネックになります。

index (show) 処理、特に出力がJSONの場合は、DB の値をほとんどそのまま格納することが多く、その場合は AR オブジェクトまでは必要としないことがよくあります。
そのため、クエリ生成まで QueryMethods で行い、to_sql を ActiveRecord::Base.connection.select_all などで Hash に格納してしまうのが有効です。

@hotchpotch さんの activerecord-row-data gem がその辺をよかれとやってくれます。

Hash だけでは足りない

シンプルな REST API などではこれで十分なのですが、複雑なモデル構造を一度に返す必要がある場合には、依然 Active Record の Relation や、モデルのメソッドを流用したくなることがあります。
freee の場合、取引の一覧表示などでは軽く 10 以上のモデルが再帰的に関与し、そうなるとさすがに Hash をゴリゴリするのはつらいです。

また、現在使っている jbuilder や ActiveRecord::Serializer をそのまま流用したいというのもあります。既に AR オブジェクトがあればそのまま渡したいですし、なにより移行によるエンバグが怖いです。

PseudoRecord: ActiveRecord のように振る舞うもの

そこで Hash をラップして ActiveRecord のように振る舞わせる PseudoRecord というライブラリを作ることにしました。
Ruby では犬のように振る舞うものは犬（として扱っていいの）です。

例えば、こんな感じで記述します。

pseudo/models.rb

PseudoRecord.define Article do
  attributes :id, :title, :content 
  belongs_to :user
  has_many :comments
end

PseudoRecord.define User do
  import :display_name  # User#display_nameを取り込む

  attributes :id, :first_name, :last_name, :email
end

PseudoRecord.define Comment do
  attributes :id, :article_id, :title, :content
  has_many :replies
end

PseudoRecord.define Reply do
  attributes :id, :comment_id, :content
end

ActiveModel::Serializer を利用する場合こんな感じで使います。

.rb

  records = Article.where(conditions)
  render json: PseudoRecord::Article.build_from(records), each_serializer: ArticleSerializer

records を渡すところを一つラップするだけです。
Serializer のコードは省略しますが通常となにも変更せずに、Article に関連する値をまとめて取ってこれます。

Ruby メタプログラミング、それは自分の足を打ち抜ける道具

ぱっと見でわかると思いますが、PseudoRecord の中はバリバリのメタプログラミングです。Rails 程ではないですが。

では、一つずつ中身を解説していきます。

Attributes

attributes は OpenStruct のように単純にメソッドを作るのと、SELECT対象を指定します。

pseudo_record/core.rb

  # ActiveModel::Serializerはデフォルトではread_attribute_for_serializationで
  # attributesの値を取得する
  def read_attribute_for_serialization(name)
    public_send(name)
  end

  module ClassMethods
    def attributes(*attrs)
      @attributes ||= []
      attrs.each do |name|
        next if @attributes.include? name
        @attributes << name
        if block_given?
          yield name
        else
          define_method name do
            @hash[name]
          end
        end
      end
    end
  end

ちょっとトリッキーなのはメソッド定義を外から与えられる点です。
これにより Time 型を TimeWithZone にするための time_attributes や enumerize を Value 型にするための enum_attributes のような拡張を容易にできます。
また、手動でアクセサメソッドをいくらでも上書き可能です。

Relation

Relation の作成はサブモデルのロードとセットになります。

pseudo_record/core.rb

  module ClassMethods
    def belongs_to(name, opts = {})
      add_relation(name, opts, id_key: :"#{name}_id", foreign_key: :id, ...)
    end
    def has_one(name, opts = {})
      add_relation(name, opts, id_key: :id, foreign_key: :"#{self.name.foreign_key}", ...)
    end
    def has_many(name, opts = {})
      add_relation(name, opts, id_key: :id, ..., to_array: true)
    end

    # relationの共通実装
    def add_relation(name, opts, params)
      updator = :"update_#{name}"
      id_key = params[:id_key]
      foreign_key = params[:foreign_key]
      model_class = params[:model_class]
      pseudo_class = params[:pseudo_class]
      to_array = params[:to_array] || false

      attributes id_key
      define_method name do
        @hash[name]
      end
      define_method updator do |id_map|
        @hash[name] = id_map[@hash[id_key]]
      end
      add_preloader do |models|
        ids = models.map{|model| model.hash[id_key]}.uniq
        next if ids.blank?
        sub_models = pseudo_class.select_by(model_class.where(foreign_key => ids), opts)
        if to_array
          id_map = ids.map{|id| [id, []]}.to_h
          sub_models.each{|model| id_map[model.public_send(foreign_key)] << model}
        else
          id_map = sub_models.map{|model| [model.public_send(foreign_key), model]}.to_h
        end
        models.each{|model| model.public_send(updator, id_map)}
        pseudo_class.preload(sub_models) unless opts[:shallow]
      end
    end

     # 擬似モデル列の生成時に行う初期化処理を登録する
     def add_preloader(&proc)
      @preloaders ||= []
      @preloaders << proc
    end

    # 登録された初期化処理を実行する
    def preload(models)
      return unless @preloaders
      @preloaders.each{|proc| proc.call(models)}
    end
  end

このライブラリの一番のコア部分。再帰的に依存関係を N+1 を回避しつつ解決し、アクセスメソッドを提供します。
長いんで省略しましたが、belongs_to、has_one、has_many は ActiveSupport::Inflector のメソッドを使って適切な値を add_relation に渡します。

Modelのメソッド追加

pseudo_record.rb

module PseudoRecord
  # 擬似モデルクラスを作成する
  def self.define(base_class,  &block)
    c = Class.new
    c.extend Capture
    c.class_eval(&block)
    if c.imported_methods.present?
      p = Class.new(base_class)
      import(p, base_class, c.imported_methods)
    else
      p = Class.new
    end
    const_set(class_name, p)
    p.include PseudoModels::Core
    p.class_eval(&block)
    p
  end

  private

  # ブロックをスキャンし、import内容を取得するだけのモジュール
  module Capture
    attr_reader :imported_methods
    def import(*args)
      @imported_methods = args
    end
    def method_missing(name, *args)
    end
  end

  # 親クラスの指定したメソッドのみを取り込む
  def self.import(pseudo_class, base_class, methods)
    (base_class.instance_methods - PseudoModels::Core.instance_methods -
     Object.instance_methods - methods).each do |name|
      # privateなundef_methodを無理やり呼ぶ
      pseudo_class.send(:undef_method, name)
    end
  end
end

メソッドの取り込みは継承で行い、不要なメソッドを undef_method します。

元々は UnboundMethod#bind でなんとかならないかと頑張ってみたものの、C コードの方でチェックしているのでどうしようもなかったのです。
（Ruby は黒魔術し放題と思っていましたが、意外と白かった。）

freee の今のコードではこの undef の手法は使わず、単純に継承だけして使っていたのですが、エラーメッセージが分かりにくい（#[] で落ちたとか言われてもどれやねんという）ので、明示的に許可したもののみにすることにして、現在レビュー申請中。

なんで const_set をわざわざしているのかというと、相互依存した場合に定数の参照がループするからです。

.rb

Pseudo::User = PseudoRecord.define do
  import User, :display_name
  ...
end

みたいな方が、git grep で探しやすいかと思ったのですが無理でした。
仕方ないのでコメントで書いておくことに。

また、他の方法としては、必要なメソッドを全て Concern にしてモジュールの include にするというのがあります。
こちらの方が安全で綺麗なのですが、問題は元のモデルの定義に手を入れる必要がある点です。
たいていの使いたいメソッドはそのモデルの極めて基本的なメソッドで、通常モデルクラスに直書きするようなものがほとんどです。
多くの人が見るコードの可読性を高速化のために下げるのは避けたいので、ちょっと強引ですがライブラリ側が頑張る方向にしました。

Construction

最後に生成部分です。

pseudo_record/core.rb

  # optsはoverride用
  def initialize(hash, opts = {})
    @hash = hash.symbolize_keys
  end

  module ClassMethods
    # ActiveRecordはnewで余計な処理を挟むので、素のnewの動作に戻す
    def new(*args, &block)
      obj = allocate
      obj.send(:initialize, *args, &block)
      obj
    end

    # Polymorphic Modelを継承した場合にtype検索にPseudoRecordを含めないようにする
    def sti_name
      superclass.sti_name
    end

    # Modelのscopeなどから擬似モデルの配列を作成する
    def build_from(relation)
      models = select_by(relation)
      preload(models)
      models
    end

    # Modelのscopeなどから擬似モデルの配列を作成する
    # ただし、関連するモデルは読み込まない
    def select_by(relation, opts = {})
      ActiveRecord::Base.connection.select_all(relation.select(@attributes).to_sql).map{|hash| self.new(hash, opts)}
    end
  end  

地味に厄介だったのが ActiveRecord の new が結構いろいろな処理をしていて継承するだけで遅くなることでした。
コードを追っても回避する手段がなかったので、超強引だけど new を上書きし、super を呼ばずに allocate、initialize で済ませることにしました。

sti_name の方は where 句で無駄な検索をさせないためのものです。
ただ、Public API ではないので、Rails のバージョンが上がると手を入れる必要があるかもしれないのが難点ですが。

Test, Test Test

このライブラリは使うのは簡単だけど、ちょっと記述が抜けているとあっさり落ちます。
このためテストをきちんと書いておくのが必須です。

幸い ActiveRecord と比較して同じならいいので、関連するレコードを全部作って、ActiveRecord 版JSONと PseudoRecord 版JSONの文字列一致を行う spec を必ず用意することで安心して開発できています。（そしてよくきちんと落ちてくれます）

どれだけ早くなったのか

開発環境で大きなリクエストでは process time で 3.5 倍くらい早くなりました。
もっとも、プロダクションだと他のボトルネックや index 対象が少ないケースもあったため、平均で 20% 程度でしたが。
今後、このライブラリを他のボトルネックにも適用していく予定です。

今後やりたいこと

• 参照元オブジェクトへのバックリンク（has_many 先の belongs_to）
• 明示的な preload
  • いまは Relation があると必ずロードするため、相互参照を記述できない
  • preload みたいに指定したものだけロードするようにしたい
• JSON生成以外への応用
• 他のレポジトリで使えるように gem 化

また、gem 化にともない、他でも使えそうならオープンソースにできないかなと考えています。
一応オープンソース化に備えて、この記事内のコードはMIT Licenseと明示しておきます。
（LICENSE.txt が必要なのかな。とっとと使いたいからライセンスきちんとしろや、という方はコメントなりで声をかけてください。）

宣伝

freee ではビジネス向けアプリの UX からサーバの CPU までのあらゆるレイヤを攻める*エンジニアを募集しています。よろしくお願いします。

明日は freee の世界遺産、ミケランジェロが彫ったモアイこと @ryosukeYamazaki よる MQ についてです。お楽しみに。