Merge pull request #380 from feature/active_support_notifications_sub…

…scriber

ActiveSupport::Notifications subscriptions

Rakefile

@@ -41,6 +41,7 @@ namespace :spec do
 
   [
     :active_record,
+    :active_support,
     :aws,
     :dalli,
     :elasticsearch,
@@ -58,7 +59,7 @@ namespace :spec do
     :sucker_punch
   ].each do |contrib|
     RSpec::Core::RakeTask.new(contrib) do |t|
-      t.pattern = "spec/ddtrace/contrib/#{contrib}/*_spec.rb"
+      t.pattern = "spec/ddtrace/contrib/#{contrib}/**/*_spec.rb"
     end
   end
 end
@@ -225,12 +226,14 @@ task :ci do
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:resque'
     # RSpec
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:active_record'
+    sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:active_support'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:dalli'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:faraday'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:graphql'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:racecar'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:redis'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:active_record'
+    sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:active_support'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:dalli'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:faraday'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:redis'

lib/ddtrace/contrib/active_support/notifications/subscriber.rb

@@ -0,0 +1,66 @@
+require 'set'
+require 'ddtrace/contrib/active_support/notifications/subscription'
+
+module Datadog
+  module Contrib
+    module ActiveSupport
+      module Notifications
+        # For classes that listen to ActiveSupport::Notification events.
+        # Creates subscriptions that are wrapped with tracing.
+        module Subscriber
+          def self.included(base)
+            base.send(:extend, ClassMethods)
+          end
+
+          # Class methods that are implemented in the inheriting class.
+          module ClassMethods
+            # Returns a list of subscriptions created for this class.
+            def subscriptions
+              @subscriptions ||= Set.new
+            end
+
+            # Returns whether subscriptions have been activated, via #subscribe!
+            def subscribed?
+              subscribed == true
+            end
+
+            protected
+
+            # Defines a callback for when subscribe! is called.
+            # Should contain subscription setup, defined by the inheriting class.
+            def on_subscribe(&block)
+              @on_subscribe_block = block
+            end
+
+            # Runs the on_subscribe callback once, to activate subscriptions.
+            # Should be triggered by the inheriting class.
+            def subscribe!
+              return subscribed? if subscribed? || on_subscribe_block.nil?
+              on_subscribe_block.call
+              @subscribed = true
+            end
+
+            # Creates a subscription and immediately activates it.
+            def subscribe(pattern, span_name, options = {}, tracer = Datadog.tracer, &block)
+              subscription(span_name, options, tracer, &block).tap do |subscription|
+                subscription.subscribe(pattern)
+              end
+            end
+
+            # Creates a subscription without activating it.
+            # Subscription is added to the inheriting class' list of subscriptions.
+            def subscription(span_name, options = {}, tracer = Datadog.tracer, &block)
+              Subscription.new(tracer, span_name, options, &block).tap do |subscription|
+                subscriptions << subscription
+              end
+            end
+
+            private
+
+            attr_reader :subscribed, :on_subscribe_block
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/active_support/notifications/subscription.rb

@@ -0,0 +1,70 @@
+module Datadog
+  module Contrib
+    module ActiveSupport
+      module Notifications
+        # An ActiveSupport::Notification subscription that wraps events with tracing.
+        class Subscription
+          attr_reader \
+            :tracer,
+            :span_name,
+            :options,
+            :block
+
+          def initialize(tracer, span_name, options, &block)
+            raise ArgumentError, 'Must be given a block!' unless block_given?
+            @tracer = tracer
+            @span_name = span_name
+            @options = options
+            @block = block
+          end
+
+          def start(_name, _id, _payload)
+            ensure_clean_context!
+            tracer.trace(@span_name, @options)
+          end
+
+          def finish(name, id, payload)
+            tracer.active_span.tap do |span|
+              return nil if span.nil?
+              block.call(span, name, id, payload)
+              span.finish
+            end
+          end
+
+          def subscribe(pattern)
+            return false if subscribers.key?(pattern)
+            subscribers[pattern] = ::ActiveSupport::Notifications.subscribe(pattern, self)
+            true
+          end
+
+          def unsubscribe(pattern)
+            return false unless subscribers.key?(pattern)
+            ::ActiveSupport::Notifications.unsubscribe(subscribers[pattern])
+            subscribers.delete(pattern)
+            true
+          end
+
+          def unsubscribe_all
+            return false if subscribers.empty?
+            subscribers.keys.each { |pattern| unsubscribe(pattern) }
+            true
+          end
+
+          protected
+
+          # Pattern => ActiveSupport:Notifications::Subscribers
+          def subscribers
+            @subscribers ||= {}
+          end
+
+          private
+
+          def ensure_clean_context!
+            return unless tracer.call_context.current_span
+            tracer.provider.context = Context.new
+          end
+        end
+      end
+    end
+  end
+end

