Changed: Apply patch on Metal inherited instead of to Base and API.

DataDog · Jan 8, 2019 · 047e43a · 047e43a
1 parent 2b46778
commit 047e43a
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 8 deletions.
diff --git a/lib/ddtrace/contrib/rails/core_extensions.rb b/lib/ddtrace/contrib/rails/core_extensions.rb
@@ -169,6 +169,22 @@ def with_datadog_span(span)
   module RailsActionPatcher
     include Datadog::Patcher
 
+    # Patch for ActionController::Metal
+    module MetalPatch
+      # For any Controller class that inherits from ActionController::Metal,
+      # make sure to include the tracing patch. We do this instead of Metal.include(Patch)
+      # because Base and API compose modules that would take precedence over the tracing patch.
+      # By adding it onto the inheriting class, we can make sure it has a higher precedence.
+      def inherited(base)
+        super
+
+        # Make sure not to include the patch twice, to avoid double measurements on controllers.
+        unless base.ancestors.include?(Datadog::Contrib::Rails::ActionControllerPatch)
+          base.send(:include, Datadog::Contrib::Rails::ActionControllerPatch)
+        end
+      end
+    end
+
     module_function
 
     def patch_action_controller
@@ -181,13 +197,7 @@ def patch_process_action
       do_once(:patch_process_action) do
         require 'ddtrace/contrib/rails/action_controller_patch'
 
-        if defined?(::ActionController::Base)
-          ::ActionController::Base.send(:include, Datadog::Contrib::Rails::ActionControllerPatch)
-        end
-
-        if defined?(::ActionController::API)
-          ::ActionController::API.send(:include, Datadog::Contrib::Rails::ActionControllerPatch)
-        end
+        ::ActionController::Metal.extend(MetalPatch)
       end
     end
   end

diff --git a/spec/ddtrace/contrib/rails/action_controller_spec.rb b/spec/ddtrace/contrib/rails/action_controller_spec.rb
@@ -48,7 +48,12 @@ def index
           let(:controller) do
             super().tap do |controller_class|
               controller_class.class_eval do
-                before_action :short_circuit
+                # Rails 3.x-5.x compatibility
+                if respond_to?(:before_action)
+                  before_action :short_circuit
+                else
+                  before_filter :short_circuit
+                end
 
                 def short_circuit
                   head :no_content

diff --git a/spec/ddtrace/contrib/rails/middleware_spec.rb b/spec/ddtrace/contrib/rails/middleware_spec.rb
@@ -32,6 +32,7 @@ def all_spans
   end
 
   before(:each) do
+    Datadog::RailsActionPatcher.patch_action_controller
     Datadog.configure do |c|
       c.use :rack, rack_options if use_rack
       c.use :rails, rails_options if use_rails