Hooks into Coaches ActiveSupport notification to instrument middleware using Skylight

skylight_middleware_logger.rb

# Coach emits profiling notifications via ActiveSupport. SkylightMiddlewareLogger
# listens for these events and logs request data via Skylight.
class SkylightMiddlewareLogger
  def self.log!
    @instance ||= new.log!
  end

  def log!
    # Tell Coach to start issuing ActiveSupport::Notifications
    Coach::Notifications.subscribe!

    subscribe('handler.start')      { |*_, event| log_handler_start(event) }
    subscribe('middleware.start')   { |*_, event| log_middleware_start(event) }
    subscribe('middleware.finish')  { |*_, event| log_middleware_finish(event) }
    subscribe('handler.finish')     { |*_, event| log_handler_finish(event) }

    self
  end

  private_class_method :new

  private

  def initialize
    @requests = {}
  end

  def subscribe(event, &block)
    ActiveSupport::Notifications.subscribe("coach.#{event}", &block)
  end

  # Fetches the current trace from Skylight gem
  def skylight_trace
    Skylight::Instrumenter.instance.try(:current_trace)
  end

  def traces_for_request(event)
    @requests[event[:request].uuid]
  end

  def log_handler_start(event)
    trace = skylight_trace
    return unless trace.present?

    # Skylight requires setting of the endpoint name in order to create a new request
    # trace.
    trace.endpoint = event[:middleware]
    @requests[event[:request].uuid] = {
      exec: trace.instrument('app.block', 'Exec'),  # entire chain trace
      chain_traces: {}  # collection of traces for each middleware
    }
  end

  def log_middleware_start(event)
    trace = skylight_trace
    return unless trace.present?

    # Create and store a Skylight trace for this middleware under the current request.
    # This allows retrieval and termination of this request on middleware.finish.
    middleware_trace = trace.instrument('app.block', event[:middleware])
    traces_for_request(event)[:chain_traces][event[:middleware]] = middleware_trace
  end

  def log_middleware_finish(event)
    trace = skylight_trace
    return unless trace.present?

    # Terminate this middlewares trace and remove from the request chain
    chain_traces = traces_for_request(event)[:chain_traces]
    trace.done(chain_traces.delete(event[:middleware]))
  end

  def log_handler_finish(event)
    trace = skylight_trace
    return unless trace.present?

    # Terminate the request trace and remove from the request cache
    trace.done(@requests.delete(event[:request].uuid)[:exec])
  end
end