datadog 2.1.0 → 2.2.0

data/lib/datadog/tracing/contrib/graphql/configuration/settings.rb

@@ -42,6 +42,11 @@ module Datadog
               o.type :bool
               o.default false
             end
+
+            option :with_unified_tracer do |o|
+              o.type :bool
+              o.default false
+            end
           end
         end
       end

data/lib/datadog/tracing/contrib/graphql/patcher.rb

@@ -4,6 +4,7 @@ require_relative '../analytics'
 require_relative '../patcher'
 require_relative 'tracing_patcher'
 require_relative 'trace_patcher'
+require_relative 'unified_trace_patcher'
 
 module Datadog
   module Tracing
@@ -23,10 +24,15 @@ module Datadog
             if configuration[:with_deprecated_tracer]
               TracingPatcher.patch!(schemas, trace_options)
             elsif Integration.trace_supported?
-              TracePatcher.patch!(schemas, trace_options)
+              if configuration[:with_unified_tracer]
+                UnifiedTracePatcher.patch!(schemas, trace_options)
+              else
+                TracePatcher.patch!(schemas, trace_options)
+              end
             else
               Datadog.logger.warn(
-                "GraphQL version (#{target_version}) does not support GraphQL::Tracing::DataDogTrace. "\
+                "GraphQL version (#{target_version}) does not support GraphQL::Tracing::DataDogTrace"\
+                'or Datadog::Tracing::Contrib::GraphQL::UnifiedTrace.'\
                 'Falling back to GraphQL::Tracing::DataDogTracing.'
               )
               TracingPatcher.patch!(schemas, trace_options)

