ddtrace 0.42.0 → 0.54.2

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

@@ -0,0 +1,33 @@
+# typed: false
+require 'ddtrace/contrib/patcher'
+require 'ddtrace/contrib/active_job/ext'
+require 'ddtrace/contrib/active_job/events'
+require 'ddtrace/contrib/active_job/log_injection'
+
+module Datadog
+  module Contrib
+    module ActiveJob
+      # Patcher enables patching of 'active_job' module.
+      module Patcher
+        include Contrib::Patcher
+
+        module_function
+
+        def target_version
+          Integration.version
+        end
+
+        def patch
+          Events.subscribe!
+          inject_log_correlation if Datadog.configuration[:active_job][:log_injection]
+        end
+
+        def inject_log_correlation
+          ::ActiveSupport.on_load(:active_job) do
+            include LogInjection
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ddtrace/contrib/active_model_serializers/configuration/settings.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/configuration/settings'
 require 'ddtrace/contrib/active_model_serializers/ext'
 

data/lib/ddtrace/contrib/active_model_serializers/event.rb

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/ext/http'
 require 'ddtrace/contrib/analytics'
 require 'ddtrace/contrib/active_support/notifications/event'
@@ -9,8 +10,8 @@ module Datadog
       # Defines basic behaviors for an ActiveModelSerializers event.
       module Event
         def self.included(base)
-          base.send(:include, ActiveSupport::Notifications::Event)
-          base.send(:extend, ClassMethods)
+          base.include(ActiveSupport::Notifications::Event)
+          base.extend(ClassMethods)
         end
 
         # Class methods for ActiveModelSerializers events.

data/lib/ddtrace/contrib/active_model_serializers/events/render.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/active_model_serializers/ext'
 require 'ddtrace/contrib/active_model_serializers/event'
 

data/lib/ddtrace/contrib/active_model_serializers/events/serialize.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/active_model_serializers/ext'
 require 'ddtrace/contrib/active_model_serializers/event'
 

data/lib/ddtrace/contrib/active_model_serializers/events.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/active_model_serializers/events/render'
 require 'ddtrace/contrib/active_model_serializers/events/serialize'
 

data/lib/ddtrace/contrib/active_model_serializers/ext.rb

@@ -1,3 +1,4 @@
+# typed: true
 module Datadog
   module Contrib
     module ActiveModelSerializers

data/lib/ddtrace/contrib/active_model_serializers/integration.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/integration'
 require 'ddtrace/contrib/active_model_serializers/configuration/settings'
 require 'ddtrace/contrib/active_model_serializers/patcher'

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

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/contrib/patcher'
 require 'ddtrace/ext/app_types'
 require 'ddtrace/contrib/active_model_serializers/ext'

data/lib/ddtrace/contrib/active_record/configuration/makara_resolver.rb

@@ -0,0 +1,31 @@
+# typed: true
+module Datadog
+  module Contrib
+    module ActiveRecord
+      module Configuration
+        # The `makara` gem has the concept of **role**, which can be
+        # inferred from the configuration `name`, in the form of:
+        # `master/0`, `replica/0`, `replica/1`, etc.
+        # The first part of this string is the database role.
+        #
+        # This allows the matching of a connection based on its role,
+        # instead of connection-specific information.
+        module MakaraResolver
+          def normalize_for_config(active_record_config)
+            hash = super
+            hash[:makara_role] = active_record_config[:makara_role]
+            hash
+          end
+
+          def normalize_for_resolve(active_record_config)
+            hash = super
+
+            hash[:makara_role] = active_record_config[:name].split('/')[0].to_s if active_record_config[:name].is_a?(String)
+
+            hash
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ddtrace/contrib/active_record/configuration/resolver.rb

@@ -1,43 +1,133 @@
+# typed: false
 require 'ddtrace/contrib/configuration/resolver'
 require 'ddtrace/vendor/active_record/connection_specification'
+require_relative 'makara_resolver'
 
 module Datadog
   module Contrib
     module ActiveRecord
       module Configuration
         # Converts Symbols, Strings, and Hashes to a normalized connection settings Hash.
