datadog 2.22.0 → 2.26.0

data/lib/datadog/tracing/contrib/waterdrop/producer.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require_relative 'ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module WaterDrop
+        # Producer integration for WaterDrop
+        module Producer
+          %i[
+            produce_many_sync
+            produce_many_async
+            produce_sync
+            produce_async
+          ].each do |method|
+            define_method(method) do |messages|
+              Datadog::Tracing.trace(Ext::SPAN_PRODUCER, resource: "waterdrop.#{__method__}") do
+                extract_span_tags(messages)
+                super(messages)
+              end
+            end
+          end
+
+          private
+
+          def extract_span_tags(messages)
+            messages = [messages] if messages.is_a?(Hash)
+            span = Datadog::Tracing.active_span
+            return unless span
+
+            topics = []
+            partitions = []
+            messages.each do |message|
+              topics << message[:topic]
+              partitions << message[:partition] if message.key?(:partition)
+            end
+
+            span.set_tag(Ext::TAG_PRODUCER, id)
+            span.set_tag(Contrib::Karafka::Ext::TAG_MESSAGE_COUNT, messages.size)
+            span.set_tag(Contrib::Ext::Messaging::TAG_SYSTEM, Contrib::Karafka::Ext::TAG_SYSTEM)
+
+            span.set_tag(Contrib::Ext::Messaging::TAG_DESTINATION, topics.uniq.sort.join(','))
+            span.set_tag(Contrib::Karafka::Ext::TAG_PARTITION, partitions.uniq.sort.join(',')) if partitions.any?
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/waterdrop.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require_relative 'component'
+require_relative 'waterdrop/integration'
+require_relative 'waterdrop/distributed/propagation'
+
+module Datadog
+  module Tracing
+    module Contrib
+      # `WaterDrop` integration public API
+      module WaterDrop
+        def self.inject(digest, data)
+          raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+          # Steep: https://github.com/soutaro/steep/issues/477
+          # @type ivar @propagation: WaterDrop::Distributed::Propagation
+          @propagation.inject!(digest, data)
+        end
+
+        def self.extract(data)
+          raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+          # Steep: https://github.com/soutaro/steep/issues/477
+          # @type ivar @propagation: WaterDrop::Distributed::Propagation
+          @propagation.extract(data)
+        end
+
+        Contrib::Component.register('waterdrop') do |config|
+          tracing = config.tracing
+          tracing.propagation_style
+
+          @propagation = WaterDrop::Distributed::Propagation.new(
+            propagation_style_inject: tracing.propagation_style_inject,
+            propagation_style_extract: tracing.propagation_style_extract,
+            propagation_extract_first: tracing.propagation_extract_first
+          )
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib.rb

@@ -55,6 +55,7 @@ require_relative 'contrib/httpclient/integration'
 require_relative 'contrib/httprb/integration'
 require_relative 'contrib/integration'
 require_relative 'contrib/kafka/integration'
+require_relative 'contrib/waterdrop'
 require_relative 'contrib/karafka'
 require_relative 'contrib/lograge/integration'
 require_relative 'contrib/mongodb/integration'

data/lib/datadog/tracing/diagnostics/environment_logger.rb

@@ -62,7 +62,7 @@ module Datadog
             # return `nil` with IO transport
             return unless transport.respond_to?(:client)
 
-            adapter = transport.client.api.adapter
+            adapter = transport.client.instance.adapter
             adapter.url
           end
 

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

@@ -175,8 +175,9 @@ module Datadog
 
           tags = {}
 
-          baggage_tag_keys.each do |key, _| # rubocop:disable Style/HashEachMethods
-            value = baggage[key]
+          baggage_tag_keys.each do |key, _|
+            # Steep: https://github.com/soutaro/steep/issues/2031
+            value = baggage[key] # steep:ignore ArgumentTypeMismatch
             next if value.nil? || value.empty?
 
             tags["baggage.#{key}"] = value

