Start measuring request sooner when it reaches the Rails engine

Author: alexisbernard

lib/rorvswild/plugin/middleware.rb

@@ -2,7 +2,7 @@
 
 module RorVsWild
   module Plugin
-    class Middleware
+    module Middleware
       module RequestQueueTime
 
         ACCEPTABLE_HEADERS = [
@@ -49,24 +49,19 @@ def parse_timestamp(timestamp)
       @installed = false
 
       def self.setup(agent)
-        return if @installed
-        Rails.application.config.middleware.unshift(RorVsWild::Plugin::Middleware, nil) if defined?(Rails)
+        return if @installed || !defined?(Rails::Engine)
+        Rails::Engine.prepend(self)
         @installed = true
       end
 
-      def initialize(app, config)
-        @app, @config = app, config
-      end
-
       def call(env)
         execution = RorVsWild::Execution::Request.new(env["ORIGINAL_FULLPATH"])
         execution.add_queue_time(calculate_queue_time(env))
         RorVsWild.agent.start_execution(execution)
         section = RorVsWild::Section.start
-        section.file, section.line = rails_engine_location
+        section.file, section.line = method(:call).super_method.source_location
         section.commands << "Rails::Engine#call"
-        code, headers, body = @app.call(env)
-        [code, headers, body]
+        super
       ensure
         RorVsWild::Section.stop
         RorVsWild.agent.stop_execution

test/plugin/middleware_test.rb

@@ -52,11 +52,15 @@ def unix_timestamp_seconds
     Time.now.to_f
   end
 
+  class App
+    def call(env)
+    end
+
+    prepend(RorVsWild::Plugin::Middleware)
+  end
+
   def middleware
     agent # Load agent
-    app = mock(call: nil)
-    middleware = RorVsWild::Plugin::Middleware.new(app, nil)
-    middleware.stubs(rails_engine_location: ["/rails/lib/engine.rb", 12])
-    middleware
+    App.new
   end
 end