Logjam Agent middleware for Sidekiq

sidekiq_middleware.rb

module Logjam
  module Sidekiq
    class Middleware
      def initialize(options={})
        TimeBandits.reset
        LogjamAgent::Forwarders.reset
      end

      def call(worker, msg, queue)
        # ugly, but makes for much nicer namespace and method names in logjam
        job_name = msg["wrapped"]
        if job_name.starts_with?("ActionMailer")
          args = msg["args"].first["arguments"]
          job_name = args[0,2].join('#')
        else
          job_name += "#perform"
        end

        TimeBandits.reset
        Thread.current.thread_variable_set(:time_bandits_completed_info, nil)

        start_time = Time.current
        start_request(start_time: start_time)

        code = 200

        yield

      rescue Exception
        code = 500
        raise
      ensure
        if completed_info = Thread.current.thread_variable_get(:time_bandits_completed_info)
          _, additions, view_time, _ = completed_info
        end

        finish_request(
          action: "Sidekiq::" + job_name,
          code: code,
          view_time: view_time || 0.0,
          total_time: elapsed(start_time),
        )
      end

      private
      def elapsed(start)
        ((Time.current - start) * 1000.0).to_i
      end

      def start_request(fields)
        LogjamAgent.start_request(LogjamAgent.application_name, LogjamAgent.environment_name, fields)
      end

      def finish_request(fields)
        # Workaround bug in TimeBandits. If consumed() isn't called the metrics aren't read.
        TimeBandits::TimeConsumers::Database.instance.consumed

        LogjamAgent.finish_request(fields.merge(TimeBandits.metrics))
      end
    end
  end
end