spec/ddtrace/contrib/active_support/notifications/subscriber_spec.rb

@@ -0,0 +1,131 @@
+require 'spec_helper'
+require 'ddtrace'
+
+require 'ddtrace/contrib/active_support/notifications/subscriber'
+
+RSpec.describe Datadog::Contrib::ActiveSupport::Notifications::Subscriber do
+  describe 'implemented' do
+    subject(:test_class) do
+      Class.new.tap do |klass|
+        klass.send(:include, described_class)
+      end
+    end
+
+    describe 'class' do
+      describe 'behavior' do
+        describe '#subscriptions' do
+          subject(:subscriptions) { test_class.subscriptions }
+
+          context 'when no subscriptions have been created' do
+            it { is_expected.to be_empty }
+          end
+
+          context 'when a subscription has been created' do
+            it do
+              subscription = test_class.send(
+                :subscription,
+                double('span name'),
+                double('options'),
+                double('tracer'),
+                &Proc.new { }
+              )
+
+              is_expected.to contain_exactly(subscription)
+            end
+          end
+        end
+
+        describe '#subscribed?' do
+          subject(:subscribed) { test_class.subscribed? }
+
+          context 'when #subscribe! hasn\'t been called' do
+            it { is_expected.to be false }
+          end
+
+          context 'after #subscribe! has been called' do
+            before(:each) do
+              test_class.send(:on_subscribe, &Proc.new { })
+              test_class.send(:subscribe!)
+            end
+
+            it { is_expected.to be true }
+          end
+        end
+
+        context 'that is protected' do
+          describe '#subscribe!' do
+            subject(:result) { test_class.send(:subscribe!) }
+
+            context 'when #on_subscribe' do
+              context 'is defined' do
+                let(:on_subscribe_block) { Proc.new { spy.call } }
+                let(:spy) { double(:spy) }
+
+                before(:each) { test_class.send(:on_subscribe, &on_subscribe_block) }
+
+                it do
+                  expect(spy).to receive(:call)
+                  is_expected.to be true
+                end
+
+                context 'but has already been called once' do
+                  before(:each) do
+                    allow(spy).to receive(:call)
+                    test_class.send(:subscribe!)
+                  end
+
+                  it do
+                    expect(spy).to_not receive(:call)
+                    is_expected.to be true
+                  end
+                end
+              end
+
+              context 'is not defined' do
+                it { is_expected.to be false }
+              end
+            end
+          end
+
+          describe '#subscribe' do
+            subject(:subscription) { test_class.send(:subscribe, pattern, span_name, options, tracer, &block) }
+            let(:pattern) { double('pattern') }
+            let(:span_name) { double('span name') }
+            let(:options) { double('options') }
+            let(:tracer) { double('tracer') }
+            let(:block) { Proc.new { } }
+
+            before(:each) do
+              expect(Datadog::Contrib::ActiveSupport::Notifications::Subscription).to receive(:new)
+                .with(tracer, span_name, options)
+                .and_call_original
+
+              expect_any_instance_of(Datadog::Contrib::ActiveSupport::Notifications::Subscription).to receive(:subscribe)
+                .with(pattern)
+            end
+
+            it { is_expected.to be_a_kind_of(Datadog::Contrib::ActiveSupport::Notifications::Subscription) }
+            it { expect(test_class.subscriptions).to contain_exactly(subscription) }
+          end
+
+          describe '#subscription' do
+            subject(:subscription) { test_class.send(:subscription, span_name, options, tracer, &block) }
+            let(:span_name) { double('span name') }
+            let(:options) { double('options') }
+            let(:tracer) { double('tracer') }
+            let(:block) { Proc.new { } }
+
+            before(:each) do
+              expect(Datadog::Contrib::ActiveSupport::Notifications::Subscription).to receive(:new)
+                .with(tracer, span_name, options)
+                .and_call_original
+            end
+
+            it { is_expected.to be_a_kind_of(Datadog::Contrib::ActiveSupport::Notifications::Subscription) }
+            it { expect(test_class.subscriptions).to contain_exactly(subscription) }
+          end
+        end
+      end
+    end
+  end
+end