Adds tracing support for ActiveRecord notifications

This change will publish distributed tracing spans related to
SQL ActiveRecord queries. These queries are unbound, so will still
contain parameter substitution strings in place of actual parameter
values.

This ensures that the SQL strings do not need to be scrubbed, since the
values are not included in the SQL.

More details of rails instrumentation can be found here:
https://guides.rubyonrails.org/active_support_instrumentation.html

5 files changed, 130 insertions, 1 deletions

diff --git a/changelogs/unreleased/an-opentracing-active-record-tracing.yml b/changelogs/unreleased/an-opentracing-active-record-tracing.yml
new file mode 100644
index 00000000000..59b480675ec
--- /dev/null
+++ b/changelogs/unreleased/an-opentracing-active-record-tracing.yml
@@ -0,0 +1,5 @@
+---
+title: Adds tracing support for ActiveRecord notifications
+merge_request: 24604
+author:
+type: other
diff --git a/config/initializers/tracing.rb b/config/initializers/tracing.rb
index 46e8125daf8..d5bef8edb43 100644
--- a/config/initializers/tracing.rb
+++ b/config/initializers/tracing.rb
@@ -23,6 +23,9 @@ if Gitlab::Tracing.enabled?
     end
   end
 
+  # Instrument Rails
+  Gitlab::Tracing::Rails::ActiveRecordSubscriber.instrument
+
   # In multi-processed clustered architectures (puma, unicorn) don't
   # start tracing until the worker processes are spawned. This works
   # around issues when the opentracing implementation spawns threads
diff --git a/lib/gitlab/tracing/common.rb b/lib/gitlab/tracing/common.rb
index 5e2b12e3f90..3a08ede8138 100644
--- a/lib/gitlab/tracing/common.rb
+++ b/lib/gitlab/tracing/common.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'opentracing'
+
 module Gitlab
   module Tracing
     module Common
@@ -32,6 +34,14 @@ module Gitlab
         end
       end
 
+      def postnotify_span(operation_name, start_time, end_time, tags: nil, child_of: nil, exception: nil)
+        span = OpenTracing.start_span(operation_name, start_time: start_time, tags: tags, child_of: child_of)
+
+        log_exception_on_span(span, exception) if exception
+
+        span.finish(end_time: end_time)
+      end
+
       def log_exception_on_span(span, exception)
         span.set_tag('error', true)
         span.log_kv(kv_tags_for_exception(exception))
@@ -44,7 +54,7 @@ module Gitlab
             'event':      'error',
             'error.kind': exception.class.to_s,
             'message':    Gitlab::UrlSanitizer.sanitize(exception.message),
-            'stack':      exception.backtrace.join("\n")
+            'stack':      exception.backtrace&.join("\n")
           }
         else
           {
diff --git a/lib/gitlab/tracing/rails/active_record_subscriber.rb b/lib/gitlab/tracing/rails/active_record_subscriber.rb
new file mode 100644
index 00000000000..214eac47e14
--- /dev/null
+++ b/lib/gitlab/tracing/rails/active_record_subscriber.rb
@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Tracing
+    module Rails
+      class ActiveRecordSubscriber
+        include Gitlab::Tracing::Common
+
+        ACTIVE_RECORD_NOTIFICATION_TOPIC = 'sql.active_record'
+        DEFAULT_OPERATION_NAME = "sqlquery"
+
+        def self.instrument
+          subscriber = new
+
+          ActiveSupport::Notifications.subscribe(ACTIVE_RECORD_NOTIFICATION_TOPIC) do |_, start, finish, _, payload|
+            subscriber.notify(start, finish, payload)
+          end
+        end
+
+        # For more information on the payloads: https://guides.rubyonrails.org/active_support_instrumentation.html
+        def notify(start, finish, payload)
+          operation_name = payload[:name].presence || DEFAULT_OPERATION_NAME
+          exception = payload[:exception]
+          tags = {
+            'component' =>        'ActiveRecord',
+            'span.kind' =>        'client',
+            'db.type' =>          'sql',
+            'db.connection_id' => payload[:connection_id],
+            'db.cached' =>        payload[:cached] || false,
+            'db.statement' =>     payload[:sql]
+          }
+
+          postnotify_span("active_record:#{operation_name}", start, finish, tags: tags, exception: exception)
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/tracing/rails/active_record_subscriber_spec.rb b/spec/lib/gitlab/tracing/rails/active_record_subscriber_spec.rb
new file mode 100644
index 00000000000..5eb5c044f84
--- /dev/null
+++ b/spec/lib/gitlab/tracing/rails/active_record_subscriber_spec.rb
@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+require 'rspec-parameterized'
+
+describe Gitlab::Tracing::Rails::ActiveRecordSubscriber do
+  using RSpec::Parameterized::TableSyntax
+
+  describe '.instrument' do
+    it 'is unsubscribable' do
+      subscription = described_class.instrument
+
+      expect(subscription).not_to be_nil
+      expect { ActiveSupport::Notifications.unsubscribe(subscription) }.not_to raise_error
+    end
+  end
+
+  describe '#notify' do
+    subject { described_class.new }
+    let(:start) { Time.now }
+    let(:finish) { Time.now }
+
+    where(:name, :operation_name, :exception, :connection_id, :cached, :cached_response, :sql) do
+      nil         | "active_record:sqlquery"  | nil               | nil | nil   | false | nil
+      ""          | "active_record:sqlquery"  | nil               | nil | nil   | false | nil
+      "User Load" | "active_record:User Load" | nil               | nil | nil   | false | nil
+      "Repo Load" | "active_record:Repo Load" | StandardError.new | nil | nil   | false | nil
+      nil         | "active_record:sqlquery"  | nil               | 123 | nil   | false | nil
+      nil         | "active_record:sqlquery"  | nil               | nil | false | false | nil
+      nil         | "active_record:sqlquery"  | nil               | nil | true  | true  | nil
+      nil         | "active_record:sqlquery"  | nil               | nil | true  | true  | "SELECT * FROM users"
+    end
+
+    with_them do
+      def payload
+        {
+          name: name,
+          exception: exception,
+          connection_id: connection_id,
+          cached: cached,
+          sql: sql
+        }
+      end
+
+      def expected_tags
+        {
+          "component" =>        "ActiveRecord",
+          "span.kind" =>        "client",
+          "db.type" =>          "sql",
+          "db.connection_id" => connection_id,
+          "db.cached" =>        cached_response,
+          "db.statement" =>     sql
+        }
+      end
+
+      it 'should notify the tracer when the hash contains null values' do
+        expect(subject).to receive(:postnotify_span).with(operation_name, start, finish, tags: expected_tags, exception: exception)
+
+        subject.notify(start, finish, payload)
+      end
+
+      it 'should notify the tracer when the payload is missing values' do
+        expect(subject).to receive(:postnotify_span).with(operation_name, start, finish, tags: expected_tags, exception: exception)
+
+        subject.notify(start, finish, payload.compact)
+      end
+
+      it 'should not throw exceptions when with the default tracer' do
+        expect { subject.notify(start, finish, payload) }.not_to raise_error
+      end
+    end
+  end
+end