Add Rails middleware option

docs/GettingStarted.md

@@ -907,6 +907,7 @@ Where `options` is an optional `Hash` that accepts the following parameters:
 | ``database_service`` | Database service name used when tracing database activity | ``<app_name>-<adapter_name>`` |
 | ``exception_controller`` | Class or Module which identifies a custom exception controller class. Tracer provides improved error behavior when it can identify custom exception controllers. By default, without this option, it 'guesses' what a custom exception controller looks like. Providing this option aids this identification. | ``nil`` |
 | ``distributed_tracing`` | Enables [distributed tracing](#distributed-tracing) so that this service trace is connected with a trace of another service if tracing headers are received | `false` |
+| ``middleware`` | Add the trace middleware to the Rails application. Set to `false` if you don't want the middleware to load. | `true` |
 | ``middleware_names`` | Enables any short-circuited middleware requests to display the middleware name as resource for the trace. | `false` |
 | ``template_base_path`` | Used when the template name is parsed. If you don't store your templates in the ``views/`` folder, you may need to change this value | ``views/`` |
 | ``tracer`` | A ``Datadog::Tracer`` instance used to instrument the application. Usually you don't need to set that. | ``Datadog.tracer`` |

lib/ddtrace/contrib/rack/patcher.rb

@@ -39,7 +39,8 @@ def patch
             @patched = true
           end
 
-          if !@middleware_patched && get_option(:middleware_names)
+          if (!instance_variable_defined?(:@middleware_patched) || !@middleware_patched) \
+             && get_option(:middleware_names)
             if get_option(:application)
               enable_middleware_names
               @middleware_patched = true

lib/ddtrace/contrib/rails/patcher.rb

@@ -1,4 +1,7 @@
 require 'ddtrace/contrib/rails/utils'
+require 'ddtrace/contrib/rails/framework'
+require 'ddtrace/contrib/rails/middlewares'
+require 'ddtrace/contrib/rack/middlewares'
 
 module Datadog
   module Contrib
@@ -17,6 +20,7 @@ module Patcher
             Datadog.configuration[:active_record][:service_name] = value
           end
         end
+        option :middleware, default: true
         option :middleware_names, default: false
         option :distributed_tracing, default: false
         option :template_base_path, default: 'views/'
@@ -28,7 +32,41 @@ module Patcher
         class << self
           def patch
             return @patched if patched? || !compatible?
-            require_relative 'framework'
+
+            # Add a callback hook to add the trace middleware before the application initializes.
+            # Otherwise the middleware stack will be frozen.
+            do_once(:rails_before_initialize_hook) do
+              ::ActiveSupport.on_load(:before_initialize) do
+                # Sometimes we don't want to activate middleware e.g. OpenTracing, etc.
+                if Datadog.configuration[:rails][:middleware]
+                  # Add trace middleware
+                  config.middleware.insert_before(0, Datadog::Contrib::Rack::TraceMiddleware)
+
+                  # 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.middleware.insert_after(
+                    ActionDispatch::ShowExceptions,
+                    Datadog::Contrib::Rails::ExceptionMiddleware
+                  )
+                end
+              end
+            end
+
+            # Add a callback hook to finish configuring the tracer after the application is initialized.
+            # We need to wait for some things, like application name, middleware stack, etc.
+            do_once(:rails_after_initialize_hook) do
+              ::ActiveSupport.on_load(:after_initialize) do
+                Datadog::Contrib::Rails::Framework.setup
+
+                # Add instrumentation to Rails components
+                Datadog::Contrib::Rails::ActionController.instrument
+                Datadog::Contrib::Rails::ActionView.instrument
+                Datadog::Contrib::Rails::ActiveSupport.instrument
+              end
+            end
+
             @patched = true
           rescue => e
             Datadog::Tracer.log.error("Unable to apply Rails integration: #{e}")
@@ -49,5 +87,3 @@ def compatible?
     end
   end
 end