data/lib/datadog/tracing/metadata/ext.rb

@@ -94,8 +94,8 @@ module Datadog
 
         # @public_api
         module HTTP
-          ERROR_RANGE = (500...600).freeze
           TAG_BASE_URL = 'http.base_url'
+          TAG_ENDPOINT = 'http.endpoint'
           TAG_METHOD = 'http.method'
           TAG_STATUS_CODE = 'http.status_code'
           TAG_USER_AGENT = 'http.useragent'

data/lib/datadog/tracing/remote.rb

@@ -7,8 +7,6 @@ module Datadog
   module Tracing
     # Remote configuration declaration
     module Remote
-      class ReadError < StandardError; end
-
       class << self
         PRODUCT = 'APM_TRACING'
 
@@ -71,13 +69,7 @@ module Datadog
         private
 
         def parse_content(content)
-          data = content.data.read
-
-          content.data.rewind
-
-          raise ReadError, 'EOF reached' if data.nil?
-
-          JSON.parse(data)
+          JSON.parse(content.data)
         end
       end
     end

data/lib/datadog/tracing/sampling/priority_sampler.rb

@@ -60,7 +60,9 @@ module Datadog
         ensure
           if trace.sampling_priority && trace.sampling_priority > 0
             # Don't modify decision if priority was set upstream.
-            if !distributed_sampling_priority && !trace.has_tag?(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER)
+            # Steep: https://github.com/soutaro/steep/issues/1971
+            if !distributed_sampling_priority && # steep:ignore FallbackAny
+                !trace.has_tag?(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER)
               # If no sampling priority being assigned at this point, a custom
               # sampler implementation is configured: this means the user has
               # full control over the sampling decision.

data/lib/datadog/tracing/span.rb

@@ -80,7 +80,7 @@ module Datadog
 
         @meta = meta || {}
         @metrics = metrics || {}
-        @metastruct = metastruct || {}
+        @metastruct = metastruct || Metastruct.new
         @status = status || 0
 
         # start_time and end_time track wall clock. In Ruby, wall clock

data/lib/datadog/tracing/span_event.rb

@@ -69,8 +69,8 @@ module Datadog
 
       private
 
-      MIN_INT64_SIGNED = -2**63
-      MAX_INT64_SIGNED = 2 << 63 - 1
+      MIN_INT64_SIGNED = -2 << 62
+      MAX_INT64_SIGNED = (2 << 62) - 1
 
       # Checks the attributes hash to ensure it only contains serializable values.
       # Invalid values are removed from the hash.

data/lib/datadog/tracing/span_operation.rb

@@ -164,7 +164,9 @@ module Datadog
             # block is application code that we don't want to hinder.
             # * We don't yield during a fatal error, as the application is likely trying to
             #   end its execution (either due to a system error or graceful shutdown).
-            return_value = yield(self) unless e && !e.is_a?(StandardError)
+            # @type var e: Exception?
+            # Steep: https://github.com/soutaro/steep/issues/919
+            return_value = yield(self) unless e && !e.is_a?(StandardError) # steep:ignore FallbackAny
           end
         # rubocop:disable Lint/RescueException
         # Here we really want to catch *any* exception, not only StandardError,
@@ -214,9 +216,6 @@ module Datadog
       end
 
       # Mark the span stopped at the current time
-      #
-      # steep:ignore:start
-      # Steep issue fixed in https://github.com/soutaro/steep/pull/1467
       def stop(stop_time = nil, exception: nil)
         # A span should not be stopped twice. Note that this is not thread-safe,
         # stop is called from multiple threads, a given span might be stopped
@@ -239,7 +238,6 @@ module Datadog
 
         self
       end
-      # steep:ignore:end
 
       # Return whether the duration is started or not
       def started?
@@ -287,8 +285,14 @@ module Datadog
       end
 
       def duration
