Trace Rails controllers that short-circuit with before_action

lib/ddtrace/contrib/rails/action_controller.rb

@@ -8,15 +8,6 @@ module Contrib
     module Rails
       # Code used to create and handle 'rails.action_controller' spans.
       module ActionController
-        include Datadog::Patcher
-
-        def self.instrument
-          # patch Rails core components
-          do_once(:instrument) do
-            Datadog::RailsActionPatcher.patch_action_controller
-          end
-        end
-
         def self.start_processing(payload)
           # trace the execution
           tracer = Datadog.configuration[:rails][:tracer]

lib/ddtrace/contrib/rails/action_controller_patch.rb

@@ -4,6 +4,8 @@ module Rails
       # Instrument ActiveController processing
       module ActionControllerPatch
         def self.included(base)
+          return if base.ancestors.include?(ProcessActionPatch)
+
           if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('2.0.0')
             base.send(:prepend, ProcessActionPatch)
           else

lib/ddtrace/contrib/rails/active_support.rb

@@ -6,15 +6,6 @@ module Contrib
     module Rails
       # Code used to create and handle 'rails.cache' spans.
       module ActiveSupport
-        include Datadog::Patcher
-
-        def self.instrument
-          do_once(:instrument) do
-            # patch Rails core components
-            Datadog::RailsCachePatcher.patch_cache_store
-          end
-        end
-
         def self.start_trace_cache(payload)
           tracer = Datadog.configuration[:rails][:tracer]
 

lib/ddtrace/contrib/rails/core_extensions.rb

@@ -181,7 +181,13 @@ def patch_process_action
       do_once(:patch_process_action) do
         require 'ddtrace/contrib/rails/action_controller_patch'
 
-        ::ActionController::Metal.send(:include, Datadog::Contrib::Rails::ActionControllerPatch)
+        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
       end
     end
   end

lib/ddtrace/contrib/rails/framework.rb

@@ -8,7 +8,6 @@
 require 'ddtrace/contrib/rails/ext'
 require 'ddtrace/contrib/rails/core_extensions'
 require 'ddtrace/contrib/rails/action_controller'
-require 'ddtrace/contrib/rails/action_view'
 require 'ddtrace/contrib/rails/active_support'
 require 'ddtrace/contrib/rails/utils'
 

lib/ddtrace/contrib/rails/patcher.rb

@@ -76,9 +76,9 @@ def setup_tracer
 
         # Add instrumentation to Rails components
         def instrument_rails
-          Datadog::Contrib::Rails::ActionController.instrument
-          Datadog::Contrib::Rails::ActionView.instrument
-          Datadog::Contrib::Rails::ActiveSupport.instrument
+          Datadog::RailsActionPatcher.patch_action_controller
+          Datadog::RailsRendererPatcher.patch_renderer
+          Datadog::RailsCachePatcher.patch_cache_store
         end
       end
     end

spec/ddtrace/contrib/rails/action_controller_spec.rb

@@ -5,7 +5,6 @@
   let(:rails_options) { { tracer: tracer } }
 
   before(:each) do
-    Datadog::RailsActionPatcher.patch_action_controller
     Datadog.configure do |c|
       c.use :rails, rails_options
     end
@@ -28,7 +27,7 @@ def index
   describe '#action' do
     subject(:result) { action.call(env) }
     let(:action) { controller.action(name) }
-    let(:env) { {} }
+    let(:env) { Rack::MockRequest.env_for('/test', {}) }
 
     shared_examples_for 'a successful dispatch' do
       it do
@@ -41,6 +40,38 @@ def index
     end
 
     describe 'for a controller' do
+      context 'that inherits from ActionController::Base' do
+        let(:base_class) { ActionController::Base }
+
+        context 'which halts an action during a #before_action' do
+          let(:controller) do
+            super().tap do |controller_class|
+              controller_class.class_eval do
+                # 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
+                end
+              end
+            end
+          end
+
+          it do
+            expect { result }.to_not raise_error
+            expect(result).to be_a_kind_of(Array)
+            expect(result).to have(3).items
+            expect(result.first).to eq(204) # Expect "No Content"
+            expect(all_spans).to have(1).items
+            expect(all_spans.first.name).to eq('rails.action_controller')
+          end
+        end
+      end
+
       context 'that inherits from ActionController::Metal' do
         let(:base_class) { ActionController::Metal }
 

spec/ddtrace/contrib/rails/railtie_spec.rb

@@ -9,17 +9,6 @@
 
   let(:tracer) { get_test_tracer }
 
-  let(:routes) { { '/' => 'test#index' } }
-  let(:controllers) { [controller] }
-
-  let(:controller) do
-    stub_const('TestController', Class.new(ActionController::Base) do
-      def index
-        head :ok
-      end
-    end)
-  end
-
   RSpec::Matchers.define :have_kind_of_middleware do |expected|
     match do |actual|
       found = 0