data/lib/datadog/tracing/contrib/graphql/unified_trace.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+
+require 'graphql/tracing'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module GraphQL
+        # These methods will be called by the GraphQL runtime to trace the execution of queries.
+        # This tracer differs from the upstream one as it follows the unified naming convention specification,
+        # which is required to use features such as API Catalog.
+        # DEV-3.0: This tracer should be the default one in the next major version.
+        module UnifiedTrace
+          # @param analytics_enabled [Boolean] Deprecated
+          # @param analytics_sample_rate [Float] Deprecated
+          # @param service [String|nil] The service name to be set on the spans
+          def initialize(*args, analytics_enabled: false, analytics_sample_rate: 1.0, service: nil, **kwargs)
+            @analytics_enabled = analytics_enabled
+            @analytics_sample_rate = analytics_sample_rate
+
+            @service_name = service
+            @has_prepare_span = respond_to?(:prepare_span)
+            super
+          end
+
+          def lex(*args, query_string:, **kwargs)
+            trace(proc { super }, 'lex', query_string, query_string: query_string)
+          end
+
+          def parse(*args, query_string:, **kwargs)
+            trace(proc { super }, 'parse', query_string, query_string: query_string) do |span|
+              span.set_tag('graphql.source', query_string)
+            end
+          end
+
+          def validate(*args, query:, validate:, **kwargs)
+            trace(proc { super }, 'validate', query.selected_operation_name, query: query, validate: validate) do |span|
+              span.set_tag('graphql.source', query.query_string)
+            end
+          end
+
+          def analyze_multiplex(*args, multiplex:, **kwargs)
+            trace(proc { super }, 'analyze_multiplex', multiplex_resource(multiplex), multiplex: multiplex)
+          end
+
+          def analyze_query(*args, query:, **kwargs)
+            trace(proc { super }, 'analyze', query.query_string, query: query)
+          end
+
+          def execute_multiplex(*args, multiplex:, **kwargs)
+            trace(proc { super }, 'execute_multiplex', multiplex_resource(multiplex), multiplex: multiplex) do |span|
+              span.set_tag('graphql.source', "Multiplex[#{multiplex.queries.map(&:query_string).join(', ')}]")
+            end
+          end
+
+          def execute_query(*args, query:, **kwargs)
+            trace(proc { super }, 'execute', query.selected_operation_name, query: query) do |span|
+              span.set_tag('graphql.source', query.query_string)
+              span.set_tag('graphql.operation.type', query.selected_operation.operation_type)
+              span.set_tag('graphql.operation.name', query.selected_operation_name) if query.selected_operation_name
+              query.variables.instance_variable_get(:@storage).each do |key, value|
+                span.set_tag("graphql.variables.#{key}", value)
+              end
+            end
+          end
+
+          def execute_query_lazy(*args, query:, multiplex:, **kwargs)
+            resource = if query
+                         query.selected_operation_name || fallback_transaction_name(query.context)
+                       else
+                         multiplex_resource(multiplex)
+                       end
+            trace(proc { super }, 'execute_lazy', resource, query: query, multiplex: multiplex)
+          end
+
+          def execute_field_span(callable, span_key, **kwargs)
+            # @platform_key_cache is initialized upstream, in ::GraphQL::Tracing::PlatformTrace
+            platform_key = @platform_key_cache[UnifiedTrace].platform_field_key_cache[kwargs[:field]]
+
+            if platform_key
+              trace(callable, span_key, platform_key, **kwargs) do |span|
+                kwargs[:arguments].each do |key, value|
+                  span.set_tag("graphql.variables.#{key}", value)
+                end
+              end
+            else
+              callable.call
+            end
+          end
+
+          def execute_field(*args, **kwargs)
+            execute_field_span(proc { super }, 'resolve', **kwargs)
+          end
+
+          def execute_field_lazy(*args, **kwargs)
+            execute_field_span(proc { super }, 'resolve_lazy', **kwargs)
+          end
+
+          def authorized_span(callable, span_key, **kwargs)
+            platform_key = @platform_key_cache[UnifiedTrace].platform_authorized_key_cache[kwargs[:type]]
+            trace(callable, span_key, platform_key, **kwargs)
+          end
+
+          def authorized(*args, **kwargs)
+            authorized_span(proc { super }, 'authorized', **kwargs)
+          end
+
+          def authorized_lazy(*args, **kwargs)
+            authorized_span(proc { super }, 'authorized_lazy', **kwargs)
+          end
+
+          def resolve_type_span(callable, span_key, **kwargs)
+            platform_key = @platform_key_cache[UnifiedTrace].platform_resolve_type_key_cache[kwargs[:type]]
+            trace(callable, span_key, platform_key, **kwargs)
+          end
+
+          def resolve_type(*args, **kwargs)
+            resolve_type_span(proc { super }, 'resolve_type', **kwargs)
+          end
+
+          def resolve_type_lazy(*args, **kwargs)
+            resolve_type_span(proc { super }, 'resolve_type_lazy', **kwargs)
+          end
+
+          include ::GraphQL::Tracing::PlatformTrace
+
+          def platform_field_key(field, *args, **kwargs)
+            field.path
+          end
+
+          def platform_authorized_key(type, *args, **kwargs)
+            "#{type.graphql_name}.authorized"
+          end
+
+          def platform_resolve_type_key(type, *args, **kwargs)
+            "#{type.graphql_name}.resolve_type"
+          end
+
+          private
+
+          def trace(callable, trace_key, resource, **kwargs)
+            Tracing.trace("graphql.#{trace_key}", resource: resource, service: @service_name, type: 'graphql') do |span|
+              yield(span) if block_given?
+
+              prepare_span(trace_key, kwargs, span) if @has_prepare_span
+
+              callable.call
+            end
+          end
+
+          def multiplex_resource(multiplex)
+            return nil unless multiplex
+
+            operations = multiplex.queries.map(&:selected_operation_name).compact.join(', ')
+            if operations.empty?
+              first_query = multiplex.queries.first
+              fallback_transaction_name(first_query && first_query.context)
+            else
+              operations
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module GraphQL
+        # Provides instrumentation for `graphql` through the GraphQL's tracing with methods defined in UnifiedTrace
+        module UnifiedTracePatcher
+          module_function
+
+          def patch!(schemas, options)
+            require_relative 'unified_trace'
+            if schemas.empty?
+              ::GraphQL::Schema.trace_with(UnifiedTrace, **options)
+            else
+              schemas.each do |schema|
+                schema.trace_with(UnifiedTrace, **options)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/kafka/consumer_event.rb

@@ -6,7 +6,7 @@ module Datadog
       module Kafka
         # Defines basic behaviors for an event for a consumer.
         module ConsumerEvent
-          def process(span, _event, _id, payload)
+          def on_start(span, _event, _id, payload)
             super
 
             span.set_tag(Ext::TAG_GROUP, payload[:group_id])

data/lib/datadog/tracing/contrib/kafka/consumer_group_event.rb

@@ -6,7 +6,7 @@ module Datadog
       module Kafka
         # Defines basic behaviors for an event for a consumer group.
         module ConsumerGroupEvent
-          def process(span, _event, _id, payload)
+          def on_start(span, _event, _id, payload)
             super
 
             span.resource = payload[:group_id]

data/lib/datadog/tracing/contrib/kafka/event.rb

@@ -29,7 +29,7 @@ module Datadog
               Datadog.configuration.tracing[:kafka]
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span.set_tag(Contrib::Ext::Messaging::TAG_SYSTEM, Ext::TAG_MESSAGING_SYSTEM)
 

