datadog 2.21.0 → 2.38.0

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

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require_relative '../patcher'
+require_relative 'ext'
+require_relative 'distributed/propagation'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module WaterDrop
+        # Patcher enables patching of 'waterdrop' module.
+        module Patcher
+          include Contrib::Patcher
+
+          module_function
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            require_relative 'producer'
+            require_relative 'middleware'
+
+            ::WaterDrop::Producer.prepend(Producer)
+            ::WaterDrop.instrumentation.subscribe('producer.configured') do |event|
+              producer = event[:producer]
+              add_middleware(producer)
+
+              if Datadog.configuration.data_streams.enabled
+                producer.monitor.subscribe('message.acknowledged') do |ack_event|
+                  if Datadog::DataStreams.enabled?
+                    payload = ack_event.payload
+                    Datadog::DataStreams.track_kafka_produce(payload[:topic], payload[:partition], payload[:offset])
+                  end
+                end
+              end
+            end
+          end
+
+          def add_middleware(producer)
+            included_middlewares = producer.middleware.instance_variable_get(:@steps)
+            producer.middleware.append(Middleware) unless included_middlewares.include?(Middleware)
+          end
+        end
+      end
+    end
+  end
+end

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,42 @@
+# 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_behavior_extract: tracing.propagation_behavior_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'
@@ -80,3 +81,11 @@ require_relative 'contrib/sneakers/integration'
 require_relative 'contrib/stripe/integration'
 require_relative 'contrib/sucker_punch/integration'
 require_relative 'contrib/trilogy/integration'
+
+# This list is used to determine if an integration is a built-in integration,
+# and prevent sending telemetry for custom integrations.
+# To do this, we get the list of registered integrations at this point and freeze it.
+# Later, when custom integrations are registered, they will not be included in this list.
+# .uniq to handle aliases, as we can register the same integration multiple times with different names.
+Datadog::Tracing::Contrib::BUILT_IN_INTEGRATIONS =
+  Datadog::Tracing::Contrib::REGISTRY.map { |entry| entry.klass }.uniq.freeze

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

@@ -24,7 +24,9 @@ module Datadog
             end
           end
         rescue => e
-          logger.warn("Failed to collect tracing environment information: #{e} Location: #{Array(e.backtrace).first}")
+          logger.warn(
+            "Failed to collect tracing environment information: #{e.class}: #{e.message} Location: #{Array(e.backtrace).first}"
+          )
         end
       end
 
@@ -62,7 +64,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

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative '../ext'
 require_relative '../metadata/ext'
 require_relative '../trace_digest'
 require_relative 'datadog_tags_codec'
@@ -77,7 +78,7 @@ module Datadog
             # Record telemetry for successful injection
             record_telemetry_metric('context_header_style.injected', 1, {'header_style' => 'baggage'})
           rescue => e
-            ::Datadog.logger.warn("Failed to encode and inject baggage header: #{e.class}: #{e}")
+            ::Datadog.logger.warn("Failed to encode and inject baggage header: #{e.class}: #{e.message}")
           end
         end
 
@@ -133,12 +134,56 @@ module Datadog
         # - Keys and values are URL-encoded
         # - Returns an empty hash if the baggage header is malformed
         #
+        # For entries with duplicate keys, only last one is preserved:
+        # https://www.w3.org/TR/2024/CR-baggage-20240530/#mutating-baggage
+        #
         # @param baggage_header [String] The W3C Baggage header string to parse
         # @return [Hash<String, String>] A hash of decoded baggage items
         def parse_baggage_header(baggage_header)
