Skip to content

Commit

Permalink
Merge pull request #380 from feature/active_support_notifications_sub…
Browse files Browse the repository at this point in the history
…scriber

ActiveSupport::Notifications subscriptions
  • Loading branch information
Emanuele Palazzetti authored Mar 26, 2018
2 parents 3a2ae39 + cce6a81 commit 9afa4a3
Show file tree
Hide file tree
Showing 5 changed files with 414 additions and 1 deletion.
5 changes: 4 additions & 1 deletion Rakefile
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ namespace :spec do

[
:active_record,
:active_support,
:aws,
:dalli,
:elasticsearch,
Expand All @@ -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
Expand Down Expand Up @@ -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'
Expand Down
66 changes: 66 additions & 0 deletions lib/ddtrace/contrib/active_support/notifications/subscriber.rb
Original file line number Diff line number Diff line change
@@ -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
70 changes: 70 additions & 0 deletions lib/ddtrace/contrib/active_support/notifications/subscription.rb
Original file line number Diff line number Diff line change
@@ -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
131 changes: 131 additions & 0 deletions spec/ddtrace/contrib/active_support/notifications/subscriber_spec.rb
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit 9afa4a3

Please sign in to comment.