+        #
+        # When matching using a Hash, these are the valid fields:
+        # ```
+        # {
+        #   adapter: ...,
+        #   host: ...,
+        #   port: ...,
+        #   database: ...,
+        #   username: ...,
+        #   role: ...,
+        # }
+        # ```
+        #
+        # Partial matching is supported: not including certain fields or setting them to `nil`
+        # will cause them to matching all values for that field. For example: `database: nil`
+        # will match any database, given the remaining fields match.
+        #
+        # Any fields not listed above are discarded.
+        #
+        # When more than one configuration could be matched, the last one to match is selected,
+        # based on addition order (`#add`).
         class Resolver < Contrib::Configuration::Resolver
-          def initialize(configurations = nil)
-            @configurations = configurations
+          prepend MakaraResolver
+
+          def initialize(active_record_configuration = nil)
+            super()
+
+            @active_record_configuration = active_record_configuration
           end
 
-          def resolve(key)
-            normalize(connection_resolver.resolve(key).symbolize_keys)
+          def active_record_configuration
+            @active_record_configuration || ::ActiveRecord::Base.configurations
           end
 
-          def configurations
-            @configurations || ::ActiveRecord::Base.configurations
+          def add(matcher, value)
+            parsed = parse_matcher(matcher)
+
+            # In case of error parsing, don't store `nil` key
+            # as it wouldn't be useful for matching configuration
+            # hashes in `#resolve`.
+            super(parsed, value) if parsed
           end
 
-          def connection_resolver
-            @resolver ||= begin
-              if defined?(::ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver)
-                ::ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(configurations)
-              else
-                ::Datadog::Vendor::ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(configurations)
+          def resolve(db_config)
+            active_record_config = resolve_connection_key(db_config).symbolize_keys
+
+            hash = normalize_for_resolve(active_record_config)
+
+            # Hashes in Ruby maintain insertion order
+            _, config = @configurations.reverse_each.find do |matcher, _|
+              matcher.none? do |key, value|
+                value != hash[key]
               end
             end
+
+            config
+          rescue => e
+            Datadog.logger.error(
+              "Failed to resolve ActiveRecord configuration key #{db_config.inspect}. " \
+              "Cause: #{e.message} Source: #{Array(e.backtrace).first}"
+            )
+
+            nil
           end
 
-          def normalize(hash)
+          protected
+
+          def parse_matcher(matcher)
+            resolved_pattern = resolve_connection_key(matcher).symbolize_keys
+            normalized = normalize_for_config(resolved_pattern)
+
+            # Remove empty fields to allow for partial matching
+            normalized.reject! { |_, v| v.nil? }
+
+            normalized
+          rescue => e
+            Datadog.logger.error(
+              "Failed to resolve ActiveRecord configuration key #{matcher.inspect}. " \
+              "Cause: #{e.message} Source: #{Array(e.backtrace).first}"
+            )
+          end
+
+          def connection_resolver
+            @resolver ||= if defined?(::ActiveRecord::Base.configurations.resolve)
+                            ::ActiveRecord::DatabaseConfigurations.new(active_record_configuration)
+                          elsif defined?(::ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver)
+                            ::ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(
+                              active_record_configuration
+                            )
+                          else
+                            ::Datadog::Vendor::ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(
+                              active_record_configuration
+                            )
+                          end
+          end
+
+          def resolve_connection_key(key)
+            result = connection_resolver.resolve(key)
+
+            if result.respond_to?(:configuration_hash) # Rails >= 6.1
+              result.configuration_hash
+            else # Rails < 6.1
+              result
+            end
+          end
+
+          # Extract only fields we'd like to match
+          # from the ActiveRecord configuration.
+          def normalize_for_config(active_record_config)
             {
-              adapter:  hash[:adapter],
-              host:     hash[:host],
-              port:     hash[:port],
-              database: hash[:database],
-              username: hash[:username]
+              adapter: active_record_config[:adapter],
+              host: active_record_config[:host],
+              port: active_record_config[:port],
+              database: active_record_config[:database],
+              username: active_record_config[:username]
             }
           end