+          # Reject inputs that would cause String operations below to raise.
+          # `valid_encoding?` is true for ASCII-8BIT regardless of byte values, so
+          # the typical Rack path is unaffected; only UTF-8-tagged-but-actually-invalid
+          # input (e.g. headers re-encoded upstream) is short-circuited here.
+          unless baggage_header.valid_encoding?
+            record_telemetry_metric('context_header_style.malformed', 1, {'header_style' => 'baggage'})
+            return {}
+          end
+
+          if baggage_header.bytesize > DD_TRACE_BAGGAGE_MAX_BYTES
+            record_telemetry_metric('context_header.truncated', 1, {'header_style' => 'baggage', 'truncation_reason' => 'baggage_byte_count_exceeded'})
+
+            # We MUST NOT propagate partial entries, but we SHOULD try
+            # to parse as much of the baggage as possible:
+            # https://www.w3.org/TR/2024/CR-baggage-20240530/#limits
+
+            # We parse 1 byte over the limit to detect if the last entry
+            # is a partial entry (toss) or ends exactly at the limit (keep).
+            remove_last_entry = baggage_header.byteslice(DD_TRACE_BAGGAGE_MAX_BYTES, 1) != ','
+
+            # To ensure we don't have a trailing partial UTF-8 codepoint, we keep one extra byte
+            # and safely remove it with `#chop`.
+            # `#chops` walks back the string until it finds a valid character boundary and deletes
+            # from there.
+            baggage_header = baggage_header.byteslice(0, DD_TRACE_BAGGAGE_MAX_BYTES + 1) #: String
+            baggage_header = baggage_header.chop
+          end
+
           baggage = {}
-          baggages = baggage_header.split(',')
-          baggages.each do |key_value|
+          # DEV: To avoid unnecessary eager string allocation, replace with
+          # DEV: `split(',') { |s| ... }` when Ruby 2.5 is no longer supported.
+          baggages = baggage_header.split(',', DD_TRACE_BAGGAGE_MAX_ITEMS + 1) # Stop splitting if we've reached max size
+          baggages.each_with_index do |key_value, index|
+            if index >= DD_TRACE_BAGGAGE_MAX_ITEMS
+              record_telemetry_metric('context_header.truncated', 1, {'header_style' => 'baggage', 'truncation_reason' => 'baggage_item_count_exceeded'})
+              break
+            end
+
+            # On truncation, remove incomplete trailing partial entry
+            next if remove_last_entry && index == baggages.size - 1
+
+            # Empty items are skipped
+            next if key_value.strip.empty?
+
             key, value = key_value.split('=', 2)
             # If baggage is malformed, return an empty hash
             if key.nil? || value.nil?
@@ -147,8 +192,17 @@ module Datadog
               return {}
             end
 
-            key = URI.decode_www_form_component(key.strip)
-            value = URI.decode_www_form_component(value.strip)
+            begin
+              key = URI.decode_www_form_component(key.strip)
+              value = URI.decode_www_form_component(value.strip)
+            rescue ArgumentError
+              # `URI.decode_www_form_component` raises on malformed percent encoding
+              # (e.g. `%XX`, lone `%`). Treat as malformed rather than letting it
+              # propagate to `Propagation#extract`'s caller.
+              record_telemetry_metric('context_header_style.malformed', 1, {'header_style' => 'baggage'})
+              return {}
+            end
+
             if key.empty? || value.empty?
               # Record telemetry for malformed header
               record_telemetry_metric('context_header_style.malformed', 1, {'header_style' => 'baggage'})
@@ -157,6 +211,7 @@ module Datadog
 
             baggage[key] = value
           end
+
           baggage
         end
 
@@ -175,8 +230,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
@@ -188,7 +244,7 @@ module Datadog
         # Record telemetry metrics for baggage operations
         def record_telemetry_metric(metric_name, value, tags)
           telemetry = ::Datadog.send(:components).telemetry
-          telemetry.inc('instrumentation_telemetry_data.tracers', metric_name, value, tags: tags)
+          telemetry.inc(Tracing::Ext::TELEMETRY_METRICS_NAMESPACE, metric_name, value, tags: tags)
         end
       end
     end

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

@@ -131,17 +131,15 @@ module Datadog
 
           encoded_tags = DatadogTagsCodec.encode(tags)
 
