unicornのタイムアウト時にもRailsのログをちゃんと出力させる

unicornはconfで

timeout 20

とかやっとくと、20秒以上かかったらそのworkerが殺される。それはいい。問題はその殺され方にあって、タイムアウトしたunicorn workerはmasterにKILLシグナルで強制的に殺される。KILLで殺されてしまうと、worker側でtrapして安全に終了処理をすることができない。
一番困るのは、Railsのloggerは（production環境のデフォルトだと）リクエストが終了するまでバッファリングしているので、リクエストの途中でKILLされてしまうとloggerがflushされない。つまり、unicornのタイムアウト時には、リクエストのログは一切production.logには出力されない。異常時のログが出ないとか、まじで困ると思うんだけど、みんなどうしてんだろ。

これに対処するためにはunicornのコードに手を入れるのが手っ取り早い。方法はいくつかあると思うけど、タイムアウトでunicorn workerを殺すときのシグナルをconfで指定できるようにしてみた。
Commit d0b24198af57383d1c09f7be79c7e59c161bcfe5 to cookpad/unicorn - GitHub

timeout 20
timeout_signal :INT

このtimeout_signalというのが今回unicornに追加したオプション。これを指定するとKILL以外のシグナルでworkerを殺すようになる（デフォルトはKILL）。たとえばINTならこんなふうにworker側でtrapできる。

after_fork do |server, worker|
  trap (:INT) do
    Rails.logger.error "Unicorn Timeout (trapped SIGINT)"
    Rails.logger.flush
    Process.kill :KILL, $$
  end
end

workerがINTをtrapしたら、Rails.loggerをflush（バッファに溜まっているログを吐かせる）してから、自分自身をKILLする。これでとりあえず目的は果たせた。

応用編：タイムアウト時に実行していたSQLをログに出す

タイムアウト時のログを出すだけじゃなくて、何をやっていたときにタイムアウトしたかをちゃんとログに出力したい。タイムアウトしがちなのはやはりDBのスロークエリだろう。というわけで、unicornタイムアウト時に実行していたSQLをログに出すようにしてみる。
これも色々やり方はありそうだけど、今回はmysql2_adapterにモンキーパッチをあてることにする。

::ActiveRecord::ConnectionAdapters::Mysql2Adapter.class_eval do
  def execute_with_timeout_logging(sql, name = nil)
    result = nil
    @@current_query_info = {:sql => sql, :start_time => Time.now}
    result = execute_without_timeout_logging(sql, name)
    @@current_query_info = nil
    result
  end
  alias_method_chain :execute, :timeout_logging

  def self.current_query_info
    unless defined?(@@current_query_info)
      @@current_query_info = nil
    end
    @@current_query_info
  end
end

で、さっきのunicornのconfに、ログの処理を追加する。

trap(:INT) do
  info = ActiveRecord::ConnectionAdapters::Mysql2adapter.current_query_info
  if info
    ms = ((Time.now - info[:start_time])*1000).to_i
    ::Rails.logger.error "Query Aborted (#{ms}ms) #{info[:sql]}"
  end
  Rails.logger.error "Unicorn Timeout (trapped SIGINT)"
  Rails.logger.flush
  Process.kill :KILL, $$
end

これでタイムアウトしたときに実行していたSQLがproduction.logに出るようになる。

おわりに

今回はunicornにtimeout_signalオプションを実装した。unicornは多くの部分がrubyで書かれているので、自分で簡単にいじれるから助かる。このtimeout_signalオプションをしばらくINTで運用してみてるけど、うまく動いているので、ちゃんとドキュメントとtestを書いて本家にpush requestしたいと思ってる。

注意しなくちゃいけないのは、KILL以外のシグナルでworkerを殺すようにすると、暴走したworkerがちゃんと死んでくれないという状況が想定される。暴走してるworkerを見つけてはKILLするcronを書くなり、そういう運用があわせて必要になると思う。