ddtrace 0.12.0.beta2 → 0.12.0.rc1

data/lib/ddtrace/contrib/active_record/utils.rb

@@ -0,0 +1,83 @@
+module Datadog
+  module Contrib
+    module ActiveRecord
+      # Common utilities for Rails
+      module Utils
+        # Return a canonical name for a type of database
+        def self.normalize_vendor(vendor)
+          case vendor
+          when nil
+            'defaultdb'
+          when 'postgresql'
+            'postgres'
+          when 'sqlite3'
+            'sqlite'
+          else
+            vendor
+          end
+        end
+
+        def self.adapter_name
+          connection_config[:adapter_name]
+        end
+
+        def self.database_name
+          connection_config[:database_name]
+        end
+
+        def self.adapter_host
+          connection_config[:adapter_host]
+        end
+
+        def self.adapter_port
+          connection_config[:adapter_port]
+        end
+
+        def self.connection_config(object_id = nil)
+          config = object_id.nil? ? default_connection_config : connection_config_by_id(object_id)
+          {
+            adapter_name: normalize_vendor(config[:adapter]),
+            adapter_host: config[:host],
+            adapter_port: config[:port],
+            database_name: config[:database]
+          }
+        end
+
+        # Attempt to retrieve the connection from an object ID.
+        def self.connection_by_id(object_id)
+          return nil if object_id.nil?
+          ObjectSpace._id2ref(object_id)
+        rescue StandardError
+          nil
+        end
+
+        # Attempt to retrieve the connection config from an object ID.
+        # Typical of ActiveSupport::Notifications `sql.active_record`
+        def self.connection_config_by_id(object_id)
+          connection = connection_by_id(object_id)
+          return {} if connection.nil?
+
+          if connection.instance_variable_defined?(:@config)
+            connection.instance_variable_get(:@config)
+          else
+            {}
+          end
+        end
+
+        def self.default_connection_config
+          return @default_connection_config if instance_variable_defined?(:@default_connection_config)
+          current_connection_name = if ::ActiveRecord::Base.respond_to?(:connection_specification_name)
+                                      ::ActiveRecord::Base.connection_specification_name
+                                    else
+                                      ::ActiveRecord::Base
+                                    end
+
+          connection_pool = ::ActiveRecord::Base.connection_handler.retrieve_connection_pool(current_connection_name)
+          connection_pool.nil? ? {} : (@default_connection_config = connection_pool.spec.config)
+        rescue StandardError
+          {}
+        end
+      end
+    end
+  end
+end

data/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

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

@@ -0,0 +1,155 @@
+module Datadog
+  module Contrib
+    module ActiveSupport
+      module Notifications
+        # An ActiveSupport::Notification subscription that wraps events with tracing.
+        class Subscription
+          attr_accessor \
+            :tracer,
+            :span_name,
+            :options
+
+          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
+            @handler = Handler.new(&block)
+            @callbacks = Callbacks.new
+          end
+
+          # ActiveSupport 3.x calls this
+          def call(name, start, finish, id, payload)
+            start_span(name, id, payload, start)
+            finish_span(name, id, payload, finish)
+          end
+
+          # ActiveSupport 4+ calls this on start
+          def start(name, id, payload)
+            start_span(name, id, payload)
+          end
+
+          # ActiveSupport 4+ calls this on finish
+          def finish(name, id, payload)
+            finish_span(name, id, payload)
+          end
+
+          def before_trace(&block)
+            callbacks.add(:before_trace, &block) if block_given?
+          end
+
+          def after_trace(&block)
+            callbacks.add(:after_trace, &block) if block_given?
+          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
+
+          attr_reader \
+            :handler,
+            :callbacks
+
+          def start_span(name, id, payload, start = nil)
+            # Run callbacks
+            callbacks.run(name, :before_trace, id, payload, start)
+
+            # Start a trace
+            tracer.trace(@span_name, @options).tap do |span|
+              # Assign start time if provided
+              span.start_time = start unless start.nil?
+            end
+          end
+
+          def finish_span(name, id, payload, finish = nil)
+            tracer.active_span.tap do |span|
+              # If no active span, return.
+              return nil if span.nil?
+
+              # Run handler for event
+              handler.run(span, name, id, payload)
+
+              # Finish the span
+              span.finish(finish)
+
+              # Run callbacks
+              callbacks.run(name, :after_trace, span, id, payload, finish)
+            end
+          end
+
+          # Pattern => ActiveSupport:Notifications::Subscribers
+          def subscribers
+            @subscribers ||= {}
+          end
+
+          # Wrapper for subscription handler
+          class Handler
+            attr_reader :block
+
+            def initialize(&block)
+              @block = block
+            end
+
+            def run(span, name, id, payload)
+              run!(span, name, id, payload)
+            rescue StandardError => e
+              Datadog::Tracer.log.debug("ActiveSupport::Notifications handler for '#{name}' failed: #{e.message}")
+            end
+
+            def run!(*args)
+              @block.call(*args)
+            end
+          end
+
+          # Wrapper for subscription callbacks
+          class Callbacks
+            attr_reader :blocks
+
+            def initialize
+              @blocks = {}
+            end
+
+            def add(key, &block)
+              blocks_for(key) << block if block_given?
+            end
+
+            def run(event, key, *args)
+              blocks_for(key).each do |callback|
+                begin
+                  callback.call(event, key, *args)
+                rescue StandardError => e
+                  Datadog::Tracer.log.debug(
+                    "ActiveSupport::Notifications '#{key}' callback for '#{event}' failed: #{e.message}"
+                  )
+                end
+              end
+            end
+
+            private
+
+            def blocks_for(key)
+              blocks[key] ||= []
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ddtrace/contrib/elasticsearch/patcher.rb

