非同期クエリの `caller_runs` フォールバックによるインストゥルメンターの汚染を修正

async_query_executor に :global_thread_pool を使用している環境で、スレッドプールの飽和時に同期クエリの sql.active_record 通知がサイレントに失われるバグが修正されました。

背景

async_load や async_pick を使った非同期クエリは、global_thread_pool のスレッドが枯渇すると呼び出し元スレッドで実行されることがあります。これは fallback_policy: :caller_runs の挙動によるもので、global_thread_pool はデフォルトの concurrency: 4 スレッドと最大キューサイズ16を合わせた計20タスクが上限となっています。

Active Recordの接続プールリーパーもこの global_thread_pool で非同期メンテナンスタスクを実行するため、Pumaのスレッド数が少ない環境でもプールは徐々に飽和していきます。PR作者が実際に観測した症状は「デプロイ直後は問題ゼロ、その後リクエストごとにクエリログが欠落する割合が増加し、次のデプロイで再びゼロに戻る」という周期的なパターンで、これはスレッドが汚染されるとワーカーが再起動されるまで永続的に汚染され続けることを示しています。

根本原因は QueryIntent#execute_or_skip にありました。このメソッドは非同期クエリを実行する際に ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] を EventBuffer に置き換えますが、終了後に元の値を復元していませんでした。バックグラウンドスレッドで実行される通常ケースでは IsolatedExecutionState がスレッドローカルであるため問題になりませんが、caller_runs により呼び出し元スレッド上で実行されると、その後の全 sql.active_record 通知が EventBuffer に吸収されたまま公開されなくなります。

技術的な変更

修正は query_intent.rb の execute_or_skip メソッドに2行を追加するだけのシンプルなものです。EventBuffer のセットアップ前に既存のインストゥルメンターを保存し、ensure ブロックで確実に復元します。

変更前:

@pool.with_connection do |connection|
  return unless @mutex.try_lock
  begin
    if pending?
      @event_buffer = EventBuffer.new(self, ActiveSupport::Notifications.instrumenter)
      # ...
    end
  rescue => error
    @error = error
  ensure
    @mutex.unlock
  end
end

変更後:

@pool.with_connection do |connection|
  return unless @mutex.try_lock
  previous_instrumenter = ActiveSupport::IsolatedExecutionState[:active_record_instrumenter]
  begin
    if pending?
      @event_buffer = EventBuffer.new(self, ActiveSupport::Notifications.instrumenter)
      # ...
    end
  rescue => error
    @error = error
  ensure
    ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] = previous_instrumenter
    @mutex.unlock
  end
end

バックグラウンドスレッドで実行される通常ケースでは previous_instrumenter が nil になるため、実質的にno-opです。caller_runs フォールバックで呼び出し元スレッド上で実行された場合は、ensure ブロックが呼び出し元の実インストゥルメンターを復元します。

テストでは、schedule_query メソッドをモンキーパッチして execute_or_skip を現在のスレッドで直接実行することで caller_runs の状況を再現しています。非同期クエリ実行後に ActiveSupport::Notifications.subscribed で同期クエリの通知が正常に届くことを検証します。

def test_execute_or_skip_does_not_contaminate_caller_thread_instrumenter
  skip unless ActiveRecord::Base.connection.async_enabled?

  pool = Post.connection_pool
  old_schedule_query = pool.method(:schedule_query)
  pool.singleton_class.undef_method(:schedule_query)
  pool.singleton_class.define_method(:schedule_query) do |future_result|
    future_result.execute_or_skip  # caller_runs をシミュレート
  end

  Post.async_count  # execute_or_skip を現在スレッドで実行

  notification_called = false
  ActiveSupport::Notifications.subscribed(->(*) { notification_called = true }, "sql.active_record") do
    Post.count  # 通知が正常に届くか確認
  end

  assert notification_called,
    "sql.active_record notification was not published after execute_or_skip ran on the caller thread"
ensure
  # ...
  ActiveSupport::IsolatedExecutionState.delete(:active_record_instrumenter)
end

設計判断

「復元しない」という設計の前提が caller_runs で崩れることへの対処 として、最小限のスコープで save/restore パターンを適用する方針が採られました。

元の実装は非同期クエリが常に独立したスレッドで実行されるという前提の下では正しく機能していました。IsolatedExecutionState はスレッド間で共有されないため、バックグラウンドスレッドが EventBuffer を設定しても呼び出し元スレッドに影響はありません。しかし caller_runs というフォールバックポリシーはその前提を崩すため、スレッド境界を越えて副作用が生じていました。修正は execute_or_skip 内の mutex のロック/アンロックと同じスコープで IsolatedExecutionState の値を管理することで、呼び出しコンテキストを問わず安全に動作させています。

ensure ブロックでの復元はエラーが発生した場合にも確実に実行されるため、例外経路でのスレッド汚染も防止されます。

まとめ

本PRは、非同期実行の前提が崩れた際に通知システムが静かに壊れるという診断困難なバグを、2行の save/restore で修正しています。IsolatedExecutionState を変更する箇所は常にその変更をスコープ内に閉じ込めるべきという教訓を、この変更は明確に示しています。