ddtrace 0.17.3 → 0.18.0

data/lib/ddtrace.rb

@@ -9,9 +9,12 @@ require 'ddtrace/quantization/http'
 require 'ddtrace/pipeline'
 require 'ddtrace/configuration'
 require 'ddtrace/patcher'
+require 'ddtrace/augmentation'
 
 # \Datadog global namespace that includes all tracing functionality for Tracer and Span classes.
 module Datadog
+  extend Augmentation
+
   @tracer = Tracer.new
   @registry = Registry.new
   @configuration = Configuration.new(registry: @registry)
@@ -68,6 +71,7 @@ require 'ddtrace/contrib/redis/integration'
 require 'ddtrace/contrib/resque/integration'
 require 'ddtrace/contrib/rest_client/integration'
 require 'ddtrace/contrib/sequel/integration'
+require 'ddtrace/contrib/shoryuken/integration'
 require 'ddtrace/contrib/sidekiq/integration'
 require 'ddtrace/contrib/sinatra/integration'
 require 'ddtrace/contrib/sucker_punch/integration'

data/lib/ddtrace/augmentation.rb

@@ -0,0 +1,13 @@
+require 'ddtrace/augmentation/method_wrapper'
+require 'ddtrace/augmentation/method_wrapping'
+require 'ddtrace/augmentation/shim'
+
+module Datadog
+  # Namespace for components that help modify
+  # existing code for instrumentation purposes.
+  module Augmentation
+    def shim(object, &block)
+      Shim.new(object, &block)
+    end
+  end
+end

data/lib/ddtrace/augmentation/method_wrapper.rb

@@ -0,0 +1,20 @@
+module Datadog
+  # Represents an wrapped method, with a reference to the original block
+  # and the block that wraps around it.
+  class MethodWrapper
+    attr_reader \
+      :original,
+      :wrapper
+
+    DEFAULT_WRAPPER = proc { |original, *args, &block| original.call(*args, &block) }
+
+    def initialize(original, &block)
+      @original = original
+      @wrapper = block_given? ? block : DEFAULT_WRAPPER
+    end
+
+    def call(*args, &block)
+      wrapper.call(original, *args, &block)
+    end
+  end
+end

data/lib/ddtrace/augmentation/method_wrapping.rb

@@ -0,0 +1,38 @@
+require 'set'
+require 'ddtrace/patcher'
+
+module Datadog
+  # Shorthands for wrapping methods
+  module MethodWrapping
+    include Datadog::Patcher
+
+    def wrapped_methods
+      @wrapped_methods ||= Set.new
+    end
+
+    # Adds method block directly to the object.
+    # Block is evaluated in the context of the object.
+    # Faster than #wrap_method!
+    def override_method!(method_name, &block)
+      return unless block_given?
+
+      without_warnings do
+        define_singleton_method(method_name, &block).tap do
+          wrapped_methods.add(method_name)
+        end
+      end
+    end
+
+    # Adds method wrapper to the object.
+    # Block is evaluated in the original context of the block.
+    # Slower than #override_method!
+    def wrap_method!(original_method, &block)
+      return unless block_given?
+      original_method = original_method.is_a?(Symbol) ? method(original_method) : original_method
+
+      override_method!(original_method.name) do |*original_args, &original_block|
+        block.call(original_method, *original_args, &original_block)
+      end
+    end
+  end
+end

data/lib/ddtrace/augmentation/shim.rb

@@ -0,0 +1,102 @@
+require 'set'
+require 'ddtrace/patcher'
+require 'ddtrace/augmentation/method_wrapping'
+
+module Datadog
+  # A "stand-in" that intercepts calls to another object. i.e. man-in-the-middle.
+  # This shim forwards all methods to object, except those overriden.
+  # Useful if you want to intercept inbound behavior to an object without modifying
+  # the object in question, especially useful if the overridding behavior shouldn't be global.
+  class Shim
+    extend Forwardable
+    include Datadog::Patcher
+    include Datadog::MethodWrapping
+
+    METHODS = Set[
+      :override_method!,
+      :shim,
+      :shim?,
+      :shim_target,
+      :wrap_method!,
+      :wrapped_methods
+    ].freeze
+
+    EXCLUDED_METHODS = Set[
+      # For all objects
+      :__binding__,
+      :__id__,
+      :__send__,
+      :extend,
+      :itself,
+      :object_id,
+      :respond_to?,
+      :tap
+    ].freeze
+
+    attr_reader :shim_target, :shim
+
+    def self.shim?(object)
+      # Check whether it responds to #shim? because otherwise the
+      # Shim forwards all method calls, including type checks to
+      # the wrapped object, to mimimize its intrusion.
+      object.respond_to?(:shim?)
+    end
+
+    # Pass this a block to override methods
+    def initialize(shim_target)
+      @shim = self
+      @shim_target = shim_target
+
+      # Save a reference to the original :define_singleton_method
+      # so methods can be defined on the shim after forwarding is applied.
+      @definition_method = method(:define_singleton_method)
+
+      # Wrap any methods
+      yield(self) if block_given?
+
+      # Forward methods
+      forwarded_methods = (
+        shim_target.public_methods.to_set \
+        - METHODS \
+        - EXCLUDED_METHODS \
+        - wrapped_methods
+      )
+      forward_methods!(*forwarded_methods)
+    end
+
+    def override_method!(method_name, &block)
+      return unless block_given?
+
+      without_warnings do
+        @definition_method.call(method_name, &block).tap do
+          wrapped_methods.add(method_name)
+        end
+      end
+    end
+
+    def wrap_method!(method_name, &block)
+      super(shim_target.method(method_name), &block)
+    end
+
+    def shim?
+      true
+    end
+
+    def respond_to?(method_name)
+      return true if METHODS.include?(method_name)
+      shim_target.respond_to?(method_name)
+    end
+
+    private
+
+    def forward_methods!(*forwarded_methods)
+      return if forwarded_methods.empty?
+
+      singleton_class.send(
+        :def_delegators,
+        :@shim_target,
+        *forwarded_methods
+      )
+    end
+  end
+end