-
-require 'ddtrace/contrib/rails/railtie' if Datadog.registry[:rails].compatible?

lib/ddtrace/contrib/rails/railtie.rb

@@ -5,11 +5,14 @@
 module Datadog
   # Railtie class initializes
   class Railtie < Rails::Railtie
-    config.app_middleware.insert_before(0, Datadog::Contrib::Rack::TraceMiddleware)
-    # 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)
+    # Add the trace middleware to the application stack
+    initializer 'datadog.add_middleware' do |app|
+      app.middleware.insert_before(0, Datadog::Contrib::Rack::TraceMiddleware)
+      # 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.)
+      app.middleware.insert_after(ActionDispatch::ShowExceptions, Datadog::Contrib::Rails::ExceptionMiddleware)
+    end
 
     config.after_initialize do
       Datadog::Contrib::Rails::Framework.setup

spec/ddtrace/contrib/rails/railtie_spec.rb

@@ -0,0 +1,60 @@
+require 'ddtrace/contrib/rails/rails_helper'
+require 'ddtrace/contrib/rails/framework'
+require 'ddtrace/contrib/rails/middlewares'
+require 'ddtrace/contrib/rack/middlewares'
+
+RSpec.describe 'Rails application' do
+  before(:each) { skip 'Test not compatible with Rails < 4.0' if Rails.version < '4.0' }
+  include_context 'Rails test application'
+
+  let(:tracer) { ::Datadog::Tracer.new(writer: FauxWriter.new) }
+
+  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|
+      while actual
+        return true if actual.class <= expected
+        without_warnings { actual = actual.instance_variable_get(:@app) }
+      end
+      false
+    end
+  end
+
+  before(:each) do
+    Datadog.registry[:rails].instance_variable_set(:@patched, false)
+    Datadog.configure do |c|
+      c.tracer hostname: ENV.fetch('TEST_DDAGENT_HOST', 'localhost')
+      c.use :rails, rails_options if use_rails
+    end
+  end
+
+  let(:use_rails) { true }
+  let(:rails_options) { { tracer: tracer } }
+
+  describe 'with Rails integration #middleware option' do
+    context 'set to true' do
+      let(:rails_options) { super().merge(middleware: true) }
+
+      it { expect(app).to have_kind_of_middleware(Datadog::Contrib::Rack::TraceMiddleware) }
+      it { expect(app).to have_kind_of_middleware(Datadog::Contrib::Rails::ExceptionMiddleware) }
+    end
+
+    context 'set to false' do
+      let(:rails_options) { super().merge(middleware: false) }
+      after(:each) { Datadog.configuration[:rails][:middleware] = true }
+
+      it { expect(app).to_not have_kind_of_middleware(Datadog::Contrib::Rack::TraceMiddleware) }
+      it { expect(app).to_not have_kind_of_middleware(Datadog::Contrib::Rails::ExceptionMiddleware) }
+    end
+  end
+end

spec/ddtrace/contrib/rails/support/application.rb

@@ -3,6 +3,11 @@
 RSpec.shared_context 'Rails test application' do
   include_context 'Rails base application'
 
+  before do
+    Datadog.registry[:rails].instance_variable_set(:@patched, false)
+    reset_rails_configuration!
+  end
+
   let(:app) do
     initialize_app!
     rails_test_application.instance

spec/ddtrace/contrib/rails/support/rails3.rb

@@ -30,7 +30,6 @@ def inherited(base)
   include_context 'Rails models'
 
   let(:rails_base_application) do
-    reset_rails_configuration!
     during_init = initialize_block
     klass = Class.new(Rails::Application) do
       redis_cache = [:redis_store, { url: ENV['REDIS_URL'] }]
@@ -113,16 +112,10 @@ def draw_test_routes!(mapper)
   def reset_rails_configuration!
     Rails.class_variable_set(:@@application, nil)
     Rails::Application.class_variable_set(:@@instance, nil)