-          return set_tags_propagation_error(reason: 'inject_max_size') if tags_too_large?(
-            encoded_tags.size,
-            scenario: 'inject'
-          )
+          return set_tags_propagation_error(reason: 'inject_max_size') if tags_too_large?(encoded_tags, scenario: 'inject')
 
           data[@tags_key] = encoded_tags
         rescue => e
           set_tags_propagation_error(reason: 'encoding_error')
           ::Datadog.logger.warn(
-            "Failed to inject x-datadog-tags: #{e.class.name} #{e.message} at #{Array(e.backtrace).first}"
+            "Failed to inject x-datadog-tags: #{e.class}: #{e.message} at #{Array(e.backtrace).first}"
           )
+          nil
         end
 
         # Import `x-datadog-tags` tags as trace distributed tags.
@@ -155,7 +153,7 @@ module Datadog
 
           return if !tags || tags.empty?
           return set_tags_propagation_error(reason: 'disabled') if tags_disabled?
-          return set_tags_propagation_error(reason: 'extract_max_size') if tags_too_large?(tags.size, scenario: 'extract')
+          return set_tags_propagation_error(reason: 'extract_max_size') if tags_too_large?(tags, scenario: 'extract')
 
           tags_hash = DatadogTagsCodec.decode(tags)
           # Only extract keys with the expected Datadog prefix
@@ -166,8 +164,9 @@ module Datadog
         rescue => e
           set_tags_propagation_error(reason: 'decoding_error')
           ::Datadog.logger.warn(
-            "Failed to extract x-datadog-tags: #{e.class.name} #{e.message} at #{Array(e.backtrace).first}"
+            "Failed to extract x-datadog-tags: #{e.class}: #{e.message} at #{Array(e.backtrace).first}"
           )
+          nil
         end
 
         def set_tags_propagation_error(reason:)
@@ -180,12 +179,15 @@ module Datadog
           ::Datadog.configuration.tracing.x_datadog_tags_max_length <= 0
         end
 
-        def tags_too_large?(size, scenario:)
-          return false if size <= ::Datadog.configuration.tracing.x_datadog_tags_max_length
+        def tags_too_large?(tags, scenario:)
+          size = tags.bytesize
+          max_length = ::Datadog.configuration.tracing.x_datadog_tags_max_length
+
+          return false if size <= max_length
 
           ::Datadog.logger.warn(
-            "Failed to #{scenario} x-datadog-tags: tags are too large for configured limit (size:#{size} >= " \
-              "limit:#{::Datadog.configuration.tracing.x_datadog_tags_max_length}). This limit can be configured " \
+            "Failed to #{scenario} x-datadog-tags: tags are too large for configured limit (bytesize:#{size} >= " \
+              "limit:#{max_length}). This limit can be configured " \
               'through the DD_TRACE_X_DATADOG_TAGS_MAX_LENGTH environment variable.'
           )
 

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

@@ -6,19 +6,6 @@ module Datadog
       # Encodes and decodes distributed 'x-datadog-tags' tags for transport
       # to and from external processes.
       module DatadogTagsCodec
-        # Backport `Regexp::match?` because it is measurably the most performant
-        # way to check if a string matches a regular expression.
-        module RefineRegexp
-          unless Regexp.method_defined?(:match?)
-            refine ::Regexp do
-              def match?(*args)
-                !match(*args).nil?
-              end
-            end
-          end
-        end
-        using RefineRegexp
-
         # ASCII characters 32-126, except `,`, `=`, and ` `. At least one character.
         VALID_KEY_CHARS = /\A(?:(?![,= ])[\u0020-\u007E])+\Z/.freeze
         # ASCII characters 32-126, except `,`. At least one character.
@@ -41,7 +28,7 @@ module Datadog
             "#{key}=#{value.strip}"
           end.join(',')
         rescue => e
-          raise EncodingError, "Error encoding tags `#{tags}`: `#{e}`"
+          raise EncodingError, "Error encoding tags `#{tags}`: `#{e.class}: #{e.message}`"
         end
 
         # Deserializes a `x-datadog-tags`-formatted String into a {Hash<String,String>}.

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

