GraphQL RubyがN+1問題の解決をFiberによってどう実現しているか

GraphQLはREST APIと異なり、クエリを使って取得するデータを柔軟に指定することが可能です。これは便利な反面、実装によっては非効率なデータアクセスを行ってしまうことがあります。

例えば、記事を格納するモデルがあり、外部キーとして作成者を持つものとします。

class Article < ApplicationRecord
  belongs_to :creator, class_name: 'User'
end

記事の一覧をGraphQLで取得する際のクエリは以下のようになるでしょう。

query getArticles {
  articles {
    nodes {
     id
     title
     body
     creator {
       name
     }
     createdAt
   }
  }
}

ArticleTypeはこのように定義したとします。

article_type.rb

module Types
  class ArticleType < Types::BaseObject
    field :id, ID, null: false
    field :title, String, null: false
    field :body, String, null: false
    field :creator, UserType, null: false
    field :created_at, GraphQL::Types::ISO8601DateTime, null: false
  end
end

この例では、ArticleをイテレートするたびにUserへのクエリが実行されてしまうので、いわゆるN+1問題が発生します。通常のRailsアプリケーションでは、includeやpreloadを使ってN+1問題を解決するわけですが、GraphQLだとアプリケーションからどのようなクエリが発行されるか予め決定できません。これを防ぐしくみが、DataLoaderです。

GraphQL - DataLoader Overview

DataLoaderを使ってみる

DataLoaderを使うには、GraphQL::Dataloader::Sourceを継承したクラスを作成します。

record_loader.rb

# frozen_string_literal: true

module Loaders
  class RecordLoader < GraphQL::Dataloader::Source
    def initialize(model, column: model.primary_key)
      @model = model
      @column = column.to_s
      @column_type = model.type_for_attribute(@column)
    end

    def load(key)
      super(@column_type.cast(key))
    end

    def fetch(ids)
      records = @model.where(@column => ids)

      # return a list with `nil` for any ID that wasn't found
      # https://graphql-ruby.org/dataloader/sources.html
      ids.map { |id| records.find { |r| r.send(@column) == id } }
    end
  end
end

そうするとBaseObjectを継承したクラスから以下のように呼び出すことができます。

user = dataloader.with(Sources::RecordLoader, ::User).load(1)

いちいちこのように呼び出すのは面倒くさいので、新たにBatchLoaderFieldというクラスを用意して、field定義から呼び出せるようにします。

field :creator, UserType, null: false, loader: :record_loader

そうすると、usersテーブルからデータをまとめて取得することができるようになります。Railsのログを見ても、SQLが一回だけ発行されているのがわかります。

User Load (0.3ms) SELECT `users`.* FROM `users` WHERE `users`.`id` IN (20001, 20000, 20002, 20003)

has_one/has_manyのアソシエーションに対しても、同様にローダーを作成することができます。

field :articles, [ArticleType], null: false, loader: :association_loader

サンプルコードはこちらにまとめてあります。

ryokdy/graphql-ruby-dataloader-example

どうやって実現しているのか

DataLoaderを使うことで無事GraphQLのN+1問題が解決できましたね！

で、ここからが本題です。このしくみってどうやって実現しているのか不思議じゃないですか？　だってRecordLoader#loadを呼び出したときに目的のレコードを取得するわけですよね。でも実際のfetch処理では複数のIDからまとめてデータを取得するわけですから、なんとなく未来に取得するであろうIDを予め予測してfetchしているような気がしてしまいます。もちろんそんなわけはなくて、RubyのFiberという軽量スレッドを使って実現しているのですが、僕がこれまでFiberの用途をよくイメージできておらず、今回ソースコードを読んでみて興味深く感じましたので、紹介したいと思います。

GraphQL Rubyの公式ドキュメントにはこのように書かれています。

GraphQL::Dataloader uses Ruby’s Fiber, a lightweight concurrency primitive which supports application-level scheduling within a Thread. By using Fiber, GraphQL::Dataloader can pause GraphQL execution when data is requested, then resume execution after the data is fetched.
At a high level, GraphQL::Dataloader’s usage of Fiber looks like this:

GraphQL::DataloaderはRubyのFiberを使用しており、Thread内のアプリケーションレベルのスケジューリングをサポートする軽量な並行処理プリミティブです。Fiberを使用することで、GraphQL::Dataloaderはデータが要求されたときにGraphQLの実行を一時停止し、データがフェッチされた後に実行を再開することができます。高レベルでは、GraphQL::DataloaderのFiberの使用方法は次のようになります：

GraphQL execution is run inside a Fiber.

GraphQLの実行はFiberの内部で実行されます。

When that Fiber returns, if the Fiber was paused to wait for data, then GraphQL execution resumes with the next (sibling) GraphQL field inside a new Fiber.

