worker_extensions.rb

module WorkerExtensions

  WorkerProcess = Struct.new(:id, :name, :app_name, :size) do
    ONE_X_WORKER_SIZE = 512.0
    TWO_X_WORKER_SIZE = 1024.0

    def memory_ceiling
      size == 1 ? ONE_X_WORKER_SIZE : TWO_X_WORKER_SIZE
    end
  end

  WORKER_PROCESS_ATTRIBUTES = ['id', 'process', 'app_name', 'size'].freeze

  def initialize(options={})
    super

    # fetch the dyno process name ie. export_worker.1
    process_id = EnvironmentHelper.string('DYNO')

    if process_id && EnvironmentHelper.string('HEROKU_API_ACCESS_KEY')
      heroku = ::Heroku::API.new(api_key: EnvironmentHelper.string('HEROKU_API_ACCESS_KEY'))

      # Heroku does not support querying for a given process by id :(
      processes_response = heroku.get_ps(EnvironmentHelper.string('APP_NAME'))

      if processes_response.data[:body].present?
        worker_context = processes_response.data[:body].find { |process_obj| process_obj['process'] == process_id }
        @process_context = WorkerProcess.new(*worker_context.slice(*WORKER_PROCESS_ATTRIBUTES).values) if worker_context
      else
        Rails.logger.warn("Unable to fetch worker process context for #{process_id}")
      end
    end
  rescue Heroku::API::Errors::RequestFailed => error
    ExceptionReporter.report(error, dyno: process_id)
  end

  def process_context
    @process_context
  end

  # If we were unable to fetch the process_context, just let workers live.
  def exceeds_memory_ceiling?(sample)
    return false unless process_context
    sample > memory_ceiling
  end

  def memory_ceiling
    process_context.memory_ceiling
  end
end