+        # Steep: https://github.com/soutaro/steep/issues/477
+        # @type ivar @duration_end: Time
+        # @type ivar @duration_start: Time
         return @duration_end - @duration_start if @duration_start && @duration_end
 
+        # Steep: https://github.com/soutaro/steep/issues/477
+        # @type ivar @end_time: Time
+        # @type ivar @start_time: Time
         @end_time - @start_time if @start_time && @end_time
       end
 
@@ -316,7 +320,9 @@ module Datadog
           'exception.stacktrace' => exc.backtrace,
         }
 
-        @span_events << SpanEvent.new('exception', attributes: event_attributes.merge!(attributes)) # steep:ignore
+        # Steep: Caused by wrong declaration, should be the same parameters as `merge`
+        # https://github.com/ruby/rbs/blob/3d0fb3a7fdde60af7120e875fe3bd7237b5b6a88/core/hash.rbs#L1468
+        @span_events << SpanEvent.new('exception', attributes: event_attributes.merge!(attributes)) # steep:ignore ArgumentTypeMismatch
       end
 
       # Return a string representation of the span.
@@ -387,7 +393,8 @@ module Datadog
       class Events
         include Tracing::Events
 
-        DEFAULT_ON_ERROR = proc { |span_op, error| span_op&.set_error(error) }
+        # Steep: https://github.com/soutaro/steep/issues/335
+        DEFAULT_ON_ERROR = proc { |span_op, error| span_op&.set_error(error) } # steep:ignore IncompatibleAssignment
 
         attr_reader \
           :logger,
@@ -395,7 +402,7 @@ module Datadog
           :after_stop,
           :before_start
 
-        def initialize(logger: Datadog.logger, on_error: nil)
+        def initialize(logger: Datadog.logger)
           @logger = logger
           @after_finish = AfterFinish.new
           @after_stop = AfterStop.new
@@ -464,7 +471,7 @@ module Datadog
               @handler.call(*args)
             rescue => e
               logger.debug do
-                "Error in on_error handler '#{@default}': #{e.class}: #{e} at #{Array(e.backtrace).first}"
+                "Error in on_error handler '#{@handler}': #{e.class}: #{e} at #{Array(e.backtrace).first}"
               end
             end
 
@@ -548,6 +555,10 @@ module Datadog
       # Used for serialization
       # @return [Integer] in nanoseconds since Epoch
       def start_time_nano
+        return 0 if @start_time.nil?
+
+        # Steep: https://github.com/soutaro/steep/issues/477
+        # @type ivar @start_time: Time
         @start_time.to_i * 1000000000 + @start_time.nsec
       end
 

data/lib/datadog/tracing/trace_operation.rb

@@ -55,6 +55,12 @@ module Datadog
         :sampled,
         :service
 
+      # Creates a new TraceOperation.
+      #
+      # @param auto_finish [Boolean] when true, automatically finishes the trace when the local root span finishes.
+      #   When false, the trace remains unfinished until {#finish!} is called.
+      #   This is useful when this {TraceOperation} represents the continuation of a remote {TraceDigest},
+      #   in which case local root spans in this {TraceOperation} are children of the {TraceDigest}'s last active span.
       def initialize(
         logger: Datadog.logger,
         agent_sample_rate: nil,
@@ -80,7 +86,8 @@ module Datadog
         trace_state_unknown_fields: nil,
         remote_parent: false,
         tracer: nil, # DEV-3.0: deprecated, remove in 3.0
-        baggage: nil
+        baggage: nil,
+        auto_finish: true
       )
         @logger = logger
 
@@ -119,6 +126,7 @@ module Datadog
         @events = events || Events.new
         @finished = false
         @spans = []
+        @auto_finish = !!auto_finish
       end
 
       def full?
@@ -318,6 +326,29 @@ module Datadog
         build_trace(spans, !finished)
       end
 
