Rack::Timeout::RequestTimeoutException の仕組み

Ruby の webserver に使われるミドルウェアである Rack に、timeout を管理する拡張 gem の rack-timeout がある。その仕事として例えば、レスポンスを返すまでに指定した秒数以上の時間がかかる場合、その秒数で Rack::Timeout::RequestTimeoutException を例外として発生させる。

github.com

タイムアウトすると聞いただけでは、処理を中断させて例外を raise しているんだろうと想像できる。そのためタイムアウトした処理の場所や環境が取得できないように思えるが、その例外の backtrace にはきちんとタイムアウトした処理の行の backtrace が載っている。どうせRubyなんだしできるんだろうなと思っていたけど、直感で方法が思いつかなかったのでコード読んでみた*1。

rack-timeout/lib/rack/timeout/core.rb

timeout = RT::Scheduler::Timeout.new do |app_thread|
  register_state_change.call :timed_out
  app_thread.raise(RequestTimeoutException.new(env))
end

response = timeout.timeout(info.timeout) do
  begin  @app.call(env)
  rescue RequestTimeoutException => e
    raise RequestTimeoutError.new(env), e.message, e.backtrace
  ensure
    register_state_change.call :completed
  end
end

github

この辺読むだけで良かった、RT::Scheduler::Timeout.new 後にアプリケーションのスレッドで RequestTimeoutException を raise しているのがわかる。 Thread#raise はそのスレッドで例外を発生させる。しかも backtrace はそのスレッドが処理をしていた行となるので完全に便利。つまるところアプリケーションの本来の処理を行うスレッドと別のスレッドで秒数をカウントし、秒数がオーバーしたらアプリケーションのスレッドに向けて例外を発生させるだけというシンプルなものだった。

ちなみにRT::Scheduler::Timeout は下のようなコードになっていて、Thread.currentでアプリケーションスレッドを参照し、 @scheduler で時間に関する処理をしつつ、タイムアウト時には @on_timeout.call を呼んでいる。@scheduler の中身に関しては本質的ではなかったので触れない。

rack-timeout/lib/rack/timeout/support/timeout.rb

ON_TIMEOUT = ->thr { thr.raise Error, "execution expired" }

def initialize(&on_timeout)
  @on_timeout = on_timeout || ON_TIMEOUT
  @scheduler  = Rack::Timeout::Scheduler.singleton
end

def timeout(secs, &block)
  return block.call if secs.nil? || secs.zero?
  thr = Thread.current
  job = @scheduler.run_in(secs) { @on_timeout.call thr }
  return block.call
ensure
  job.cancel! if job
end

github

ちなみに簡素な作りで良ければこういう感じに書ける。

class TimeoutError < StandardError; end

def call_with_timeout(time, &block)
  thr = Thread.current

  Thread.new do
    sleep(time)
    thr.raise(TimeoutError.new)
  end

  return block.call
end

begin
  call_with_timeout(2) do
    3.times do
      sleep(1)
      puts "Hello Thread!"
    end
  end
rescue TimeoutError => e
  puts "timeout!!"
end

実行結果

$ ruby call_with_timeout.rb
Hello Thread!
timeout!!

*1:この記事のコード例は github から引用し、不要な部分を削減しているため注意