+
+          # Both resolvers perform the same operations for this implementation, but can be specialized
+          alias_method :normalize_for_resolve, :normalize_for_config
         end
       end
     end

data/lib/ddtrace/contrib/active_record/configuration/settings.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/configuration/settings'
 require 'ddtrace/contrib/active_record/ext'
 require 'ddtrace/contrib/active_record/utils'

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

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/contrib/active_support/notifications/event'
 
 module Datadog
@@ -6,8 +7,8 @@ module Datadog
       # Defines basic behaviors for an ActiveRecord event.
       module Event
         def self.included(base)
-          base.send(:include, ActiveSupport::Notifications::Event)
-          base.send(:extend, ClassMethods)
+          base.include(ActiveSupport::Notifications::Event)
+          base.extend(ClassMethods)
         end
 
         # Class methods for ActiveRecord events.

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

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/analytics'
 require 'ddtrace/contrib/active_record/ext'
 require 'ddtrace/contrib/active_record/event'

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

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/ext/integration'
 require 'ddtrace/ext/net'
 require 'ddtrace/contrib/analytics'

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

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/active_record/events/instantiation'
 require 'ddtrace/contrib/active_record/events/sql'
 

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

@@ -1,3 +1,4 @@
+# typed: true
 module Datadog
   module Contrib
     module ActiveRecord

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

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/ext/sql'
 require 'ddtrace/ext/app_types'
 require 'ddtrace/contrib/integration'
@@ -5,6 +6,7 @@ require 'ddtrace/contrib/active_record/events'
 require 'ddtrace/contrib/active_record/configuration/resolver'
 require 'ddtrace/contrib/active_record/configuration/settings'
 require 'ddtrace/contrib/active_record/patcher'
+require 'ddtrace/contrib/rails/utils'
 
 module Datadog
   module Contrib
@@ -29,6 +31,12 @@ module Datadog
           super && version >= MINIMUM_VERSION
         end
 
+        # enabled by rails integration so should only auto instrument
+        # if detected that it is being used without rails
+        def auto_instrument?
+          !Datadog::Contrib::Rails::Utils.railtie_supported?
+        end
+
         def default_configuration
           Configuration::Settings.new
         end

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

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/contrib/patcher'
 require 'ddtrace/contrib/active_record/events'
 

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

@@ -1,3 +1,6 @@
+# typed: false
+require 'datadog/core/environment/ext'
+
 module Datadog
   module Contrib
     module ActiveRecord
@@ -21,44 +24,80 @@ module Datadog
           connection_config[:port]
         end
 
-        # In newer Rails versions, the `payload` contains both the `connection` and its `object_id` named `connection_id`.
-        #
-        # So, if rails is recent we'll have a direct access to the connection.
-        # Else, we'll find it thanks to the passed `connection_id`.
+        # Returns the connection configuration hash from the
+        # current connection
         #
-        # See this PR for more details: https://github.com/rails/rails/pull/34602
+        # Since Rails 6.0, we have direct access to the object,
+        # while older versions of Rails only provide us the
+        # connection id.
         #
+        # @see https://github.com/rails/rails/pull/34602
         def self.connection_config(connection = nil, connection_id = nil)
           return default_connection_config if connection.nil? && connection_id.nil?
 
           conn = if !connection.nil?
+                   # Since Rails 6.0, the connection object
+                   # is directly available.
                    connection
-                 # Rails 3.0 - 3.2
-                 elsif Gem.loaded_specs['activerecord'].version < Gem::Version.new('4.0')
-                   ::ActiveRecord::Base
-                     .connection_handler
-                     .connection_pools
-                     .values
-                     .flat_map(&:connections)
-                     .find { |c| c.object_id == connection_id }
-                 # Rails 4.2+
                  else