+      # When automatic context management is disabled (@auto_finish is false),
+      # this method finishes the trace, marking it as completed.
+      #
+      # The trace will **not** automatically finish when its local root span
+      # when @auto_finish is false, thus calling this method is mandatory
+      # in such scenario.
+      #
+      # Unfinished spans are discarded.
+      #
+      # This method is idempotent and safe to call after the trace is finished.
+      # It is also a no-op when @auto_finish is true, to prevent misuse.
+      #
+      # @!visibility private
+      def finish!
+        return if @auto_finish || finished?
+
+        @finished = true
+        @active_span = nil
+        @active_span_count = 0
+
+        events.trace_finished.publish(self)
+      end
+
       # Returns a set of trace headers used for continuing traces.
       # Used for propagation across execution contexts.
       # Data should reflect the active state of the trace.
@@ -460,7 +491,7 @@ module Datadog
 
         @active_span = span_op
 
-        set_root_span!(span_op) unless root_span
+        set_local_root_span!(span_op)
       end
 
       def deactivate_span!(span_op)
@@ -483,6 +514,12 @@ module Datadog
         logger.debug { "Error starting span on trace: #{e} Backtrace: #{e.backtrace.first(3)}" }
       end
 
+      # For traces with automatic context management (auto_finish),
+      # when the local root span finishes, the trace also finishes.
+      # The trace cannot receive new spans after finished.
+      #
+      # Without auto_finish, the trace can still receive spans
+      # until explicitly finished.
       def finish_span(span, span_op, parent)
         # Save finished span & root span
         @spans << span unless span.nil?
@@ -490,8 +527,9 @@ module Datadog
         # Deactivate the span, re-activate parent.
         deactivate_span!(span_op)
 
-        # Set finished, to signal root span has completed.
-        @finished = true if span_op == root_span
+        # Finish if the local root span is finished and automatic
+        # context management is enabled.
+        @finished = true if span_op == root_span && @auto_finish
 
         # Update active span count
         @active_span_count -= 1
@@ -505,8 +543,8 @@ module Datadog
         logger.debug { "Error finishing span on trace: #{e} Backtrace: #{e.backtrace.first(3)}" }
       end
 
-      # Track the root span
-      def set_root_span!(span)
+      # Track the root {SpanOperation} object from the current execution context.
+      def set_local_root_span!(span)
         return if span.nil? || root_span
 
         @root_span = span

data/lib/datadog/tracing/tracer.rb

@@ -129,7 +129,6 @@ module Datadog
       # @yield Optional block where new newly created {Datadog::Tracing::SpanOperation} captures the execution.
       # @yieldparam [Datadog::Tracing::SpanOperation] span_op the newly created and active [Datadog::Tracing::SpanOperation]
       # @yieldparam [Datadog::Tracing::TraceOperation] trace_op the active [Datadog::Tracing::TraceOperation]
