Merge pull request #409 from DataDog/feature/add_rake_instrumentation

Add Rake instrumentation

Appraisals

@@ -132,6 +132,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
@@ -154,6 +155,7 @@ else
     gem 'aws-sdk', '~> 2.0'
     gem 'sucker_punch'
     gem 'dalli'
+    gem 'rake', '< 12.3'
     gem 'resque', '< 2.0'
     gem 'mysql2', '0.3.21', platform: :ruby
     gem 'activerecord-mysql-adapter', platform: :ruby

Rakefile

@@ -55,6 +55,7 @@ namespace :spec do
     :mongodb,
     :racecar,
     :rack,
+    :rake,
     :redis,
     :resque,
     :sequel,
@@ -232,6 +233,7 @@ task :ci do
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:mongodb'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake spec:grpc'
     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'
@@ -242,6 +244,7 @@ task :ci do
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake spec:excon'
     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'

docs/GettingStarted.md

@@ -37,6 +37,7 @@ For descriptions of terminology used in APM, take a look at the [official docume
      - [Racecar](#racecar)
      - [Rack](#rack)
      - [Rails](#rails)
+     - [Rake](#rake)
      - [Redis](#redis)
      - [Resque](#resque)
      - [Sequel](#sequel)
@@ -257,6 +258,7 @@ For a list of available integrations, and their configuration options, please re
 | Racecar        | `racecar`       | `>= 0.3.5`             | *[Link](#racecar)*        | *[Link](https://github.com/zendesk/racecar)*                                   |
 | Rack           | `rack`          | `>= 1.4.7`             | *[Link](#rack)*           | *[Link](https://github.com/rack/rack)*                                         |
 | Rails          | `rails`         | `>= 3.2, < 5.2`        | *[Link](#rails)*          | *[Link](https://github.com/rails/rails)*                                       |
+| Rake           | `rake`          | `>= 12.0`              | *[Link](#rake)*           | *[Link](https://github.com/ruby/rake)*                                         |
 | Redis          | `redis`         | `>= 3.2, < 4.0`        | *[Link](#redis)*          | *[Link](https://github.com/redis/redis-rb)*                                    |
 | Resque         | `resque`        | `>= 1.0, < 2.0`        | *[Link](#resque)*         | *[Link](https://github.com/resque/resque)*                                     |
 | Sequel         | `sidekiq`       | `>= 3.41`              | *[Link](#sequel)*         | *[Link](https://github.com/jeremyevans/sequel)*                                |
@@ -754,6 +756,71 @@ Where `options` is an optional `Hash` that accepts the following parameters:
 | ``template_base_path`` | Used when the template name is parsed. If you don't store your templates in the ``views/`` folder, you may need to change this value | ``views/`` |
 | ``tracer`` | A ``Datadog::Tracer`` instance used to instrument the application. Usually you don't need to set that. | ``Datadog.tracer`` |
 
+### Rake
+
+You can add instrumentation around your Rake tasks by activating the `rake` integration. Each task and its subsequent subtasks will be traced.
+
+To activate Rake task tracing, add the following to your `Rakefile`:
+
+```ruby
+# At the top of your Rakefile:
+require 'rake'
+require 'ddtrace'
+
+Datadog.configure do |c|
+  c.use :rake, options
+end
+
+task :my_task do
+  # Do something task work here...
+end
+
+Rake::Task['my_task'].invoke
+```
+
+Where `options` is an optional `Hash` that accepts the following parameters:
+
+| Key | Description | Default |
+| --- | --- | --- |
+| ``enabled`` | Defines whether Rake tasks should be traced. Useful for temporarily disabling tracing. `true` or `false` | ``true`` |
+| ``quantize`` | Hash containing options for quantization of task arguments. See below for more details and examples. | ``{}`` |
+| ``service_name`` | Service name which the Rake task traces should be grouped under. | ``rake`` |
+| ``tracer`` | A ``Datadog::Tracer`` instance used to instrument the application. Usually you don't need to set that. | ``Datadog.tracer`` |
+
+**Configuring task quantization behavior**
+
+```ruby
+Datadog.configure do |c|
+  # Given a task that accepts :one, :two, :three...
+  # Invoked with 'foo', 'bar', 'baz'.
+
+  # Default behavior: all arguments are quantized.
+  # `rake.invoke.args` tag  --> ['?']
+  # `rake.execute.args` tag --> { one: '?', two: '?', three: '?' }
+  c.use :rake
+
+  # Show values for any argument matching :two exactly
+  # `rake.invoke.args` tag  --> ['?']
+  # `rake.execute.args` tag --> { one: '?', two: 'bar', three: '?' }
+  c.use :rake, quantize: { args: { show: [:two] } }
+
+  # Show all values for all arguments.
+  # `rake.invoke.args` tag  --> ['foo', 'bar', 'baz']
+  # `rake.execute.args` tag --> { one: 'foo', two: 'bar', three: 'baz' }
+  c.use :rake, quantize: { args: { show: :all } }
+
+  # Totally exclude any argument matching :three exactly
+  # `rake.invoke.args` tag  --> ['?']
+  # `rake.execute.args` tag --> { one: '?', two: '?' }
+  c.use :rake, quantize: { args: { exclude: [:three] } }
+
+  # Remove the arguments entirely
+  # `rake.invoke.args` tag  --> ['?']
+  # `rake.execute.args` tag --> {}
+  c.use :rake, quantize: { args: { exclude: :all } }
+end
+```
+
 ### Redis
 
 The Redis integration will trace simple calls as well as pipelines.

gemfiles/contrib.gemfile

@@ -21,6 +21,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

gemfiles/contrib_old.gemfile

@@ -19,6 +19,7 @@ gem "sidekiq", "4.0.0"
 gem "aws-sdk", "~> 2.0"
 gem "sucker_punch"
 gem "dalli"
+gem "rake", "< 12.3"
 gem "resque", "< 2.0"
 gem "mysql2", "0.3.21", platform: :ruby
 gem "activerecord-mysql-adapter", platform: :ruby

lib/ddtrace.rb

@@ -71,6 +71,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'

lib/ddtrace/contrib/rake/instrumentation.rb

@@ -0,0 +1,70 @@
+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)
+            return super unless enabled?
+
+            tracer.trace(SPAN_NAME_INVOKE) do |span|
+              super
+              annotate_invoke!(span, args)
+            end
+          end
+
+          def execute(args = nil)
+            return super unless enabled?
+
+            tracer.trace(SPAN_NAME_EXECUTE) do |span|
+              super
+              annotate_execute!(span, args)
+            end
+          end
+
+          private
+
+          def annotate_invoke!(span, args)
+            span.resource = name
+            span.set_tag('rake.task.arg_names', arg_names)
+            span.set_tag('rake.invoke.args', quantize_args(args)) unless args.nil?
+          rescue StandardError => e
+            Datadog::Tracer.log.debug("Error while tracing Rake invoke: #{e.message}")
+          end
+
+          def annotate_execute!(span, args)
+            span.resource = name
+            span.set_tag('rake.execute.args', quantize_args(args.to_hash)) unless args.nil?
+          rescue StandardError => e
+            Datadog::Tracer.log.debug("Error while tracing Rake execute: #{e.message}")
+          end
+
+          def quantize_args(args)
+            quantize_options = Datadog.configuration[:rake][:quantize][:args]
+            Datadog::Quantization::Hash.format(args, quantize_options)
+          end
+
+          def enabled?
+            configuration[:enabled] == true
+          end
+
+          def tracer
+            configuration[:tracer]
+          end
+
+          def configuration
+            Datadog.configuration[:rake]
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/rake/patcher.rb

@@ -0,0 +1,53 @@
+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
+        option :quantize, default: {}
+
+        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

lib/ddtrace/quantization/hash.rb

@@ -14,21 +14,24 @@ module Hash
       module_function
 
       def format(hash_obj, options = {})
+        options ||= {}
         format!(hash_obj, options)
       rescue StandardError
         options[:placeholder] || PLACEHOLDER
       end
 
       def format!(hash_obj, options = {})
+        options ||= {}
         options = merge_options(DEFAULT_OPTIONS, options)
         format_hash(hash_obj, options)
       end
 
       def format_hash(hash_obj, options = {})
-        return hash_obj if options[:show] == :all
-
         case hash_obj
         when ::Hash
+          return {} if options[:exclude] == :all
+          return hash_obj if options[:show] == :all
+
           hash_obj.each_with_object({}) do |(key, value), quantized|
             if options[:show].include?(key.to_sym)
               quantized[key] = value
@@ -71,7 +74,13 @@ def merge_options(original, additional)
                            end
 
           # Exclude
-          options[:exclude] = (original[:exclude] || []).dup.concat(additional[:exclude] || []).uniq
+          # If either is :all, value becomes :all
+          options[:exclude] = if original[:exclude] == :all || additional[:exclude] == :all
+                                :all
+                              else
+                                (original[:exclude] || []).dup.concat(additional[:exclude] || []).uniq
+                              end
+
           options[:placeholder] = additional[:placeholder] || original[:placeholder]
         end
       end