data/lib/datadog/tracing/contrib/kafka/events/connection/request.rb

@@ -15,7 +15,9 @@ module Datadog
 
               EVENT_NAME = 'request.connection.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.resource = payload[:api]
@@ -24,8 +26,6 @@ module Datadog
                 span.set_tag(Ext::TAG_RESPONSE_SIZE, payload[:response_size]) if payload.key?(:response_size)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_CONNECTION_REQUEST
               end

data/lib/datadog/tracing/contrib/kafka/events/consumer/process_batch.rb

@@ -17,7 +17,9 @@ module Datadog
 
               EVENT_NAME = 'process_batch.consumer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.resource = payload[:topic]
@@ -31,8 +33,6 @@ module Datadog
                 span.set_tag(Ext::TAG_OFFSET_LAG, payload[:offset_lag]) if payload.key?(:offset_lag)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_PROCESS_BATCH
               end

data/lib/datadog/tracing/contrib/kafka/events/consumer/process_message.rb

@@ -17,7 +17,9 @@ module Datadog
 
               EVENT_NAME = 'process_message.consumer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.resource = payload[:topic]
@@ -29,8 +31,6 @@ module Datadog
                 span.set_tag(Ext::TAG_OFFSET_LAG, payload[:offset_lag]) if payload.key?(:offset_lag)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_PROCESS_MESSAGE
               end

data/lib/datadog/tracing/contrib/kafka/events/consumer_group/heartbeat.rb

@@ -19,7 +19,9 @@ module Datadog
 
               EVENT_NAME = 'heartbeat.consumer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 if payload.key?(:topic_partitions)
@@ -29,8 +31,6 @@ module Datadog
                 end
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_CONSUMER_HEARTBEAT
               end

data/lib/datadog/tracing/contrib/kafka/events/produce_operation/send_messages.rb

@@ -15,7 +15,9 @@ module Datadog
 
               EVENT_NAME = 'send_messages.producer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.set_tag(Ext::TAG_MESSAGE_COUNT, payload[:message_count]) if payload.key?(:message_count)
@@ -23,8 +25,6 @@ module Datadog
                 span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_PRODUCER)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_SEND_MESSAGES
               end

data/lib/datadog/tracing/contrib/kafka/events/producer/deliver_messages.rb

@@ -15,7 +15,9 @@ module Datadog
 
               EVENT_NAME = 'deliver_messages.producer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.set_tag(Ext::TAG_ATTEMPTS, payload[:attempts]) if payload.key?(:attempts)
@@ -26,8 +28,6 @@ module Datadog
                 span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_PRODUCER)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_DELIVER_MESSAGES
               end

data/lib/datadog/tracing/contrib/racecar/event.rb

@@ -19,7 +19,7 @@ module Datadog
           # Class methods for Racecar events.
           # Note, they share the same process method and before_trace method.
           module ClassMethods
-            def subscription(*args)
+            def subscription(*args, **kwargs)
               super.tap do |subscription|
                 subscription.before_trace { ensure_clean_context! }
               end
@@ -33,7 +33,7 @@ module Datadog
               Datadog.configuration.tracing[:racecar]
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.service = configuration[:service_name]
               span.resource = payload[:consumer_class]
 

data/lib/datadog/tracing/contrib/rails/ext.rb

@@ -14,6 +14,15 @@ module Datadog
           ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_RAILS_ANALYTICS_SAMPLE_RATE'
           ENV_DISABLE = 'DISABLE_DATADOG_RAILS'
 
+          SPAN_RUNNER_FILE = 'rails.runner.file'
+          SPAN_RUNNER_INLINE = 'rails.runner.inline'
+          SPAN_RUNNER_STDIN = 'rails.runner.stdin'
+          TAG_COMPONENT = 'rails'
+          TAG_OPERATION_FILE = 'runner.file'
+          TAG_OPERATION_INLINE = 'runner.inline'
+          TAG_OPERATION_STDIN = 'runner.stdin'
+          TAG_RUNNER_SOURCE = 'source'
+
           # @!visibility private
           MINIMUM_VERSION = Gem::Version.new('4')
         end

data/lib/datadog/tracing/contrib/rails/patcher.rb

@@ -5,6 +5,7 @@ require_relative '../rack/middlewares'
 require_relative 'framework'
 require_relative 'log_injection'
 require_relative 'middlewares'
+require_relative 'runner'
 require_relative 'utils'
 require_relative '../semantic_logger/patcher'
 
