Merge pull request #641 from DataDog/fix/rails_load_hooks_run_twice

Prevent Rails load hooks applying patch twice

lib/ddtrace/contrib/rails/patcher.rb

@@ -32,10 +32,12 @@ def patch_before_intialize
         end
 
         def before_intialize(app)
-          # Middleware must be added before the application is initialized.
-          # Otherwise the middleware stack will be frozen.
-          # Sometimes we don't want to activate middleware e.g. OpenTracing, etc.
-          add_middleware(app) if Datadog.configuration[:rails][:middleware]
+          do_once(:rails_before_initialize, for: app) do
+            # Middleware must be added before the application is initialized.
+            # Otherwise the middleware stack will be frozen.
+            # Sometimes we don't want to activate middleware e.g. OpenTracing, etc.
+            add_middleware(app) if Datadog.configuration[:rails][:middleware]
+          end
         end
 
         def add_middleware(app)
@@ -59,10 +61,12 @@ def patch_after_intialize
         end
 
         def after_intialize(app)
-          # Finish configuring the tracer after the application is initialized.
-          # We need to wait for some things, like application name, middleware stack, etc.
-          setup_tracer
-          instrument_rails
+          do_once(:rails_after_initialize, for: app) do
+            # Finish configuring the tracer after the application is initialized.
+            # We need to wait for some things, like application name, middleware stack, etc.
+            setup_tracer
+            instrument_rails
+          end
         end
 
         # Configure Rails tracing with settings

lib/ddtrace/patcher.rb

@@ -21,21 +21,23 @@ def without_warnings
         end
       end
 
-      def do_once(key = nil)
+      def do_once(key = nil, options = {})
         # If already done, don't do again
-        @done_once ||= {}
-        return @done_once[key] if @done_once.key?(key)
+        @done_once ||= Hash.new { |h, k| h[k] = {} }
+        if @done_once.key?(key) && @done_once[key].key?(options[:for])
+          return @done_once[key][options[:for]]
+        end
 
         # Otherwise 'do'
         yield.tap do
           # Then add the key so we don't do again.
-          @done_once[key] = true
+          @done_once[key][options[:for]] = true
         end
       end
 
-      def done?(key)
+      def done?(key, options = {})
         return false unless instance_variable_defined?(:@done_once)
-        !@done_once.nil? && @done_once.key?(key)
+        !@done_once.nil? && @done_once.key?(key) && @done_once[key].key?(options[:for])
       end
     end
 

spec/ddtrace/contrib/concurrent_ruby/integration_spec.rb

@@ -9,7 +9,7 @@
     example.run
     ::Concurrent.send(:remove_const, :Future)
     ::Concurrent.const_set('Future', unmodified_future)
-    Datadog.registry[:concurrent_ruby].patcher.instance_variable_set(:@done_once, {})
+    remove_patch!(:concurrent_ruby)
   end
 
   let(:tracer) { ::Datadog::Tracer.new(writer: FauxWriter.new) }

spec/ddtrace/contrib/rails/railtie_spec.rb

@@ -22,12 +22,19 @@ def index
 
   RSpec::Matchers.define :have_kind_of_middleware do |expected|
     match do |actual|
+      found = 0
       while actual
-        return true if actual.class <= expected
+        found += 1 if actual.class <= expected
         without_warnings { actual = actual.instance_variable_get(:@app) }
       end
-      false
+      found == (count || 1)
     end
+
+    chain :once do
+      @count = 1
+    end
+
+    chain :copies, :count
   end
 
   before(:each) do
@@ -44,8 +51,8 @@ def index
     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) }
+      it { expect(app).to have_kind_of_middleware(Datadog::Contrib::Rack::TraceMiddleware).once }
+      it { expect(app).to have_kind_of_middleware(Datadog::Contrib::Rails::ExceptionMiddleware).once }
     end
 
     context 'set to false' do
@@ -56,4 +63,28 @@ def index
       it { expect(app).to_not have_kind_of_middleware(Datadog::Contrib::Rails::ExceptionMiddleware) }
     end
   end
+
+  describe 'when load hooks run twice' do
+    before(:each) do
+      # Set expectations
+      expect(Datadog::Contrib::Rails::Patcher).to receive(:add_middleware)
+        .with(a_kind_of(Rails::Application))
+        .once
+        .and_call_original
+
+      without_warnings do
+        # Then load the app, which run load hooks
+        app
+
+        # Then manually re-run load hooks
+        ActiveSupport.run_load_hooks(:before_initialize, app)
+        ActiveSupport.run_load_hooks(:after_initialize, app)
+      end
+    end
+
+    it 'only includes the middleware once' do
+      expect(app).to have_kind_of_middleware(Datadog::Contrib::Rack::TraceMiddleware).once
+      expect(app).to have_kind_of_middleware(Datadog::Contrib::Rails::ExceptionMiddleware).once
+    end
+  end
 end

spec/ddtrace/patcher_spec.rb

@@ -105,20 +105,81 @@
           end
         end
       end
+
+      context 'when called with a key and :for' do
+        subject(:result) { patcher.do_once(key, for: for_key) { integration.patch } }
+
+        let(:key) { double('key') }
+        let(:for_key) { double('for key') }
+
+        it do
+          expect(integration).to receive(:patch).once.and_return(patch_result)
+          expect(result).to be(patch_result)
+        end
+
+        context 'then called a second time' do
+          context 'with a matching key and :for' do
+            context 'that is the same' do
+              subject(:result) do
+                patcher.do_once(key, for: for_key) { integration.patch }
+                patcher.do_once(key, for: for_key) { integration.patch }
+              end
+
+              it do
+                expect(integration).to receive(:patch).once.and_return(patch_result)
+                expect(result).to be true # Because second block doesn't run
+              end
+            end
+
+            context 'that is different' do
+              subject(:result) do
+                patcher.do_once(key, for: for_key) { integration.patch }
+                patcher.do_once(key, for: for_key_two) { integration.patch }
+              end
+
+              let(:for_key_two) { double('for key two') }
+
+              it do
+                expect(integration).to receive(:patch).twice.and_return(patch_result)
+                expect(result).to be(patch_result)
+              end
+            end
+          end
+        end
+      end
     end
 
     describe '#done?' do
       context 'when called before do_once' do
-        subject(:done) { patcher.done?(key) }
         let(:key) { double('key') }
-        it { is_expected.to be false }
+
+        context 'with a key' do
+          subject(:done) { patcher.done?(key) }
+          it { is_expected.to be false }
+        end
+
+        context 'with a key and :for' do
+          subject(:done) { patcher.done?(key, for: for_key) }
+          let(:for_key) { double('for key') }
+          it { is_expected.to be false }
+        end
       end
 
       context 'when called after do_once' do
-        subject(:done) { patcher.done?(key) }
         let(:key) { double('key') }
-        before(:each) { patcher.do_once(key) { 'Perform patch' } }
-        it { is_expected.to be true }
+
+        context 'with a key' do
+          subject(:done) { patcher.done?(key) }
+          before(:each) { patcher.do_once(key) { 'Perform patch' } }
+          it { is_expected.to be true }
+        end
+
+        context 'with a key and :for' do
+          subject(:done) { patcher.done?(key, for: for_key) }
+          let(:for_key) { double('key') }
+          before(:each) { patcher.do_once(key, for: for_key) { 'Perform patch' } }
+          it { is_expected.to be true }
+        end
       end
     end
   end