lib/plain_apm/extensions/context/railtie.rb in plain_apm-0.7.1 vs lib/plain_apm/extensions/context/railtie.rb in plain_apm-0.8.0

@@ -7,15 +7,21 @@
 #
 # See LICENSE.txt in the current directory for the license.
 module PlainApm
   module Extensions
     module Context
-      class Railtie < Rails::Railtie
-        initializer "plain_apm.insert_context_middleware" do |app|
-          if defined?(ActionDispatch::RequestId)
-            app.config.middleware.insert_after ActionDispatch::RequestId, PlainApm::Extensions::Context::Middleware
-          else
-            app.config.middleware.insert_after Rack::MethodOverride, PlainApm::Extensions::Context::Middleware
+      if defined?(Rails::Railtie)
+        class Railtie < Rails::Railtie
+          initializer "plain_apm.initialize_context" do |app|
+            ActiveSupport.on_load(:active_job) do |klass|
+              klass.prepend(PlainApm::Extensions::Context::ActiveJob)
+            end
+
+            if defined?(ActionDispatch::RequestId)
+              app.config.middleware.insert_after ActionDispatch::RequestId, PlainApm::Extensions::Context::Rack
+            else
+              app.config.middleware.insert_after Rack::MethodOverride, PlainApm::Extensions::Context::Rack
+            end
           end
         end
       end
     end
   end