@@ -28,6 +29,7 @@ module Datadog
           def patch
             patch_before_initialize
             patch_after_initialize
+            patch_rails_runner
           end
 
           def patch_before_initialize
@@ -81,6 +83,11 @@ module Datadog
           def setup_tracer
             Contrib::Rails::Framework.setup
           end
+
+          # Instruments the `bin/rails runner` command.
+          def patch_rails_runner
+            ::Rails::Command.singleton_class.prepend(Command) if defined?(::Rails::Command)
+          end
         end
       end
     end

data/lib/datadog/tracing/contrib/rails/runner.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Rails
+        # Instruments the `bin/rails runner` command.
+        # This command executes the provided code with the host Rails application loaded.
+        # The command can be either:
+        # * `-`: for code provided through the STDIN.
+        # * File path: for code provided through a local file.
+        # * `inline code`: for code provided directly as a command line argument.
+        # @see https://guides.rubyonrails.org/v6.1/command_line.html#bin-rails-runner
+        module Runner
+          # Limit the maximum size of the source code captured in the source tag.
+          MAX_TAG_VALUE_SIZE = 4096
+          private_constant :MAX_TAG_VALUE_SIZE
+
+          def runner(code_or_file = nil, *_command_argv)
+            if code_or_file == '-'
+              name = Ext::SPAN_RUNNER_STDIN
+              resource = nil
+              operation = Ext::TAG_OPERATION_STDIN
+              # The source is not yet available for STDIN, but it will be captured in `eval`.
+            elsif File.exist?(code_or_file)
+              name = Ext::SPAN_RUNNER_FILE
+              resource = code_or_file
+              operation = Ext::TAG_OPERATION_FILE
+              source = File.read(code_or_file)
+            else
+              name = Ext::SPAN_RUNNER_INLINE
+              resource = nil
+              operation = Ext::TAG_OPERATION_INLINE
+              source = code_or_file
+            end
+
+            Tracing.trace(
+              name,
+              service: Datadog.configuration.tracing[:rails][:service_name],
+              resource: resource,
+              tags: {
+                Tracing::Metadata::Ext::TAG_COMPONENT => Ext::TAG_COMPONENT,
+                Tracing::Metadata::Ext::TAG_OPERATION => operation,
+              }
+            ) do |span|
+              if source
+                span.set_tag(
+                  Ext::TAG_RUNNER_SOURCE,
+                  Core::Utils.truncate(source, MAX_TAG_VALUE_SIZE)
+                )
+              end
+              Contrib::Analytics.set_rate!(span, Datadog.configuration.tracing[:rails])
+
+              super
+            end
+          end
+
+          # Capture the executed source code when provided from STDIN.
+          def eval(*args)
+            span = Datadog::Tracing.active_span
+            if span.name == Ext::SPAN_RUNNER_STDIN
+              source = args[0]
+              span.set_tag(
+                Ext::TAG_RUNNER_SOURCE,
+                Core::Utils.truncate(source, MAX_TAG_VALUE_SIZE)
+              )
+            end
+
+            super
+          end
+
+          ruby2_keywords :eval if respond_to?(:ruby2_keywords, true)
+        end
+
+        # The instrumentation target, {Rails::Command::RunnerCommand} is only loaded
+        # right before `bin/rails runner` is executed. This means there's not much
+        # opportunity to patch it ahead of time.
+        # To ensure we can patch it successfully, we patch it's caller, {Rails::Command}
+        # and promptly patch {Rails::Command::RunnerCommand} when it is loaded.
+        module Command
+          def find_by_namespace(*args)
+            ret = super
+            # Patch RunnerCommand if it is loaded and not already patched.
+            if defined?(::Rails::Command::RunnerCommand) && !(::Rails::Command::RunnerCommand < Runner)
+              ::Rails::Command::RunnerCommand.prepend(Runner)
+            end
+            ret
+          end
+
+          ruby2_keywords :find_by_namespace if respond_to?(:ruby2_keywords, true)
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/distributed/b3_multi.rb

@@ -31,7 +31,7 @@ module Datadog
 
           # DEV: We need these to be hex encoded
           data[@trace_id_key] = format('%032x', digest.trace_id)
