Changed: Inject :connection to event instead of using #_id2ref.

lib/ddtrace/contrib/active_record/events/sql.rb

@@ -24,7 +24,7 @@ def span_name
           end
 
           def process(span, event, _id, payload)
-            config = Utils.connection_config(payload[:connection_id])
+            config = Utils.connection_config(payload[:connection])
             settings = Datadog.configuration[:active_record, config]
             adapter_name = Datadog::Utils::Database.normalize_vendor(config[:adapter])
             service_name = !settings.nil? ? settings.service_name : configuration[:service_name]

lib/ddtrace/contrib/active_record/patcher.rb

@@ -1,4 +1,5 @@
 require 'ddtrace/contrib/patcher'
+require 'ddtrace/contrib/active_record/patches/abstract_adapter'
 require 'ddtrace/contrib/active_record/events'
 
 module Datadog
@@ -17,6 +18,7 @@ def patched?
         def patch
           do_once(:active_record) do
             begin
+              ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, Patches::AbstractAdapter)
               Events.subscribe!
             rescue StandardError => e
               Datadog::Tracer.log.error("Unable to apply Active Record integration: #{e}")

lib/ddtrace/contrib/active_record/patches/abstract_adapter.rb

@@ -0,0 +1,72 @@
+require 'set'
+require 'ddtrace/augmentation/shim'
+
+module Datadog
+  module Contrib
+    module ActiveRecord
+      # Defines basic behaviors for an ActiveRecord event.
+      module Patches
+        # Adds patch to AbstractAdapter to make it pass more information through
+        # ActiveSupport notifications, for better instrumentation.
+        module AbstractAdapter
+          def self.included(base)
+            if Gem::Version.new(RUBY_VERSION) < Gem::Version.new('2.0.0')
+              base.class_eval do
+                alias_method :log_without_datadog, :log
+                remove_method :log
+                include InstanceMethods
+              end
+            else
+              base.send(:prepend, InstanceMethods)
+            end
+          end
+
+          # Compatibility shim for Rubies not supporting `.prepend`
+          module InstanceMethodsCompatibility
+            def log(*args, &block)
+              log_without_datadog(*args, &block)
+            end
+          end
+
+          # InstanceMethods - implementing instrumentation
+          module InstanceMethods
+            include InstanceMethodsCompatibility unless Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('2.0.0')
+
+            EVENT_ACTIVERECORD_SQL = 'sql.active_record'.freeze
+
+            # Override #log since sometimes connections are initialized prior
+            # to when the patch is applied; this will allow existing connections
+            # to receive the Shim as well.
+            def log(*args, &block)
+              insert_shim! unless shim_inserted?
+              super
+            end
+
+            private
+
+            def shim_inserted?
+              instance_variable_defined?(:@instrumenter) \
+                && Datadog::Shim.shim?(@instrumenter)
+            end
+
+            def insert_shim!
+              @instrumenter = Datadog::Shim.new(@instrumenter) do |shim|
+                connection = self
+
+                shim.override_method!(:instrument) do |*args, &block|
+                  # Inject connection into arguments
+                  if args[0] == EVENT_ACTIVERECORD_SQL && args[1].is_a?(Hash)
+                    args[1][:connection] ||= connection
+                  end
+
+                  # Call original
+                  shim_target.instrument(*args, &block)
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/active_record/utils.rb

@@ -19,24 +19,13 @@ def self.adapter_port
           connection_config[:port]
         end
 
-        def self.connection_config(object_id = nil)
-          object_id.nil? ? default_connection_config : connection_config_by_id(object_id)
-        end
-
-        # Attempt to retrieve the connection from an object ID.
-        def self.connection_by_id(object_id)
-          return nil if object_id.nil?
-          ObjectSpace._id2ref(object_id)
-        rescue StandardError
-          nil
+        def self.connection_config(connection = nil)
+          connection.nil? ? default_connection_config : connection_config_from_connection(connection)
         end
 
         # Attempt to retrieve the connection config from an object ID.
         # Typical of ActiveSupport::Notifications `sql.active_record`
-        def self.connection_config_by_id(object_id)
-          connection = connection_by_id(object_id)
-          return {} if connection.nil?
-
+        def self.connection_config_from_connection(connection)
           if connection.instance_variable_defined?(:@config)
             connection.instance_variable_get(:@config)
           else

spec/ddtrace/contrib/active_record/performance_spec.rb

@@ -0,0 +1,55 @@
+require 'spec_helper'
+require 'ddtrace'
+
+require 'active_record'
+require 'sqlite3'
+
+RSpec.describe 'ActiveRecord tracing performance' do
+  let(:tracer) { ::Datadog::Tracer.new(writer: FauxWriter.new) }
+  let(:options) { { tracer: tracer } }
+  let(:spans) { tracer.writer.spans }
+
+  before(:each) do
+    skip('Performance test does not run in CI.')
+
+    # Configure the tracer
+    Datadog.configure do |c|
+      c.use :active_record, options
+    end
+  end
+
+  after(:each) { Datadog.registry[:active_record].reset_configuration! }
+
+  describe 'for an in-memory database' do
+    let!(:connection) do
+      ActiveRecord::Base.establish_connection(adapter: 'sqlite3', database: ':memory:')
+    end
+
+    describe 'when queried with a simple select' do
+      subject(:measurement) { measure(iterations) }
+      let(:iterations) { 100_000 }
+
+      def measure(iterations = 1)
+        Benchmark.measure do
+          iterations.times do
+            connection.connection.execute('SELECT 42')
+          end
+        end
+      end
+
+      before(:each) do
+        # Perform a measurement to warm up
+        measure(10)
+
+        # Discard warm-up spans
+        tracer.writer.spans
+      end
+
+      it 'produces a measurement' do
+        expect { measurement }.to_not raise_error
+        expect(spans).to have(iterations).items
+        puts "\nRun time for #{iterations} iterations: #{measurement.utime}\n"
+      end
+    end
+  end
+end