-                   ::ActiveRecord::Base
-                     .connection_handler
-                     .connection_pool_list
-                     .flat_map(&:connections)
-                     .find { |c| c.object_id == connection_id }
+                   # For Rails < 6.0, only the `connection_id`
+                   # is available. We have to find the connection
+                   # object from it.
+                   connection_from_id(connection_id)
                  end
 
-          if conn.instance_variable_defined?(:@config)
+          if conn && conn.instance_variable_defined?(:@config)
             conn.instance_variable_get(:@config)
           else
             EMPTY_CONFIG
           end
         end
 
+        # DEV: JRuby responds to {ObjectSpace._id2ref}, despite raising an error
+        # DEV: when invoked. Thus, we have to explicitly check for Ruby runtime.
+        if Datadog::Core::Environment::Ext::RUBY_ENGINE != 'jruby'
+          # CRuby has access to {ObjectSpace._id2ref}, which allows for
+          # direct look up of the connection object.
+          def self.connection_from_id(connection_id)
+            # `connection_id` is the `#object_id` of the
+            # connection. We can perform an ObjectSpace
+            # lookup to find it.
+            #
+            # This works not only for ActiveRecord, but for
+            # extensions that might have their own connection
+            # pool (e.g. https://rubygems.org/gems/makara).
+            ObjectSpace._id2ref(connection_id)
+          rescue => e
+            # Because `connection_id` references a live connection
+            # present in the current stack, it is very unlikely that
+            # `_id2ref` will fail, but we add this safeguard just
+            # in case.
+            Datadog.logger.debug(
+              "connection_id #{connection_id} does not represent a valid object. " \
+                      "Cause: #{e.message} Source: #{Array(e.backtrace).first}"
+            )
+          end
+        else
+          # JRuby does not enable {ObjectSpace._id2ref} by default,
+          # as it has large performance impact:
+          # https://github.com/jruby/jruby/wiki/PerformanceTuning/cf155dd9#dont-enable-objectspace
+          #
+          # This fallback code does not support the makara gem,
+          # as its connections don't live in the ActiveRecord
+          # connection pool.
+          def self.connection_from_id(connection_id)
+            ::ActiveRecord::Base
+              .connection_handler
+              .connection_pool_list
+              .flat_map(&:connections)
+              .find { |c| c.object_id == connection_id }
+          end
+        end
+
+        # @return [Hash]
         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
@@ -66,10 +105,19 @@ module Datadog
                                     end
 
           connection_pool = ::ActiveRecord::Base.connection_handler.retrieve_connection_pool(current_connection_name)
-          connection_pool.nil? ? EMPTY_CONFIG : (@default_connection_config = connection_pool.spec.config)
+          connection_pool.nil? ? EMPTY_CONFIG : (@default_connection_config = db_config(connection_pool))
         rescue StandardError
           EMPTY_CONFIG
         end
+
+        # @return [Hash]
+        def self.db_config(connection_pool)
+          if ::Rails::VERSION::MAJOR >= 6 && ::Rails::VERSION::MINOR >= 1
+            connection_pool.db_config.configuration_hash
+          else
+            connection_pool.spec.config
+          end
+        end
       end
     end
   end

data/lib/ddtrace/contrib/active_support/cache/instrumentation.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/contrib/active_support/ext'
 
 module Datadog
@@ -18,9 +19,13 @@ module Datadog
             # NOTE: the ``finish_trace_cache()`` is fired but it already has a safe-guard
             # to avoid any kind of issue.
             current_span = tracer.active_span
-            return if payload[:action] == Ext::RESOURCE_CACHE_GET &&
-                      current_span.try(:name) == Ext::SPAN_CACHE &&
-                      current_span.try(:resource) == Ext::RESOURCE_CACHE_GET
+            return if current_span.try(:name) == Ext::SPAN_CACHE &&
+                      (
+                        payload[:action] == Ext::RESOURCE_CACHE_GET &&
+                        current_span.try(:resource) == Ext::RESOURCE_CACHE_GET ||
+                        payload[:action] == Ext::RESOURCE_CACHE_MGET &&
+                        current_span.try(:resource) == Ext::RESOURCE_CACHE_MGET
+                      )
 
             tracing_context = payload.fetch(:tracing_context)
 
