Merge pull request #345 from DataDog/feature/improve_rails_middleware…

…_error_handling

Improve Rails middleware error handling

Rakefile

@@ -258,6 +258,13 @@ task :ci do
     sh 'rvm $RAILS5_VERSIONS --verbose do appraisal rails5-postgres rake test:railsdisableenv'
     # RSpec
     sh 'rvm $LAST_STABLE --verbose do rake benchmark'
+    sh 'rvm $RAILS3_VERSIONS --verbose do appraisal rails30-postgres rake spec:rails'
+    sh 'rvm $RAILS3_VERSIONS --verbose do appraisal rails32-mysql2 rake spec:rails'
+    sh 'rvm $RAILS3_VERSIONS --verbose do appraisal rails32-postgres rake spec:rails'
+    sh 'rvm $RAILS4_VERSIONS --verbose do appraisal rails4-mysql2 rake spec:rails'
+    sh 'rvm $RAILS4_VERSIONS --verbose do appraisal rails4-postgres rake spec:rails'
+    sh 'rvm $RAILS5_VERSIONS --verbose do appraisal rails5-mysql2 rake spec:rails'
+    sh 'rvm $RAILS5_VERSIONS --verbose do appraisal rails5-postgres rake spec:rails'
   else
     puts 'Too many workers than parallel tasks'
   end

lib/ddtrace/contrib/rails/action_controller.rb

@@ -50,20 +50,14 @@ def self.finish_processing(payload)
             span.set_tag('rails.route.action', payload.fetch(:action))
             span.set_tag('rails.route.controller', payload.fetch(:controller))
 
-            if payload[:exception].nil?
+            exception = payload[:exception_object]
+            if exception.nil?
               # [christian] in some cases :status is not defined,
               # rather than firing an error, simply acknowledge we don't know it.
               status = payload.fetch(:status, '?').to_s
               span.status = 1 if status.starts_with?('5')
-            else
-              error = payload[:exception]
-              if defined?(::ActionDispatch::ExceptionWrapper)
-                status = ::ActionDispatch::ExceptionWrapper.status_code_for_exception(error[0])
-                status = status ? status.to_s : '?'
-              else
-                status = '500'
-              end
-              span.set_error(error) if status.starts_with?('5')
+            elsif Utils.exception_is_error?(exception)
+              span.set_error(exception)
             end
           ensure
             span.finish()

lib/ddtrace/contrib/rails/middlewares.rb

@@ -22,8 +22,13 @@ def call(env)
         # SignalException::Interrupt would still bubble up.
         rescue Exception => e
           tracer = Datadog.configuration[:rails][:tracer]
-          span = tracer.active_span()
-          span.set_error(e) unless span.nil?
+          span = tracer.active_span
+          unless span.nil?
+            # Only set error if it's supposed to be flagged as such
+            # e.g. we don't want to flag 404s.
+            # You can add custom errors via `config.action_dispatch.rescue_responses`
+            span.set_error(e) if Utils.exception_is_error?(e)
+          end
           raise e
         end
       end

lib/ddtrace/contrib/rails/railtie.rb

@@ -6,7 +6,10 @@ module Datadog
   # Railtie class initializes
   class Railtie < Rails::Railtie
     config.app_middleware.insert_before(0, Datadog::Contrib::Rack::TraceMiddleware)
-    config.app_middleware.use(Datadog::Contrib::Rails::ExceptionMiddleware)
+    # Insert right after Rails exception handling middleware, because if it's before,
+    # it catches and swallows the error. If it's too far after, custom middleware can find itself
+    # between, and raise exceptions that don't end up getting tagged on the request properly (e.g lost stack trace.)
+    config.app_middleware.insert_after(ActionDispatch::ShowExceptions, Datadog::Contrib::Rails::ExceptionMiddleware)
 
     config.after_initialize do
       Datadog::Contrib::Rails::Framework.setup

lib/ddtrace/contrib/rails/utils.rb

@@ -72,6 +72,18 @@ def self.connection_config
           end
         end
 
+        def self.exception_is_error?(exception)
+          if defined?(::ActionDispatch::ExceptionWrapper)
+            # Gets the equivalent status code for the exception (not all are 5XX)
+            # You can add custom errors via `config.action_dispatch.rescue_responses`
+            status = ::ActionDispatch::ExceptionWrapper.status_code_for_exception(exception.class.name)
+            # Only 5XX exceptions are actually errors (e.g. don't flag 404s)
+            status.to_s.starts_with?('5')
+          else
+            true
+          end
+        end
+
         private_class_method :connection_config
       end
     end

spec/ddtrace/contrib/rails/middleware_spec.rb

