`execute_or_skip`実行後にインストゥルメンターが汚染される問題を修正

FutureResult#execute_or_skipがcaller_runsフォールバックポリシーによって呼び出しスレッド上で実行された際に、スレッドのactive_record_instrumenterがEventBufferで置き換えられたまま復元されないバグが修正されました。これにより、非同期クエリ実行後もsql.active_record通知が正常に配信されるようになります。

背景

async_query_executorを:global_thread_poolに設定して非同期クエリを使用するアプリケーションで、一部リクエストの同期クエリがインストゥルメント・ログ記録されなくなる問題が確認されていました。

原因はグローバルスレッドプールの飽和にあります。global_thread_poolのfallback_policyは:caller_runsに設定されており、スレッドプールのキャパシティが上限に達した際にはタスクを呼び出しスレッドで実行します。

caller_runsによって呼び出しスレッド（リクエストスレッド）上でタスクが実行されると問題が顕在化します。execute_or_skipはActiveSupport::IsolatedExecutionState[:active_record_instrumenter]をEventBufferに差し替えますが、元のインストゥルメンターを復元しないため、スレッドのIsolatedExecutionStateに古いEventBufferが永続します。Railsはリクエスト間でこのキーをクリアしないため、そのスレッドが生きている限り以降のsql.active_record通知がすべてEventBufferに飲み込まれ、購読者には配信されなくなります。

技術的な変更

future_result.rbのexecute_or_skipメソッドに、インストゥルメンターの保存と復元処理が追加されました。変更はわずか2行で、既存のensureブロックを活用しています。

変更前:

@pool.with_connection do |connection|
  return unless @mutex.try_lock
  begin
    if pending?
      @event_buffer = EventBuffer.new(self, @instrumenter)
      # ...
      execute_query(connection, async: true)
    end
  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, @instrumenter)
      # ...
      execute_query(connection, async: true)
    end
  ensure
    ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] = previous_instrumenter
    @mutex.unlock
  end
end

ミューテックスのロック取得後に即座にprevious_instrumenterを保存し、ensureブロックでミューテックスのアンロック前に復元しています。バックグラウンドスレッドで実行される通常ケースではnilを保存してnilに戻すだけであり、実質的にno-opです。

テストではschedule_queryメソッドをモンキーパッチする手法でcaller_runsフォールバックを再現しています。pool.singleton_class上でメソッドを差し替えてexecute_or_skipを呼び出しスレッドで直接実行させ、その後のPost.countがsql.active_record通知を正常に発火するかをActiveSupport::Notifications.subscribedで検証しています。

設計判断

EventBufferの有効期間をensureブロックで厳密に囲むアプローチが採用されました。EventBufferはもともとバックグラウンドスレッドでのみ使用される前提で設計されており、スレッドローカルな状態の復元は不要とみなされていました。今回の修正は、その前提がcaller_runsによって崩れる場合があるという事実を受けて、最小限の変更で安全性を確保する方針をとっています。

インストゥルメンターの保存・復元をミューテックスのロック取得直後に行うことで、EventBufferが有効な期間とIsolatedExecutionStateへの書き込み期間を一致させています。これにより、caller_runsでも通常のバックグラウンドスレッド実行でも同一のコードパスが正しく動作する対称的な設計になっています。

まとめ

この修正は、非同期クエリとスレッドプール飽和が重なるという本番環境でのみ再現する複合的な条件を、2行のコード追加で解消しています。IsolatedExecutionStateのような「スレッドローカルに見える」状態を一時的に変更するコードは、caller_runsのようなポリシーによってスレッドの境界が曖昧になる場面でも正しく動作するよう、必ず保存と復元をセットで実装するという教訓が凝縮された変更です。