-    Rails::Railtie::Configuration.class_variable_set(:@@app_middleware, app_middleware)
+    if Rails::Railtie::Configuration.class_variable_defined?(:@@app_middleware)
+      Rails::Railtie::Configuration.class_variable_set(:@@app_middleware, Rails::Configuration::MiddlewareStackProxy.new)
+    end
     Rails::Railtie::Configuration.class_variable_set(:@@app_generators, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@to_prepare_blocks, nil)
   end
-
-  def app_middleware
-    current = Rails::Railtie::Configuration.class_variable_get(:@@app_middleware)
-    Datadog::Contrib::Rails::Test::Configuration.fetch(:app_middleware, current).dup.tap do |copy|
-      copy.instance_variable_set(:@operations, (copy.instance_variable_get(:@operations) || []).dup)
-      copy.instance_variable_set(:@delete_operations, (copy.instance_variable_get(:@delete_operations) || []).dup)
-    end
-  end
 end

spec/ddtrace/contrib/rails/support/rails4.rb

@@ -16,7 +16,6 @@
   include_context 'Rails models'
 
   let(:rails_base_application) do
-    reset_rails_configuration!
     klass = Class.new(Rails::Application) do
       def config.database_configuration
         parsed = super
@@ -108,16 +107,10 @@ def reset_rails_configuration!
     Rails::Railtie::Configuration.class_variable_set(:@@eager_load_namespaces, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@watchable_files, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@watchable_dirs, nil)
-    Rails::Railtie::Configuration.class_variable_set(:@@app_middleware, app_middleware)
+    if Rails::Railtie::Configuration.class_variable_defined?(:@@app_middleware)
+      Rails::Railtie::Configuration.class_variable_set(:@@app_middleware, Rails::Configuration::MiddlewareStackProxy.new)
+    end
     Rails::Railtie::Configuration.class_variable_set(:@@app_generators, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@to_prepare_blocks, nil)
   end
-
-  def app_middleware
-    current = Rails::Railtie::Configuration.class_variable_get(:@@app_middleware)
-    Datadog::Contrib::Rails::Test::Configuration.fetch(:app_middleware, current).dup.tap do |copy|
-      copy.instance_variable_set(:@operations, (copy.instance_variable_get(:@operations) || []).dup)
-      copy.instance_variable_set(:@delete_operations, (copy.instance_variable_get(:@delete_operations) || []).dup)
-    end
-  end
 end

spec/ddtrace/contrib/rails/support/rails5.rb

@@ -16,7 +16,6 @@
   include_context 'Rails models'
 
   let(:rails_base_application) do
-    reset_rails_configuration!
     klass = Class.new(Rails::Application) do
       def config.database_configuration
         parsed = super
@@ -86,16 +85,10 @@ def reset_rails_configuration!
     Rails::Railtie::Configuration.class_variable_set(:@@eager_load_namespaces, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@watchable_files, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@watchable_dirs, nil)
-    Rails::Railtie::Configuration.class_variable_set(:@@app_middleware, app_middleware)
+    if Rails::Railtie::Configuration.class_variable_defined?(:@@app_middleware)
+      Rails::Railtie::Configuration.class_variable_set(:@@app_middleware, Rails::Configuration::MiddlewareStackProxy.new)
+    end
     Rails::Railtie::Configuration.class_variable_set(:@@app_generators, nil)
     Rails::Railtie::Configuration.class_variable_set(:@@to_prepare_blocks, nil)
   end
-
-  def app_middleware
-    current = Rails::Railtie::Configuration.class_variable_get(:@@app_middleware)
-    Datadog::Contrib::Rails::Test::Configuration.fetch(:app_middleware, current).dup.tap do |copy|
-      copy.instance_variable_set(:@operations, (copy.instance_variable_get(:@operations) || []).dup)
-      copy.instance_variable_set(:@delete_operations, (copy.instance_variable_get(:@delete_operations) || []).dup)
-    end
-  end
 end