-          data[@span_id_key] = format('%016x', digest.span_id)
+          data[@span_id_key] = format('%016x', digest.span_id || 0) # # Fall back to zero (invalid) if not present
 
           if digest.trace_sampling_priority
             sampling_priority = Helpers.clamp_sampling_priority(

data/lib/datadog/tracing/distributed/b3_single.rb

@@ -25,8 +25,10 @@ module Datadog
         def inject!(digest, env)
           return if digest.nil?
 
+          span_id = digest.span_id || 0 # Fall back to zero (invalid) if not present
+
           # DEV: We need these to be hex encoded
-          value = "#{format('%032x', digest.trace_id)}-#{format('%016x', digest.span_id)}"
+          value = "#{format('%032x', digest.trace_id)}-#{format('%016x', span_id)}"
 
           if digest.trace_sampling_priority
             sampling_priority = Helpers.clamp_sampling_priority(

data/lib/datadog/tracing/distributed/datadog.rb

@@ -42,7 +42,7 @@ module Datadog
 
           data[@trace_id_key] = Tracing::Utils::TraceId.to_low_order(digest.trace_id).to_s
 
-          data[@parent_id_key] = digest.span_id.to_s
+          data[@parent_id_key] = digest.span_id.to_s if digest.span_id
           data[@sampling_priority_key] = digest.trace_sampling_priority.to_s if digest.trace_sampling_priority
           data[@origin_key] = digest.trace_origin.to_s if digest.trace_origin
 
@@ -109,7 +109,7 @@ module Datadog
 
           return tags if high_order == 0
 
-          tags.merge(Tracing::Metadata::Ext::Distributed::TAG_TID => high_order.to_s(16))
+          tags.merge(Tracing::Metadata::Ext::Distributed::TAG_TID => format('%016x', high_order))
         end
 
         # Side effect: Remove high order 64 bit hex-encoded `tid` tag from distributed tags

data/lib/datadog/tracing/distributed/propagation.rb

@@ -32,7 +32,9 @@ module Datadog
 
         # inject! populates the env with span ID, trace ID and sampling priority
         #
-        # This method will never raise errors, but instead log them to `Datadog.logger`.
+        # This method will never raise errors.
+        # It can propagate partial data, if deemed useful, instead of failing.
+        # In case of unrecoverable errors, it will log them to `Datadog.logger`.
         #
         # DEV-2.0: inject! should work without arguments, injecting the active_trace's digest
         # DEV-2.0: and returning a new Hash with the injected data.
@@ -45,7 +47,7 @@ module Datadog
         # @param digest [TraceDigest]
         # @param data [Hash]
         # @return [Boolean] `true` if injected successfully, `false` if no propagation style is configured
-        # @return [nil] in case of error, see `Datadog.logger` output for details.
+        # @return [nil] in case of unrecoverable errors, see `Datadog.logger` output for details.
         def inject!(digest, data)
           if digest.nil?
             ::Datadog.logger.debug('Cannot inject distributed trace data: digest is nil.')
@@ -54,6 +56,11 @@ module Datadog
 
           digest = digest.to_digest if digest.respond_to?(:to_digest)
 
+          if digest.trace_id.nil?
+            ::Datadog.logger.debug('Cannot inject distributed trace data: digest.trace_id is nil.')
+            return nil
+          end
+
           result = false
 
           # Inject all configured propagation styles

data/lib/datadog/tracing/distributed/trace_context.rb

@@ -107,7 +107,7 @@ module Datadog
         def build_traceparent(digest)
           build_traceparent_string(
             digest.trace_id,
-            digest.span_id,
+            digest.span_id || 0, # Fall back to zero (invalid) if not present
             build_trace_flags(digest)
           )
         end
@@ -198,7 +198,8 @@ module Datadog
 
         def last_dd_parent_id(digest)
           if !digest.span_remote
-            "p:#{format('%016x', digest.span_id)};"
+            span_id = digest.span_id || 0 # Fall back to zero (invalid) if not present
+            "p:#{format('%016x', span_id)};"
           elsif digest.trace_distributed_tags&.key?(Tracing::Metadata::Ext::Distributed::TAG_DD_PARENT_ID)
             "p:#{digest.trace_distributed_tags[Tracing::Metadata::Ext::Distributed::TAG_DD_PARENT_ID]};"
           else

data/lib/datadog/tracing/span_operation.rb

@@ -48,7 +48,8 @@ module Datadog
         tags: nil,
         trace_id: nil,
         type: nil,
-        links: nil
+        links: nil,
+        id: nil
       )
         # Ensure dynamically created strings are UTF-8 encoded.
         #
@@ -60,7 +61,7 @@ module Datadog
         self.type = type
         self.resource = resource
 
-        @id = Tracing::Utils.next_id
+        @id = id.nil? ? Tracing::Utils.next_id : id
         @parent_id = parent_id || 0
         @trace_id = trace_id || Tracing::Utils::TraceId.next_id