@@ -15,6 +15,7 @@ module Datadog
         include Base
         register_as :elasticsearch, auto_patch: true
         option :service_name, default: SERVICE
+        option :quantize, default: {}
 
         @patched = false
 
@@ -89,7 +90,11 @@ module Datadog
                   span.set_tag(METHOD, method)
                   span.set_tag(URL, url)
                   span.set_tag(PARAMS, params) if params
-                  span.set_tag(BODY, body) if body
+                  if body
+                    quantize_options = Datadog.configuration[:elasticsearch][:quantize]
+                    quantized_body = Datadog::Contrib::Elasticsearch::Quantize.format_body(body, quantize_options)
+                    span.set_tag(BODY, quantized_body)
+                  end
                   span.set_tag('out.host', host) if host
                   span.set_tag('out.port', port) if port
 

data/lib/ddtrace/contrib/elasticsearch/quantize.rb

@@ -3,6 +3,11 @@ module Datadog
     module Elasticsearch
       # Quantize contains ES-specific resource quantization tools.
       module Quantize
+        PLACEHOLDER = '?'.freeze
+        EXCLUDE_KEYS = [].freeze
+        SHOW_KEYS = [:_index, :_type, :_id].freeze
+        DEFAULT_OPTIONS = { exclude: EXCLUDE_KEYS, show: SHOW_KEYS }.freeze
+
         ID_REGEXP = %r{\/([0-9]+)([\/\?]|$)}
         ID_PLACEHOLDER = '/?\2'.freeze
 
@@ -16,6 +21,90 @@ module Datadog
           quantized_url = url.gsub(ID_REGEXP, ID_PLACEHOLDER)
           quantized_url.gsub(INDEX_REGEXP, INDEX_PLACEHOLDER)
         end
+
+        def format_body(body, options = {})
+          format_body!(body, options)
+        rescue StandardError
+          PLACEHOLDER
+        end
+
+        def format_body!(body, options = {})
+          options = merge_options(DEFAULT_OPTIONS, options)
+
+          # Determine if bulk query or not, based on content
+          statements = body.end_with?("\n") ? body.split("\n") : [body]
+
+          # Parse each statement and quantize them.
+          statements.collect do |string|
+            reserialize_json(string) do |obj|
+              format_statement(obj, options)
+            end
+          end.join("\n")
+        end
+
+        def format_statement(statement, options = {})
+          return statement if options[:show] == :all
+
+          case statement
+          when Hash
+            statement.each_with_object({}) do |(key, value), quantized|
+              if options[:show].include?(key.to_sym)
+                quantized[key] = value
+              elsif !options[:exclude].include?(key.to_sym)
+                quantized[key] = format_value(value, options)
+              end
+            end
+          else
+            format_value(statement, options)
+          end
+        end
+
+        def format_value(value, options = {})
+          return value if options[:show] == :all
+
+          case value
+          when Hash
+            format_statement(value, options)
+          when Array
+            # If any are objects, format them.
+            if value.any? { |v| v.class <= Hash || v.class <= Array }
+              value.collect { |i| format_value(i, options) }
+            # Otherwise short-circuit and return single placeholder
+            else
+              PLACEHOLDER
+            end
+          else
+            PLACEHOLDER
+          end
+        end
+
+        def merge_options(original, additional)
+          {}.tap do |options|
+            # Show
+            # If either is :all, value becomes :all
+            options[:show] = if original[:show] == :all || additional[:show] == :all
+                               :all
+                             else
+                               (original[:show] || []).dup.concat(additional[:show] || []).uniq
+                             end
+
+            # Exclude
+            options[:exclude] = (original[:exclude] || []).dup.concat(additional[:exclude] || []).uniq
+          end
+        end
+
+        # Parses a JSON object from a string, passes its value
+        # to the block provided, and dumps its result back to JSON.
+        # If JSON parsing fails, it prints fail_value.
+        def reserialize_json(string, fail_value = PLACEHOLDER)
+          return string unless block_given?
+          begin
+            JSON.dump(yield(JSON.parse(string)))
+          rescue JSON::ParserError
+            # If it can't parse/dump, don't raise an error.
+            fail_value
+          end
+        end
       end
     end
   end

