From ff035590a263c8e97257d88365649c2849440dc5 Mon Sep 17 00:00:00 2001 From: David Elner Date: Wed, 25 Apr 2018 17:00:24 -0400 Subject: [PATCH] Added: Rake integration. --- Appraisals | 2 + Rakefile | 3 + gemfiles/contrib.gemfile | 1 + lib/ddtrace.rb | 1 + lib/ddtrace/contrib/rake/instrumentation.rb | 63 ++++++ lib/ddtrace/contrib/rake/patcher.rb | 52 +++++ .../contrib/rake/instrumentation_spec.rb | 190 ++++++++++++++++++ 7 files changed, 312 insertions(+) create mode 100644 lib/ddtrace/contrib/rake/instrumentation.rb create mode 100644 lib/ddtrace/contrib/rake/patcher.rb create mode 100644 spec/ddtrace/contrib/rake/instrumentation_spec.rb diff --git a/Appraisals b/Appraisals index 1e135e82ed4..7d6f5cbe20f 100644 --- a/Appraisals +++ b/Appraisals @@ -129,6 +129,7 @@ if RUBY_VERSION >= '2.2.2' && RUBY_PLATFORM != 'java' gem 'aws-sdk' gem 'sucker_punch' gem 'dalli' + gem 'rake' gem 'resque', '< 2.0' gem 'racecar', '>= 0.3.5' gem 'mysql2', '< 0.5', platform: :ruby @@ -149,6 +150,7 @@ else gem 'aws-sdk', '~> 2.0' gem 'sucker_punch' gem 'dalli' + gem 'rake' gem 'resque', '< 2.0' gem 'mysql2', '0.3.21', platform: :ruby gem 'activerecord-mysql-adapter', platform: :ruby diff --git a/Rakefile b/Rakefile index 42e7f3b8ea1..8fba3d5158e 100644 --- a/Rakefile +++ b/Rakefile @@ -52,6 +52,7 @@ namespace :spec do :mongodb, :racecar, :rack, + :rake, :redis, :resque, :sequel, @@ -226,6 +227,7 @@ task :ci do sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:graphql' sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:mongodb' sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:racecar' + sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:rake' sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:redis' sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:resque' sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:sequel' @@ -234,6 +236,7 @@ task :ci do 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:mongodb' + sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:rake' sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:redis' sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:resque' sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:sequel' diff --git a/gemfiles/contrib.gemfile b/gemfiles/contrib.gemfile index 997ae5c580d..9ecc51ef9e6 100644 --- a/gemfiles/contrib.gemfile +++ b/gemfiles/contrib.gemfile @@ -19,6 +19,7 @@ gem "sidekiq" gem "aws-sdk" gem "sucker_punch" gem "dalli" +gem "rake" gem "resque", "< 2.0" gem "racecar", ">= 0.3.5" gem "mysql2", "< 0.5", platform: :ruby diff --git a/lib/ddtrace.rb b/lib/ddtrace.rb index f6d1df4dfe9..b8dba3246a1 100644 --- a/lib/ddtrace.rb +++ b/lib/ddtrace.rb @@ -68,6 +68,7 @@ def configure(target = configuration, opts = {}) require 'ddtrace/contrib/sucker_punch/patcher' require 'ddtrace/contrib/mongodb/patcher' require 'ddtrace/contrib/dalli/patcher' +require 'ddtrace/contrib/rake/patcher' require 'ddtrace/contrib/resque/patcher' require 'ddtrace/contrib/racecar/patcher' require 'ddtrace/contrib/sidekiq/patcher' diff --git a/lib/ddtrace/contrib/rake/instrumentation.rb b/lib/ddtrace/contrib/rake/instrumentation.rb new file mode 100644 index 00000000000..c4427f7f605 --- /dev/null +++ b/lib/ddtrace/contrib/rake/instrumentation.rb @@ -0,0 +1,63 @@ +module Datadog + module Contrib + module Rake + # Instrumentation for Rake tasks + module Instrumentation + SPAN_NAME_INVOKE = 'rake.invoke'.freeze + SPAN_NAME_EXECUTE = 'rake.execute'.freeze + + def self.included(base) + base.send(:prepend, InstanceMethods) + end + + # Instance methods for Rake instrumentation + module InstanceMethods + def invoke(*args) + if enabled? + tracer.trace(SPAN_NAME_INVOKE) do |span| + super + annotate!(span) + # TODO: Add quantization + span.set_tag('rake.args', args) + end + else + super + end + end + + def execute(args = nil) + if enabled? + tracer.trace(SPAN_NAME_EXECUTE) do |span| + super + annotate!(span) + # TODO: Add quantization + span.set_tag('rake.args', args.to_hash) unless args.nil? + end + else + super + end + end + + private + + def annotate!(span) + span.resource = name + span.set_tag('rake.arg_names', arg_names) + end + + def enabled? + configuration[:enabled] == true + end + + def tracer + configuration[:tracer] + end + + def configuration + Datadog.configuration[:rake] + end + end + end + end + end +end diff --git a/lib/ddtrace/contrib/rake/patcher.rb b/lib/ddtrace/contrib/rake/patcher.rb new file mode 100644 index 00000000000..80dac5a1cd7 --- /dev/null +++ b/lib/ddtrace/contrib/rake/patcher.rb @@ -0,0 +1,52 @@ +require 'ddtrace/ext/app_types' +require 'ddtrace/contrib/rake/instrumentation' + +module Datadog + module Contrib + module Rake + # Patcher for Rake instrumentation + module Patcher + include Base + + register_as :rake + option :service_name, default: 'rake' + option :tracer, default: Datadog.tracer + option :enabled, default: true + + module_function + + def patch + return patched? if patched? || !compatible? + + patch_rake + + # Set service info + configuration[:tracer].set_service_info( + configuration[:service_name], + 'rake', + Ext::AppTypes::WORKER + ) + + @patched = true + end + + def patched? + return @patched if defined?(@patched) + @patched = false + end + + def patch_rake + ::Rake::Task.send(:include, Instrumentation) + end + + def compatible? + RUBY_VERSION >= '2.0.0' && defined?(::Rake) + end + + def configuration + Datadog.configuration[:rake] + end + end + end + end +end diff --git a/spec/ddtrace/contrib/rake/instrumentation_spec.rb b/spec/ddtrace/contrib/rake/instrumentation_spec.rb new file mode 100644 index 00000000000..b7e77fc9067 --- /dev/null +++ b/spec/ddtrace/contrib/rake/instrumentation_spec.rb @@ -0,0 +1,190 @@ +require 'spec_helper' + +require 'securerandom' +require 'rake' +require 'rake/tasklib' +require 'ddtrace' +require 'ddtrace/contrib/rake/patcher' + +RSpec.describe Datadog::Contrib::Rake::Instrumentation do + let(:tracer) { Datadog::Tracer.new(writer: FauxWriter.new) } + let(:configuration_options) { { tracer: tracer, enabled: true } } + let(:spans) { tracer.writer.spans } + let(:span) { spans.first } + + before(:each) do + skip unless Datadog::Contrib::Rake::Patcher.compatible? + + # Reset options (that might linger from other tests) + Datadog.configuration[:rake].reset_options! + + # Patch Rake + Datadog.configure do |c| + c.use :rake, configuration_options + end + end + + after(:each) do + # We don't want instrumentation enabled during the rest of the test suite... + Datadog.configure do |c| + c.use :rake, enabled: false + end + end + + def reset_task!(task_name) + if Rake::Task.task_defined?(task_name) + Rake::Task[task_name].reenable + Rake::Task[task_name].clear + end + end + + let(:task_name) { :test_rake_instrumentation } + let(:task_body) { Proc.new { |task, args| spy.call(task, args) } } + let(:task_arg_names) { [] } + let(:task_class) do + stub_const('RakeInstrumentationTestTask', Class.new(Rake::TaskLib)).tap do |task_class| + tb = task_body + task_class.send(:define_method, :initialize) do |name = task_name, *args| + task(name, *args, &tb) + end + end + end + let(:task) { Rake::Task[task_name] } + let(:spy) { double('spy') } + + describe '#invoke' do + shared_examples_for 'a single task execution' do + before(:each) do + expect(spy).to receive(:call) do |invocation_task, invocation_args| + expect(invocation_task).to eq(task) + expect(invocation_args.to_hash).to eq(args_hash) + end + task.invoke(*args) + end + + let(:invoke_span) { spans.find { |s| s.name == described_class::SPAN_NAME_INVOKE } } + let(:execute_span) { spans.find { |s| s.name == described_class::SPAN_NAME_EXECUTE } } + + it do + expect(spans).to have(2).items + end + + describe '\'rake.invoke\' span' do + it do + expect(invoke_span.name).to eq(described_class::SPAN_NAME_INVOKE) + expect(invoke_span.resource).to eq(task_name.to_s) + expect(invoke_span.parent_id).to eq(0) + end + end + + describe '\'rake.execute\' span' do + it do + expect(execute_span.name).to eq(described_class::SPAN_NAME_EXECUTE) + expect(execute_span.resource).to eq(task_name.to_s) + expect(execute_span.parent_id).to eq(invoke_span.span_id) + end + end + end + + context 'for a task' do + let(:args_hash) { {} } + let(:task_arg_names) { args_hash.keys } + let(:args) { args_hash.values } + + def define_task! + reset_task!(task_name) + Rake::Task.define_task(task_name, *task_arg_names, &task_body) + end + + before(:each) { define_task! } + + context 'without args' do + it_behaves_like 'a single task execution' + end + + context 'with args' do + let(:args_hash) { { one: 1, two: 2, three: 3 } } + it_behaves_like 'a single task execution' + end + + context 'with a prerequisite task' do + let(:prerequisite_task_name) { :test_rake_instrumentation_prerequisite } + let(:prerequisite_task_body) { Proc.new { |task, args| prerequisite_spy.call(task, args) } } + let(:prerequisite_spy) { double('prerequisite spy') } + let(:prerequisite_task) { Rake::Task[prerequisite_task_name] } + + def define_task! + reset_task!(task_name) + reset_task!(prerequisite_task_name) + Rake::Task.define_task(prerequisite_task_name, &prerequisite_task_body) + Rake::Task.define_task(task_name => prerequisite_task_name, &task_body) + end + + before(:each) do + expect(prerequisite_spy).to receive(:call) do |invocation_task, invocation_args| + expect(invocation_task).to eq(prerequisite_task) + expect(invocation_args.to_hash).to eq({}) + end.ordered + + expect(spy).to receive(:call) do |invocation_task, invocation_args| + expect(invocation_task).to eq(task) + expect(invocation_args.to_hash).to eq(args_hash) + end.ordered + + task.invoke(*args) + end + + let(:invoke_span) { spans.find { |s| s.name == described_class::SPAN_NAME_INVOKE } } + let(:prerequisite_task_execute_span) do + spans.find do |s| + s.name == described_class::SPAN_NAME_EXECUTE \ + && s.resource == prerequisite_task_name.to_s + end + end + let(:task_execute_span) do + spans.find do |s| + s.name == described_class::SPAN_NAME_EXECUTE \ + && s.resource == task_name.to_s + end + end + + it do + expect(spans).to have(3).items + end + + describe '\'rake.invoke\' span' do + it do + expect(invoke_span.name).to eq(described_class::SPAN_NAME_INVOKE) + expect(invoke_span.resource).to eq(task_name.to_s) + expect(invoke_span.parent_id).to eq(0) + end + end + + describe 'prerequisite \'rake.execute\' span' do + it do + expect(prerequisite_task_execute_span.name).to eq(described_class::SPAN_NAME_EXECUTE) + expect(prerequisite_task_execute_span.resource).to eq(prerequisite_task_name.to_s) + expect(prerequisite_task_execute_span.parent_id).to eq(invoke_span.span_id) + end + end + + describe 'task \'rake.execute\' span' do + it do + expect(task_execute_span.name).to eq(described_class::SPAN_NAME_EXECUTE) + expect(task_execute_span.resource).to eq(task_name.to_s) + expect(task_execute_span.parent_id).to eq(invoke_span.span_id) + end + end + end + + context 'defined by a class' do + def define_task! + reset_task!(task_name) + task_class.new(task_name, *task_arg_names) + end + + it_behaves_like 'a single task execution' + end + end + end +end