data/lib/ddtrace/contrib/active_record/events/sql.rb

@@ -24,10 +24,14 @@ module Datadog
           end
 
           def process(span, event, _id, payload)
-            config = Utils.connection_config(payload[:connection_id])
+            config = Utils.connection_config(payload[:connection])
             settings = Datadog.configuration[:active_record, config]
             adapter_name = Datadog::Utils::Database.normalize_vendor(config[:adapter])
-            service_name = !settings.nil? ? settings.service_name : configuration[:service_name]
+            service_name = if settings.service_name != Datadog::Utils::Database::VENDOR_DEFAULT
+                             settings.service_name
+                           else
+                             adapter_name
+                           end
 
             span.name = "#{adapter_name}.query"
             span.service = service_name

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

@@ -1,4 +1,5 @@
 require 'ddtrace/contrib/patcher'
+require 'ddtrace/contrib/active_record/patches/abstract_adapter'
 require 'ddtrace/contrib/active_record/events'
 
 module Datadog
@@ -17,6 +18,7 @@ module Datadog
         def patch
           do_once(:active_record) do
             begin
+              ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, Patches::AbstractAdapter)
               Events.subscribe!
             rescue StandardError => e
               Datadog::Tracer.log.error("Unable to apply Active Record integration: #{e}")

data/lib/ddtrace/contrib/active_record/patches/abstract_adapter.rb

@@ -0,0 +1,72 @@
+require 'set'
+require 'ddtrace/augmentation/shim'
+
+module Datadog
+  module Contrib
+    module ActiveRecord
+      # Defines basic behaviors for an ActiveRecord event.
+      module Patches
+        # Adds patch to AbstractAdapter to make it pass more information through
+        # ActiveSupport notifications, for better instrumentation.
+        module AbstractAdapter
+          def self.included(base)
+            if Gem::Version.new(RUBY_VERSION) < Gem::Version.new('2.0.0')
+              base.class_eval do
+                alias_method :log_without_datadog, :log
+                remove_method :log
+                include InstanceMethods
+              end
+            else
+              base.send(:prepend, InstanceMethods)
+            end
+          end
+
+          # Compatibility shim for Rubies not supporting `.prepend`
+          module InstanceMethodsCompatibility
+            def log(*args, &block)
+              log_without_datadog(*args, &block)
+            end
+          end
+
+          # InstanceMethods - implementing instrumentation
+          module InstanceMethods
+            include InstanceMethodsCompatibility unless Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('2.0.0')
+
+            EVENT_ACTIVERECORD_SQL = 'sql.active_record'.freeze
+
+            # Override #log since sometimes connections are initialized prior
+            # to when the patch is applied; this will allow existing connections
+            # to receive the Shim as well.
+            def log(*args, &block)
+              insert_shim! unless shim_inserted?
+              super
+            end
+
+            private
+
+            def shim_inserted?
+              instance_variable_defined?(:@instrumenter) \
+                && Datadog::Shim.shim?(@instrumenter)
+            end
+
+            def insert_shim!
+              @instrumenter = Datadog::Shim.new(@instrumenter) do |shim|
+                connection = self
+
+                shim.override_method!(:instrument) do |*args, &block|
+                  # Inject connection into arguments
+                  if args[0] == EVENT_ACTIVERECORD_SQL && args[1].is_a?(Hash)
+                    args[1][:connection] ||= connection
+                  end
+
+                  # Call original
+                  shim_target.instrument(*args, &block)
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -3,6 +3,8 @@ module Datadog
     module ActiveRecord
       # Common utilities for Rails
       module Utils