data/lib/ddtrace/contrib/grape/endpoint.rb

@@ -71,7 +71,7 @@ module Datadog
             span.resource = resource
 
             # set the request span resource if it's a `rack.request` span
-            request_span = payload[:env][:datadog_rack_request_span]
+            request_span = payload[:env][Datadog::Contrib::Rack::TraceMiddleware::RACK_REQUEST_SPAN]
             if !request_span.nil? && request_span.name == 'rack.request'
               request_span.resource = resource
             end

data/lib/ddtrace/contrib/racecar/patcher.rb

@@ -1,4 +1,5 @@
 require 'ddtrace/ext/app_types'
+require 'ddtrace/contrib/active_support/notifications/subscriber'
 
 module Datadog
   module Contrib
@@ -6,19 +7,50 @@ module Datadog
       # Provides instrumentation for `racecar` through ActiveSupport instrumentation signals
       module Patcher
         include Base
+        include ActiveSupport::Notifications::Subscriber
+
         NAME_MESSAGE = 'racecar.message'.freeze
         NAME_BATCH = 'racecar.batch'.freeze
         register_as :racecar
-        option :tracer, default: Datadog.tracer
         option :service_name, default: 'racecar'
+        option :tracer, default: Datadog.tracer do |value|
+          (value || Datadog.tracer).tap do |v|
+            # Make sure to update tracers of all subscriptions
+            subscriptions.each do |subscription|
+              subscription.tracer = v
+            end
+          end
+        end
+
+        on_subscribe do
+          # Subscribe to single messages
+          subscription(
+            self::NAME_MESSAGE,
+            { service: configuration[:service_name] },
+            configuration[:tracer],
+            &method(:process)
+          ).tap do |subscription|
+            subscription.before_trace { ensure_clean_context! }
+            subscription.subscribe('process_message.racecar')
+          end
+
+          # Subscribe to batch messages
+          subscription(
+            self::NAME_BATCH,
+            { service: configuration[:service_name] },
+            configuration[:tracer],
+            &method(:process)
+          ).tap do |subscription|
+            subscription.before_trace { ensure_clean_context! }
+            subscription.subscribe('process_batch.racecar')
+          end
+        end
 
         class << self
           def patch
             return patched? if patched? || !compatible?
 
-            ::ActiveSupport::Notifications.subscribe('process_batch.racecar', self)
-            ::ActiveSupport::Notifications.subscribe('process_message.racecar', self)
-
+            subscribe!
             configuration[:tracer].set_service_info(
               configuration[:service_name],
               'racecar',
@@ -33,28 +65,17 @@ module Datadog
             @patched = false
           end
 
-          def start(event, _, payload)
-            ensure_clean_context!
-
-            name = event[/message/] ? NAME_MESSAGE : NAME_BATCH
-            span = configuration[:tracer].trace(name)
+          def process(span, event, _, payload)
             span.service = configuration[:service_name]
             span.resource = payload[:consumer_class]
+
             span.set_tag('kafka.topic', payload[:topic])
             span.set_tag('kafka.consumer', payload[:consumer_class])
             span.set_tag('kafka.partition', payload[:partition])
             span.set_tag('kafka.offset', payload[:offset]) if payload.key?(:offset)
             span.set_tag('kafka.first_offset', payload[:first_offset]) if payload.key?(:first_offset)
             span.set_tag('kafka.message_count', payload[:message_count]) if payload.key?(:message_count)
-          end
-
-          def finish(_, _, payload)
-            current_span = configuration[:tracer].call_context.current_span
-
-            return unless current_span
-
-            current_span.set_error(payload[:exception_object]) if payload[:exception_object]
-            current_span.finish
+            span.set_error(payload[:exception_object]) if payload[:exception_object]
           end
 
           private
@@ -67,9 +88,12 @@ module Datadog
             defined?(::Racecar) && defined?(::ActiveSupport::Notifications)
           end
 
+          # Context objects are thread-bound.
+          # If Racecar re-uses threads, context from a previous trace
+          # could leak into the new trace. This "cleans" current context,
+          # preventing such a leak.
           def ensure_clean_context!
             return unless configuration[:tracer].call_context.current_span
-
             configuration[:tracer].provider.context = Context.new
           end
         end