@@ -3,6 +3,7 @@
 require_relative '../configuration/ext'
 require_relative '../trace_digest'
 require_relative '../trace_operation'
+require_relative '../span_link'
 require_relative '../../core/telemetry/logger'
 require_relative 'baggage'
 
@@ -13,6 +14,8 @@ module Datadog
       class Propagation
         # @param propagation_styles [Hash<String,Object>]
         #  a map of propagation styles to their corresponding implementations
+        # @param propagation_behavior_extract [String]
+        #   the behavior to apply when extracting distributed trace data
         # @param propagation_style_inject [Array<String>]
         #   a list of styles to use when injecting distributed trace data
         # @param propagation_style_extract [Array<String>]
@@ -23,9 +26,11 @@ module Datadog
           propagation_styles:,
           propagation_style_inject:,
           propagation_style_extract:,
+          propagation_behavior_extract:,
           propagation_extract_first:
         )
           @propagation_styles = propagation_styles
+          @propagation_behavior_extract = propagation_behavior_extract
           @propagation_extract_first = propagation_extract_first
           @propagation_style_inject = propagation_style_inject.map { |style| propagation_styles[style] }
           @propagation_style_extract = propagation_style_extract.map { |style| propagation_styles[style] }
@@ -76,7 +81,7 @@ module Datadog
           rescue => e
             result = nil
             ::Datadog.logger.error(
-              "Error injecting distributed trace data. Cause: #{e} Location: #{Array(e.backtrace).first}"
+              "Error injecting distributed trace data. Cause: #{e.class}: #{e.message} Location: #{Array(e.backtrace).first}"
             )
             ::Datadog::Core::Telemetry::Logger.report(
               e,
@@ -97,12 +102,20 @@ module Datadog
           return unless data
           return if data.empty?
 
+          if @propagation_behavior_extract == Tracing::Configuration::Ext::Distributed::PROPAGATION_BEHAVIOR_EXTRACT_IGNORE
+            return nil
+          end
+
           extracted_trace_digest = nil
+          # Name of the propagation style that first produced a trace context.
+          # Used as the `context_headers` span-link attribute when restarting a trace.
+          extracted_style_name = nil
 
           @propagation_style_extract.each do |propagator|
             # First extraction?
             unless extracted_trace_digest
               extracted_trace_digest = propagator.extract(data)
+              extracted_style_name = @propagation_styles.key(propagator) if extracted_trace_digest
               next
             end
 
@@ -139,13 +152,32 @@ module Datadog
           rescue => e
             # TODO: Not to report Telemetry logs for now
             ::Datadog.logger.error(
-              "Error extracting distributed trace data. Cause: #{e} Location: #{Array(e.backtrace).first}"
+              "Error extracting distributed trace data. Cause: #{e.class}: #{e.message} Location: #{Array(e.backtrace).first}"
             )
           end
           # Handle baggage after all other styles if present
           extracted_trace_digest = propagate_baggage(data, extracted_trace_digest) if @baggage_propagator
 
-          extracted_trace_digest
+          # Default mode: Current digest continues the upstream trace.
+          return extracted_trace_digest if @propagation_behavior_extract != Tracing::Configuration::Ext::Distributed::PROPAGATION_BEHAVIOR_EXTRACT_RESTART || extracted_trace_digest&.trace_id.nil?
+
+          # Restart mode: start a new local digest and create a Span Link on
+          # the service-entry span with the original propagation data.
+          link = SpanLink.new(
+            extracted_trace_digest,
+            attributes: {
+              'reason' => 'propagation_behavior_extract',
+              'context_headers' => extracted_style_name,
+            }
+          )
+          baggage_tags = extracted_trace_digest.trace_distributed_tags&.select { |k, _| k.start_with?('baggage.') }
+          baggage_tags = nil if baggage_tags&.empty?
+          TraceDigest.new(
+            span_links: [link],
+            baggage: extracted_trace_digest.baggage,
+            trace_distributed_tags: baggage_tags,
+            span_remote: false,
+          )
         end
 
         private