-      # rubocop:disable Metrics/MethodLength
       def trace(
         name,
         continue_from: nil,
@@ -241,10 +240,21 @@ module Datadog
         trace.to_correlation
       end
 
-      # Setup a new trace to continue from where another
+      # Setup a new trace execution context to continue from where another
       # trace left off.
+      # This is useful to continue distributed or async traces.
       #
-      # Used to continue distributed or async traces.
+      # The first span created in the restored context is a direct child of the
+      # active span from when the {Datadog::Tracing::TraceDigest} was created.
+      #
+      # When no block is given, the trace context is restored in the current thread.
+      # It remains active until the first span created in this restored context is finished.
+      # After that, if a new span is created, it start a new, unrelated trace.
+      #
+      # When a block is given, the trace context is restored inside the block execution.
+      # It remains active until the block ends, even when the first span created inside
+      # the block finishes. This means that multiple spans can be direct children of the
+      # active span from when the {Datadog::Tracing::TraceDigest} was created.
       #
       # @param [Datadog::Tracing::TraceDigest] digest continue from the {Datadog::Tracing::TraceDigest}.
       # @param [Thread] key Thread to retrieve trace from. Defaults to current thread. For internal use only.
@@ -260,13 +270,32 @@ module Datadog
         # Start a new trace from the digest
         context = call_context(key)
         original_trace = active_trace(key)
-        trace = start_trace(continue_from: digest)
+        # When we want the trace to be bound to a block, we cannot let
+        # it auto finish when the local root span finishes. This would
+        # create mutiple traces inside the block. Instead, we'll
+        # expliclity finish the trace after the block finishes.
+        auto_finish = !block
+
+        trace = start_trace(continue_from: digest, auto_finish: auto_finish)
 
         # If block hasn't been given; we need to manually deactivate
         # this trace. Subscribe to the trace finished event to do this.
         subscribe_trace_deactivation!(context, trace, original_trace) unless block
 
-        context.activate!(trace, &block)
+        if block
+          # When a block is given, the trace will be active until the block finishes.
+          context.activate!(trace) do
+            yield
+          ensure # We have to flush even when an error occurs
+            # On block completion, force the trace to finish and flush its finished spans.
+            # Unfinished spans are lost as the {TraceOperation} has ended.
+            trace.finish!
+            flush_trace(trace)
+          end
+        else
+          # Otherwise, the trace will be bound to the current thread after this point
+          context.activate!(trace)
+        end
       end
 
       # Sample a span, tagging the trace as appropriate.
@@ -329,7 +358,7 @@ module Datadog
         @provider.context(key)
       end
 
-      def build_trace(digest = nil)
+      def build_trace(digest, auto_finish)
         # Resolve hostname if configured
         hostname = Core::Environment::Socket.hostname if Datadog.configuration.tracing.report_hostname
         hostname = (hostname && !hostname.empty?) ? hostname : nil
@@ -353,7 +382,8 @@ module Datadog
             trace_state_unknown_fields: digest.trace_state_unknown_fields,
             remote_parent: digest.span_remote,
             tracer: self,
-            baggage: digest.baggage
+            baggage: digest.baggage,
+            auto_finish: auto_finish
           )
         else
           TraceOperation.new(
@@ -362,13 +392,12 @@ module Datadog
             profiling_enabled: profiling_enabled,
             apm_tracing_enabled: apm_tracing_enabled,
             remote_parent: false,
-            tracer: self
+            tracer: self,
+            auto_finish: auto_finish
           )
         end
       end
-      # rubocop:enable Metrics/MethodLength
 
-      # rubocop:disable Metrics/MethodLength
       def bind_trace_events!(trace_op)
         events = trace_op.send(:events)
 
@@ -387,13 +416,12 @@ module Datadog
           flush_trace(event_trace_op)
         end
       end
-      # rubocop:enable Metrics/MethodLength
 
       # Creates a new TraceOperation, with events bounds to this Tracer instance.
       # @return [TraceOperation]
-      def start_trace(continue_from: nil)
+      def start_trace(continue_from: nil, auto_finish: true)
         # Build a new trace using digest if provided.
-        trace = build_trace(continue_from)
+        trace = build_trace(continue_from, auto_finish)
 
         # Bind trace events: sample trace, set default service, flush spans.
         bind_trace_events!(trace)
@@ -402,7 +430,6 @@ module Datadog
       end
 
       # rubocop:disable Lint/UnderscorePrefixedVariableName