@@ -0,0 +1,218 @@
+require 'ddtrace/contrib/rails/rails_helper'
+
+RSpec.describe 'Rails request' do
+  include Rack::Test::Methods
+  include_context 'Rails test application'
+
+  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
+
+  let(:tracer) { ::Datadog::Tracer.new(writer: FauxWriter.new) }
+
+  def all_spans
+    tracer.writer.spans(:keep)
+  end
+
+  before(:each) do
+    Datadog.configure do |c|
+      c.use :rack, tracer: tracer
+      c.use :rails, tracer: tracer
+    end
+  end
+
+  context 'with middleware' do
+    before(:each) { get '/' }
+
+    let(:rails_middleware) { [middleware] }
+
+    context 'that raises an exception' do
+      let(:middleware) do
+        stub_const('RaiseExceptionMiddleware', Class.new do
+          def initialize(app)
+            @app = app
+          end
+
+          def call(env)
+            @app.call(env)
+            raise NotImplementedError.new
+          end
+        end)
+      end
+
+      it do
+        expect(last_response).to be_server_error
+        expect(all_spans.length).to be >= 2
+      end
+
+      context 'rack span' do
+        subject(:span) { all_spans.first }
+
+        it do
+          expect(span.name).to eq('rack.request')
+          expect(span.span_type).to eq('http')
+          expect(span.resource).to eq('TestController#index')
+          expect(span.get_tag('http.url')).to eq('/')
+          expect(span.get_tag('http.method')).to eq('GET')
+          expect(span.get_tag('http.status_code')).to eq('500')
+          expect(span.get_tag('error.type')).to eq('NotImplementedError')
+          expect(span.get_tag('error.msg')).to eq('NotImplementedError')
+          expect(span.status).to eq(Datadog::Ext::Errors::STATUS)
+          expect(span.get_tag('error.stack')).to_not be nil
+        end
+      end
+    end
+
+    context 'that raises a known NotFound exception' do
+      let(:middleware) do
+        stub_const('RaiseNotFoundMiddleware', Class.new do
+          def initialize(app)
+            @app = app
+          end
+
+          def call(env)
+            @app.call(env)
+            raise ActionController::RoutingError.new('/missing_route')
+          end
+        end)
+      end
+
+      it do
+        expect(last_response).to be_not_found
+        expect(all_spans.length).to be >= 2
+      end
+
+      context 'rack span' do
+        subject(:span) { all_spans.first }
+
+        it do
+          expect(span.name).to eq('rack.request')
+          expect(span.span_type).to eq('http')
+          expect(span.resource).to eq('TestController#index')
+          expect(span.get_tag('http.url')).to eq('/')
+          expect(span.get_tag('http.method')).to eq('GET')
+          expect(span.get_tag('http.status_code')).to eq('404')
+
+          if Rails.version >= '3.2'
+            expect(span.get_tag('error.type')).to be nil
+            expect(span.get_tag('error.msg')).to be nil
+            expect(span.status).to_not eq(Datadog::Ext::Errors::STATUS)
+            expect(span.get_tag('error.stack')).to be nil
+          else
+            # Rails 3.0 raises errors for 404 routing errors
+            expect(span.get_tag('error.type')).to eq('ActionController::RoutingError')
+            expect(span.get_tag('error.msg')).to eq('/missing_route')
+            expect(span.status).to eq(Datadog::Ext::Errors::STATUS)
+            expect(span.get_tag('error.stack')).to_not be nil
+          end
+        end
+      end
+    end
+
+    context 'that raises a custom exception' do
+      let(:error_class) do
+        stub_const('CustomError', Class.new(StandardError) do
+          def message
+            'Custom error message!'
+          end
+        end)
+      end
+
+      let(:middleware) do
+        # Run this to define the error class
+        error_class
+
+        stub_const('RaiseCustomErrorMiddleware', Class.new do
+          def initialize(app)
+            @app = app
+          end
+
+          def call(env)
+            @app.call(env)
+            raise CustomError.new
+          end
+        end)
+      end
+
+      it do
+        expect(last_response).to be_server_error
+        expect(all_spans.length).to be >= 2
+      end
+
+      context 'rack span' do
+        subject(:span) { all_spans.first }
+
+        it do
+          expect(span.name).to eq('rack.request')
+          expect(span.span_type).to eq('http')
+          expect(span.resource).to eq('TestController#index')
+
+          if Rails.version >= '3.2'
+            expect(span.get_tag('http.url')).to eq('/')
+          else
+          end
+
+          expect(span.get_tag('http.method')).to eq('GET')
+          expect(span.get_tag('http.status_code')).to eq('500')
+          expect(span.get_tag('error.type')).to eq('CustomError')
+          expect(span.get_tag('error.msg')).to eq('Custom error message!')
+          expect(span.status).to eq(Datadog::Ext::Errors::STATUS)
+          expect(span.get_tag('error.stack')).to_not be nil
+        end
+      end
+
+      if Rails.version >= '3.2'
+        context 'that is flagged as a custom 404' do
+          # TODO: Make a cleaner API for injecting into Rails application configuration
+          let(:initialize_block) do
+            super_block = super()
+            Proc.new do
+              self.instance_exec(&super_block)
+              config.action_dispatch.rescue_responses.merge!(
+                'CustomError' => :not_found
+              )
+            end
+          end
+
+          after(:each) do
+            # Be sure to delete configuration after, so it doesn't carry over to other examples.
+            # TODO: Clear this configuration automatically via rails_helper shared examples
+            ActionDispatch::Railtie.config.action_dispatch.rescue_responses.delete('CustomError')
+            ActionDispatch::ExceptionWrapper.class_variable_get(:@@rescue_responses).tap do |resps|
+              resps.delete('CustomError')
+            end
+          end
+
+          it do
+            expect(last_response).to be_not_found
+            expect(all_spans.length).to be >= 2
+          end
+
+          context 'rack span' do
+            subject(:span) { all_spans.first }
+
+            it do
+              expect(span.name).to eq('rack.request')
+              expect(span.span_type).to eq('http')
+              expect(span.resource).to eq('TestController#index')
+              expect(span.get_tag('http.url')).to eq('/')
+              expect(span.get_tag('http.method')).to eq('GET')
+              expect(span.get_tag('http.status_code')).to eq('404')
+              expect(span.get_tag('error.type')).to be nil
+              expect(span.get_tag('error.msg')).to be nil
+              expect(span.status).to_not eq(Datadog::Ext::Errors::STATUS)
+              expect(span.get_tag('error.stack')).to be nil
+            end
+          end
+        end
+      end
+    end
+  end
+end