shopify_deadlines.rb

class Deadline
  def self.current=(value)
    ActiveRecord::Base.deadline = value
  end

  def initialize(duration_seconds)
    @deadline = now + duration_seconds.to_i
  end

  def extend_deadline(duration_seconds)
    @deadline += duration_seconds
  end

  def exceeded?
    now >= @deadline
  end

  def time_left_seconds
    [@deadline - now, 0].max.to_f
  end

  def time_left_ms
    (time_left_seconds * 1000).to_i
  end

  private

  def now
    Process.clock_gettime(Process::CLOCK_MONOTONIC).to_f
  end
end

module ActiveRecord
  class DeadlineExceededError < ActiveRecord::ActiveRecordError
    def initialize
      super("The query was cancelled because the request or a job timeout has been hit")
    end
  end

  class Base
    def self.deadline
      Thread.current.thread_variable_get(:active_record_deadline)
    end

    def self.deadline=(value)
      Thread.current.thread_variable_set(:active_record_deadline, value)
    end
  end

  module DeadlinePatch
    def execute(sql, name = nil)
      if ActiveRecord::Base.deadline && ActiveRecord::Base.deadline.exceeded?
        raise DeadlineExceededError
      end
      super
    end
  end
end

ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(ActiveRecord::DeadlinePatch)

class DeadlineMiddleware
  # We want the deadline to kick off slightly before the Unicorn killer will activate
  DEADLINE_SECONDS = UnicornKiller.default_timeout - 2

  def initialize(app)
    @app = app
  end

  def call(env)
    deadline = Deadline.new(DEADLINE_SECONDS)
    env['deadline'] = deadline
    Deadline.current = deadline
    @app.call(env)
  ensure
    env.delete('deadline')
    Deadline.current = nil
  end
end