-      # rubocop:disable Metrics/MethodLength
       def start_span(
         name,
         continue_from: nil,
@@ -454,9 +481,8 @@ module Datadog
           span
         end
       end
-      # rubocop:enable Lint/UnderscorePrefixedVariableName
-      # rubocop:enable Metrics/MethodLength
 
+      # rubocop:enable Lint/UnderscorePrefixedVariableName
       def resolve_tags(tags, service)
         merged_tags = if @tags.any? && tags
           # Combine default tags with provided tags,

data/lib/datadog/tracing/transport/http/client.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative 'statistics'
+require_relative '../../../core/transport/http/client'
 require_relative '../../../core/transport/http/env'
 require_relative '../../../core/transport/http/response'
 
@@ -9,33 +10,24 @@ module Datadog
     module Transport
       module HTTP
         # Routes, encodes, and sends tracer data to the trace agent via HTTP.
-        class Client
+        class Client < Core::Transport::HTTP::Client
           include Datadog::Tracing::Transport::HTTP::Statistics
 
-          attr_reader :api, :logger
+          private
 
-          def initialize(api, logger: Datadog.logger)
-            @api = api
-            @logger = logger
-          end
-
-          def send_request(request, &block)
-            # Build request into env
-            env = build_env(request)
-
-            # Get responses from API
-            response = yield(api, env)
+          def on_response(response)
+            super
 
             # Update statistics
             update_stats_from_response!(response)
+          end
 
-            response
-          rescue => e
-            message =
-              "Internal error during #{self.class.name} request. Cause: #{e.class.name} #{e.message} " \
-                "Location: #{Array(e.backtrace).first}"
+          def on_exception(exception)
+            # Note: this method does NOT call super - it has replacement
+            # logic for how to log the exception.
+
+            message = build_exception_message(exception)
 
-            # Log error
             if stats.consecutive_errors > 0
               logger.debug(message)
             else
@@ -44,13 +36,7 @@ module Datadog
             end
 
             # Update statistics
-            update_stats_from_exception!(e)
-
-            Datadog::Core::Transport::InternalErrorResponse.new(e)
-          end
-
-          def build_env(request)
-            Datadog::Core::Transport::HTTP::Env.new(request)
+            update_stats_from_exception!(exception)
           end
         end
       end

data/lib/datadog/tracing/transport/http/traces.rb

@@ -2,12 +2,10 @@
 
 require 'json'
 
-require_relative '../traces'
-require_relative 'client'
 require_relative '../../../core/transport/http/response'
 require_relative '../../../core/transport/http/api/endpoint'
-require_relative '../../../core/transport/http/api/spec'
 require_relative '../../../core/transport/http/api/instance'
+require_relative '../traces'
 
 module Datadog
   module Tracing
@@ -27,58 +25,15 @@ module Datadog
             end
           end
 
-          # Extensions for HTTP client
-          module Client
-            def send_traces_payload(request)
-              send_request(request) do |api, env|
-                api.send_traces(env)
-              end
-            end
-          end
-
           module API
-            # HTTP API Spec
-            class Spec < Core::Transport::HTTP::API::Spec
-              attr_accessor :traces
-
-              def send_traces(env, &block)
-                raise Core::Transport::HTTP::API::Spec::EndpointNotDefinedError.new('traces', self) if traces.nil?
-
-                traces.call(env, &block)
-              end
-
-              def encoder
-                traces.encoder
-              end
-            end
-
-            # HTTP API Instance
-            class Instance < Core::Transport::HTTP::API::Instance
-              def send_traces(env)
-                unless spec.is_a?(Traces::API::Spec)
-                  raise Core::Transport::HTTP::API::Instance::EndpointNotSupportedError.new(
-                    'traces', self
-                  )
-                end
-
-                spec.send_traces(env) do |request_env|
-                  call(request_env)
-                end
-              end
-            end
-
             # Endpoint for submitting trace data
             class Endpoint < Datadog::Core::Transport::HTTP::API::Endpoint
               HEADER_CONTENT_TYPE = 'Content-Type'
               HEADER_TRACE_COUNT = 'X-Datadog-Trace-Count'
               SERVICE_RATE_KEY = 'rate_by_service'
 
-              attr_reader \
-                :encoder
-
               def initialize(path, encoder, options = {})
-                super(:post, path)
-                @encoder = encoder
+                super(:post, path, encoder: encoder)
                 @service_rates = options.fetch(:service_rates, false)
               end
 
@@ -111,9 +66,6 @@ module Datadog
               end
             end
           end
-
-          # Add traces behavior to transport components
-          HTTP::Client.include(Traces::Client)
         end
       end
     end