arproxy を puma と使う場合の注意点

ActiveRecord でクエリ実行時に任意の処理を挟むことができる arproxy について、puma で...というかマルチスレッドで使う場合に注意すべき点があった。

問題がある実装

class FooController < ApplicationController
  before_action { Arproxy.enable! }
  after_action { Arproxy.disable! }
end

puma ではデフォルトでマルチスレッドでリクエストを処理するので、複数スレッドで Arproxy.enable! Arproxy.disable! が呼ばれる可能性がある。

しかし、これらのメソッドは Arproxy のインスタンス変数 @proxy_chain を変更したり、DB アダプターのメソッドを書き換えたりするため複数スレッドで実行すると問題が発生する。

回避策

マルチスレッドで実行されうる箇所では Arproxy.enable! Arproxy.disable! は呼ばない。

リクエストごとに実行を切り替えたい場合は、スレッドローカル変数を使い proxy の実装で切り替えるなどする。

class AllQueryLogger < Arproxy::Base
  def execute(sql, name=nil)
    # スレッドローカル変数を参照して実行切り替え
    return super unless Thread.current[:arproxy_all_query_logger_enabled]

    warn sql
    super(sql, name)
  end
end

# デフォルトで有効にする
# Executor によりスレッド生成されたときに実行される
Rails.application.executor.to_run do
  Thread.current[:arproxy_all_query_logger_enabled] = true
end
# メインスレッドでも有効に
Thread.current[:arproxy_all_query_logger_enabled] = true

より構造化した例 ↓

require "arproxy"

module Arproxy
  module MultiThread
    def self.included(klass)
      klass.extend ClassMethods
    end

    module ClassMethods
      def enable_proxy!(klass)
        proxy_enabled[klass] = true
      end

      def disable_proxy!(klass)
        proxy_enabled[klass] = false
      end

      def enable_proxy?(klass)
        proxy_enabled[klass]
      end

      private
      def thread_local
        Thread.current[:arproxy] ||= {}
      end

      def proxy_enabled
        thread_local[:proxy_enabled] ||= Hash.new {|h, k| h[k] = true }
      end
    end
  end

  class Base
    module MultiThread
      def self.included(klass)
        klass.extend ClassMethods
        klass.prepend InstanceMethods
      end

      module ClassMethods
        def enable!
          Arproxy.enable_proxy!(self)
        end

        def disable!
          Arproxy.disable_proxy!(self)
        end

        def enable?
          Arproxy.enable_proxy?(self)
        end
      end

      module InstanceMethods
        # 次の proxy を呼ぶ前に enable? をチェックしてスキップする処理を追加
        def execute(sql, name = nil, **_kwargs)
          n = next_proxy
          loop do
            if n.class.enable?
              break n.execute sql, name
            else
              n = n.next_proxy
            end
          end
        end
      end
    end
  end
end

Arproxy.include Arproxy::MultiThread
Arproxy::Base.include Arproxy::Base::MultiThread

# 最初になにもしない (次の proxy の enable? をチェックして実行する) クラスを設定
# これがないと最初の proxy の enable? がチェックできない
class ChainHead < Arproxy::Base
end

class AllQueryLogger < Arproxy::Base
  def execute(sql, name=nil)
    warn sql
    super(sql, name)
  end
end

Arproxy.configure do |config|
  config.adapter = "mysql2"
  config.use ChainHead
  config.use AllQueryLogger
end

Rails.application.executor.to_run do
  AllQueryLogger.enable!
end
AllQueryLogger.enable!

Arproxy.enable!

# スレッド中で AllQueryLogger.(enable!|disable!) は安全に呼ぶことができる