+        EMPTY_CONFIG = {}.freeze
+
         def self.adapter_name
           Datadog::Utils::Database.normalize_vendor(connection_config[:adapter])
         end
@@ -19,28 +21,16 @@ module Datadog
           connection_config[:port]
         end
 
-        def self.connection_config(object_id = nil)
-          object_id.nil? ? default_connection_config : connection_config_by_id(object_id)
+        def self.connection_config(connection = nil)
+          connection.nil? ? default_connection_config : connection_config_from_connection(connection)
         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?
-
+        def self.connection_config_from_connection(connection)
           if connection.instance_variable_defined?(:@config)
             connection.instance_variable_get(:@config)
           else
-            {}
+            EMPTY_CONFIG
           end
         end
 
@@ -53,9 +43,9 @@ module Datadog
                                     end
 
           connection_pool = ::ActiveRecord::Base.connection_handler.retrieve_connection_pool(current_connection_name)
-          connection_pool.nil? ? {} : (@default_connection_config = connection_pool.spec.config)
+          connection_pool.nil? ? EMPTY_CONFIG : (@default_connection_config = connection_pool.spec.config)
         rescue StandardError
-          {}
+          EMPTY_CONFIG
         end
       end
     end

data/lib/ddtrace/contrib/aws/instrumentation.rb

@@ -13,22 +13,18 @@ module Datadog
       # Generates Spans for all interactions with AWS
       class Handler < Seahorse::Client::Handler
         def call(context)
-          pin = Datadog::Pin.get_from(::Aws)
-
-          return @handler.call(context) unless pin && pin.enabled?
-
-          pin.tracer.trace(Ext::SPAN_COMMAND) do |span|
-            result = @handler.call(context)
-            annotate!(span, pin, ParsedContext.new(context))
-            result
+          tracer.trace(Ext::SPAN_COMMAND) do |span|
+            @handler.call(context).tap do
+              annotate!(span, ParsedContext.new(context))
+            end
           end
         end
 
         private
 
-        def annotate!(span, pin, context)
-          span.service = pin.service
-          span.span_type = pin.app_type
+        def annotate!(span, context)
+          span.service = configuration[:service_name]
+          span.span_type = Datadog::Ext::AppTypes::WEB
           span.name = Ext::SPAN_COMMAND
           span.resource = context.safely(:resource)
           span.set_tag(Ext::TAG_AGENT, Ext::TAG_DEFAULT_AGENT)
@@ -39,6 +35,14 @@ module Datadog
           span.set_tag(Datadog::Ext::HTTP::METHOD, context.safely(:http_method))
           span.set_tag(Datadog::Ext::HTTP::STATUS_CODE, context.safely(:status_code))
         end
+
+        def tracer
+          configuration[:tracer]
+        end
+
+        def configuration
+          Datadog.configuration[:aws]
+        end
       end
     end
   end

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

@@ -22,7 +22,6 @@ module Datadog
               require 'ddtrace/contrib/aws/instrumentation'
               require 'ddtrace/contrib/aws/services'
 
-              add_pin
               add_plugin(Seahorse::Client::Base, *loaded_constants)
             rescue StandardError => e
               Datadog::Tracer.log.error("Unable to apply AWS integration: #{e}")
@@ -30,16 +29,6 @@ module Datadog
           end
         end
 
-        def add_pin
-          Pin
-            .new(
-              get_option(:service_name),
-              app: Ext::APP,
-              app_type: Datadog::Ext::AppTypes::WEB,
-              tracer: get_option(:tracer)
-            ).onto(::Aws)
-        end
-
         def add_plugin(*targets)
           targets.each { |klass| klass.add_plugin(Instrumentation) }
         end

data/lib/ddtrace/contrib/configurable.rb

@@ -20,11 +20,8 @@ module Datadog
           @resolver = nil
         end
 
-        def configuration(name = :default)
-          name = :default if name.nil?
-          name = resolver.resolve(name)
-          return nil unless configurations.key?(name)
-          configurations[name]
+        def configuration(key = :default)
+          configurations[resolve_configuration_key(key)]
         end
 
         def configurations
@@ -33,12 +30,12 @@ module Datadog
           end
         end
 
-        def configure(name, options = {}, &block)
-          name = resolver.resolve(name || :default)
+        def configure(key, options = {}, &block)
+          key = resolver.resolve(key || :default)
 
-          configurations[name].tap do |settings|
+          configurations[key].tap do |settings|
             settings.configure(options, &block)
-            configurations[name] = settings
+            configurations[key] = settings
           end
         end
 
@@ -49,6 +46,13 @@ module Datadog
         def resolver
           @resolver ||= Configuration::Resolver.new
         end
+
+        def resolve_configuration_key(key = :default)
+          key = :default if key.nil?
+          key = resolver.resolve(key)
+          key = :default unless configurations.key?(key)
+          key
+        end
       end
     end
   end