@@ -59,6 +64,32 @@ module Datadog
             Datadog.logger.debug(e.message)
           end
 
+          def finish_trace_cache_multi(payload)
+            # retrieve the tracing context and continue the trace
+            tracing_context = payload.fetch(:tracing_context)
+            span = tracing_context[:dd_cache_span]
+            return unless span && !span.finished?
+
+            begin
+              # discard parameters from the cache_store configuration
+              if defined?(::Rails)
+                store, = *Array.wrap(::Rails.configuration.cache_store).flatten
+                span.set_tag(Ext::TAG_CACHE_BACKEND, store)
+              end
+              normalized_keys = payload.fetch(:keys, []).map do |key|
+                ::ActiveSupport::Cache.expand_cache_key(key)
+              end
+              cache_keys = Datadog::Utils.truncate(normalized_keys, Ext::QUANTIZE_CACHE_MAX_KEY_SIZE)
+              span.set_tag(Ext::TAG_CACHE_KEY_MULTI, cache_keys)
+
+              span.set_error(payload[:exception]) if payload[:exception]
+            ensure
+              span.finish
+            end
+          rescue StandardError => e
+            Datadog.logger.debug(e.message)
+          end
+
           # Defines instrumentation for ActiveSupport cache reading
           module Read
             def read(*args, &block)
@@ -82,6 +113,29 @@ module Datadog
             end
           end
 
+          # Defines instrumentation for ActiveSupport cache reading of multiple keys
+          module ReadMulti
+            def read_multi(*keys, &block)
+              payload = {
+                action: Ext::RESOURCE_CACHE_MGET,
+                keys: keys,
+                tracing_context: {}
+              }
+
+              begin
+                # process and catch cache exceptions
+                Instrumentation.start_trace_cache(payload)
+                super
+              rescue Exception => e
+                payload[:exception] = [e.class.name, e.message]
+                payload[:exception_object] = e
+                raise e
+              end
+            ensure
+              Instrumentation.finish_trace_cache_multi(payload)
+            end
+          end
+
           # Defines instrumentation for ActiveSupport cache fetching
           module Fetch
             def fetch(*args, &block)
@@ -105,6 +159,31 @@ module Datadog
             end
           end
 
+          # Defines instrumentation for ActiveSupport cache fetching of multiple keys
+          module FetchMulti
+            def fetch_multi(*args, &block)
+              # extract options hash
+              keys = args[-1].instance_of?(Hash) ? args[0..-2] : args
+              payload = {
+                action: Ext::RESOURCE_CACHE_MGET,
+                keys: keys,
+                tracing_context: {}
+              }
+
+              begin
+                # process and catch cache exceptions
+                Instrumentation.start_trace_cache(payload)
+                super
+              rescue Exception => e
+                payload[:exception] = [e.class.name, e.message]
+                payload[:exception_object] = e
+                raise e
+              end
+            ensure
+              Instrumentation.finish_trace_cache_multi(payload)
+            end
+          end
+
           # Defines instrumentation for ActiveSupport cache writing
           module Write
             def write(*args, &block)
@@ -128,6 +207,29 @@ module Datadog
             end
           end
 
+          # Defines instrumentation for ActiveSupport cache writing of multiple keys
+          module WriteMulti
+            def write_multi(hash, options = nil)
+              payload = {
+                action: Ext::RESOURCE_CACHE_MSET,
+                keys: hash.keys,
+                tracing_context: {}
+              }
+
+              begin
+                # process and catch cache exceptions
+                Instrumentation.start_trace_cache(payload)
+                super
+              rescue Exception => e
+                payload[:exception] = [e.class.name, e.message]
+                payload[:exception_object] = e
+                raise e
+              end
+            ensure
+              Instrumentation.finish_trace_cache_multi(payload)
+            end
+          end
+
           # Defines instrumentation for ActiveSupport cache deleting
           module Delete
             def delete(*args, &block)