その Fiber が戻ると、Fiber がデータを待つために一時停止していた場合は、新しい Fiber 内の次の（兄弟） GraphQL フィールドで GraphQL の実行が再開されます。

That cycle continues until no further sibling fields are available and all known Fibers are paused.

このサイクルは、それ以上兄弟フィールドが利用できなくなり、すべての既知の Fiber が一時停止されるまで続きます。

GraphQL::Dataloader takes the first paused Fiber and resumes it, causing the GraphQL::Dataloader::Source to execute its #fetch(...) call. That Fiber continues execution as far as it can.

GraphQL::Dataloaderは最初に一時停止されたFiberを取得して再開し、GraphQL::Dataloader::Sourceに#fetch(...)呼び出しを実行させます。そのファイバーは、可能な限り実行を継続します。

Likewise, paused Fibers are resumed, causing GraphQL execution to continue, until all paused Fibers are evaluated completely.

同様に、一時停止していたファイバーが再開され、一時停止していたすべてのファイバーが完全に評価されるまで、GraphQLの実行が続行されます。

こちらはDataLoaderのメインルーチンです。

dataloader.rb

def run
  job_fibers = []
  next_job_fibers = []
  source_fibers = []
  next_source_fibers = []
  first_pass = true
  manager = spawn_fiber do
    while first_pass || job_fibers.any?
      first_pass = false
      # GraphQLのリクエストを処理するためのFiberを実行する
      while (f = (job_fibers.shift || spawn_job_fiber))
        if f.alive?
          finished = run_fiber(f)
          if !finished
            next_job_fibers << f
          end
        end
      end
      join_queues(job_fibers, next_job_fibers)

      # データを取得するためのFiberを実行する
      while source_fibers.any? || @source_cache.each_value.any? { |group_sources| group_sources.each_value.any?(&:pending?) }
        while (f = source_fibers.shift || spawn_source_fiber)
          if f.alive?
            finished = run_fiber(f)
            if !finished
              next_source_fibers << f
            end
          end
        end
        join_queues(source_fibers, next_source_fibers)
      end
    end
  end

  # メインルーチンのFiberを生成する
  run_fiber(manager)

メインルーチンから、通常のGraphQLの処理を担当するFiberが実行されます。このとき処理の中断が何もなければ、すべての兄弟のGraphQLフィールドが処理されることになります。

ところが、GraphQLの処理の中でRecordLoader#loadが呼び出されると、内部でGraphQL::DataLoader::Source#syncが呼び出されます。

source.rb

def sync(pending_result_keys)
    # 親スレッドに制御を戻す
    @dataloader.yield
  iterations = 0
  while pending_result_keys.any? { |key| !@results.key?(key) }
    iterations += 1
    if iterations > 1000
      raise "#{self.class}#sync tried 1000 times to load pending keys (#{pending_result_keys}), but they still weren't loaded. There is likely a circular dependency."
    end
    # 親スレッドに制御を戻す
    @dataloader.yield
  end
  nil
end

@dataloader.yieldが呼び出されると、GraphQLのフィールドの処理を中断して、親スレッドに制御が戻ります。親スレッドは兄弟のフィールドの処理を続行しますが、そこでも#loadが呼び出されるので、このタイミングで処理が中断します。結果として、すべての兄弟フィールドの処理がloadで中断している状態で、データフェッチのためのFiberが実行されることになるわけです。この時点で取得すべきレコードのIDはすべて揃っているわけですから、一括でレコードを取得することができます。

def run_pending_keys
  if !@fetching.empty?
    @fetching.each_key { |k| @pending.delete(k) }
  end
  return if @pending.empty?
  fetch_h = @pending
  @pending = {}
  @fetching.merge!(fetch_h)
  # Pending状態のキーを使ってデータを一括取得する
  results = fetch(fetch_h.values)
  fetch_h.each_with_index do |(key, _value), idx|
    @results[key] = results[idx]
  end
  nil
rescue StandardError => error
  fetch_h.each_key { |key| @results[key] = error }
ensure
  fetch_h && fetch_h.each_key { |k| @fetching.delete(k) }
end

データをfetchしたら、ふたたびGraphQLの処理に制御が戻りますので、中断していたloadの処理が再開されることになります。ファイバーはスレッドと違って、プログラマーが処理の切り替えをコントロールするので、それってふつうのサブルーチンとどう違うのって思いがちですが、このように遅延実行のためにFiberを使って処理を中断させることでうれしいケースがあることがよくわかりました。

GraphQL RubyがN+1問題の解決をFiberによってどう実現しているか

DataLoaderを使ってみる

どうやって実現しているのか

関連する記事