Add support for ActiveRecord async queries

lib/datadog/tracing/contrib/active_support/notifications/subscription.rb

@@ -6,6 +6,7 @@ module Contrib
       module ActiveSupport
         module Notifications
           # An ActiveSupport::Notification subscription that wraps events with tracing.
+          # TODO: Inherit from {ActiveSupport::Subscriber}, to adhere to the public API.
           class Subscription
             attr_accessor \
               :span_name,
@@ -29,6 +30,11 @@ def initialize(span_name, span_options, on_start: nil, on_finish: nil, trace: ni
               @callbacks = Callbacks.new
             end
 
+            def publish(name, _time, _end, id, payload)
+              start(name, id, payload)
+              finish(name, id, payload)
+            end
+
             # Called by ActiveSupport on event start
             def start(name, id, payload)
               start_span(name, id, payload) if @trace&.call(name, payload)

spec/datadog/tracing/contrib/active_record/app.rb

@@ -9,6 +9,21 @@
 logger = Logger.new($stdout)
 logger.level = Logger::INFO
 
+# Enable the async query executor, so we can test Relation#load_async.
+# It does not affect non-async queries.
+if defined?(ActiveRecord) && ActiveRecord.respond_to?(:async_query_executor=)
+  ActiveRecord.async_query_executor = :global_thread_pool
+  #
+  # REMOVE ME if all tests pass
+  #
+  # if defined?(ActiveRecord::ConnectionAdapters::ConnectionPool) &&
+  #     ActiveRecord::ConnectionAdapters::ConnectionPool.respond_to?(:install_executor_hooks)
+  #   ActiveRecord::ConnectionAdapters::ConnectionPool.install_executor_hooks
+  # end
+  # # Force initialization of global thread pool
+  # ActiveRecord.global_thread_pool_async_query_executor if ActiveRecord.respond_to?(:global_thread_pool_async_query_executor)
+end
+
 # connecting to any kind of database is enough to test the integration
 root_pw = ENV.fetch('TEST_MYSQL_ROOT_PASSWORD', 'root')
 host = ENV.fetch('TEST_MYSQL_HOST', '127.0.0.1')

spec/datadog/tracing/contrib/active_record/tracer_spec.rb

@@ -33,8 +33,6 @@
   end
 
   context 'when query is made' do
-    before { Article.count }
-
     it_behaves_like 'analytics for integration' do
       let(:analytics_enabled_var) { Datadog::Tracing::Contrib::ActiveRecord::Ext::ENV_ANALYTICS_ENABLED }
       let(:analytics_sample_rate_var) { Datadog::Tracing::Contrib::ActiveRecord::Ext::ENV_ANALYTICS_SAMPLE_RATE }
@@ -160,5 +158,28 @@
         end
       end
     end
+
+    context 'with adapter supporting background execution' do
+      it 'parents the database span to the calling context' do
+        Datadog::Tracing.trace('root-span') do
+          relation = Article.limit(1).load_async
+
+          # Confirm async execution (there's no public API to confirm it).
+          expect(relation.instance_variable_get(:@future_result)).to_not be_nil
+
+          # Ensure we didn't break the query
+          expect(relation.to_a).to be_a(Array)
+        end
+
+        # Remove internal AR queries
+        spans.reject! { |s| s.resource.start_with?('SET ') }
+
+        expect(spans).to have(2).items
+
+        select = spans.find { |s| s.resource.include?('articles') }
+
+        expect(select).to_not be_root_span
+      end
+    end
   end
 end

spec/datadog/tracing/contrib/active_support/notifications/subscription_spec.rb

@@ -85,6 +85,52 @@
         end
       end
 
+      describe '#publish' do
+        subject(:result) { subscription.publish(name, time, end_time, id, payload) }
+
+        let(:name) { double('name') }
+        let(:time) { double('time') }
+        let(:end_time) { double('end_time') }
+        let(:id) { double('id') }
+        let(:span_op) { double('span_op') }
+
+        before do
+          allow(on_start_spy).to receive(:call)
+          allow(on_finish_spy).to receive(:call)
+        end
+
+        it 'calls both start and finish methods' do
+          expect(subscription).to receive(:start).with(name, id, payload).ordered
+          expect(subscription).to receive(:finish).with(name, id, payload).ordered
+          subject
+        end
+
+        context 'with a complete event lifecycle' do
+          let(:span_op) { instance_double(Datadog::Tracing::SpanOperation) }
+
+          it 'creates and finishes a span' do
+            expect(Datadog::Tracing).to receive(:trace).with(span_name, **options).and_return(span_op)
+            expect(on_start_spy).to receive(:call).with(span_op, name, id, payload)
+            expect(on_finish_spy).to receive(:call).with(span_op, name, id, payload)
+            expect(span_op).to receive(:finish).with(nil)
+
+            subject
+            expect(payload[:datadog_span]).to eq(span_op)
+          end
+        end
+
+        context 'with trace? returning false' do
+          let(:trace) { proc { |_name, _payload| false } }
+
+          it 'does not create a span but still calls finish' do
+            expect(Datadog::Tracing).not_to receive(:trace)
+            expect(on_start_spy).not_to receive(:call)
+            expect(on_finish_spy).not_to receive(:call)
+            subject
+          end
+        end
+      end
+
       describe '#before_trace' do
         context 'given a block' do
           let(:callback_block) { proc { callback_spy.call } }