cw-datadog 2.23.0.2

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

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Kafka
+        # Kafka integration constants
+        # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
+        module Ext
+          ENV_ENABLED = 'DD_TRACE_KAFKA_ENABLED'
+          # @!visibility private
+          ENV_ANALYTICS_ENABLED = 'DD_TRACE_KAFKA_ANALYTICS_ENABLED'
+          ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_KAFKA_ANALYTICS_SAMPLE_RATE'
+          SPAN_CONNECTION_REQUEST = 'kafka.connection.request'
+          SPAN_CONSUMER_HEARTBEAT = 'kafka.consumer.heartbeat'
+          SPAN_CONSUMER_JOIN_GROUP = 'kafka.consumer.join_group'
+          SPAN_CONSUMER_LEAVE_GROUP = 'kafka.consumer.leave_group'
+          SPAN_CONSUMER_SYNC_GROUP = 'kafka.consumer.sync_group'
+          SPAN_DELIVER_MESSAGES = 'kafka.producer.deliver_messages'
+          SPAN_PROCESS_BATCH = 'kafka.consumer.process_batch'
+          SPAN_PROCESS_MESSAGE = 'kafka.consumer.process_message'
+          SPAN_SEND_MESSAGES = 'kafka.producer.send_messages'
+          TAG_ATTEMPTS = 'kafka.attempts'
+          TAG_API = 'kafka.api'
+          TAG_CLIENT = 'kafka.client'
+          TAG_GROUP = 'kafka.group'
+          TAG_HIGHWATER_MARK_OFFSET = 'kafka.highwater_mark_offset'
+          TAG_MESSAGE_COUNT = 'kafka.message_count'
+          TAG_MESSAGE_KEY = 'kafka.message_key'
+          TAG_DELIVERED_MESSAGE_COUNT = 'kafka.delivered_message_count'
+          TAG_OFFSET = 'kafka.offset'
+          TAG_OFFSET_LAG = 'kafka.offset_lag'
+          TAG_PARTITION = 'kafka.partition'
+          TAG_REQUEST_SIZE = 'kafka.request_size'
+          TAG_RESPONSE_SIZE = 'kafka.response_size'
+          TAG_SENT_MESSAGE_COUNT = 'kafka.sent_message_count'
+          TAG_TOPIC = 'kafka.topic'
+          TAG_TOPIC_PARTITIONS = 'kafka.topic_partitions'
+          TAG_COMPONENT = 'kafka'
+          TAG_OPERATION_CONNECTION_REQUEST = 'connection.request'
+          TAG_OPERATION_CONSUMER_HEARTBEAT = 'consumer.heartbeat'
+          TAG_OPERATION_CONSUMER_JOIN_GROUP = 'consumer.join_group'
+          TAG_OPERATION_CONSUMER_LEAVE_GROUP = 'consumer.leave_group'
+          TAG_OPERATION_CONSUMER_SYNC_GROUP = 'consumer.sync_group'
+          TAG_OPERATION_DELIVER_MESSAGES = 'producer.deliver_messages'
+          TAG_OPERATION_PROCESS_BATCH = 'consumer.process_batch'
+          TAG_OPERATION_PROCESS_MESSAGE = 'consumer.process_message'
+          TAG_OPERATION_SEND_MESSAGES = 'producer.send_messages'
+          TAG_MESSAGING_SYSTEM = 'kafka'
+          TAG_KAFKA_BOOTSTRAP_SERVERS = 'messaging.kafka.bootstrap.servers'
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/kafka/instrumentation/consumer.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Kafka
+        module Instrumentation
+          # Instrumentation for Kafka::Consumer
+          module Consumer
+            def self.prepended(base)
+              base.prepend(InstanceMethods)
+            end
+
+            # Instance methods for consumer instrumentation
+            module InstanceMethods
+              def each_message(**kwargs, &block)
+                return super unless Datadog::DataStreams.enabled?
+
+                wrapped_block = proc do |message|
+                  Datadog.logger.debug { "Kafka each_message: DSM enabled for topic #{message.topic}" }
+
+                  begin
+                    headers = message.headers || {}
+                    Datadog::DataStreams.set_consume_checkpoint(
+                      type: 'kafka',
+                      source: message.topic,
+                      auto_instrumentation: true
+                    ) { |key| headers[key] }
+                  rescue => e
+                    Datadog.logger.debug("Error setting DSM checkpoint: #{e.class}: #{e}")
+                  end
+
+                  yield(message) if block
+                end
+
+                super(**kwargs, &wrapped_block)
+              end
+
+              def each_batch(**kwargs, &block)
+                return super unless Datadog::DataStreams.enabled?
+
+                wrapped_block = proc do |batch|
+                  Datadog.logger.debug { "Kafka each_batch: DSM enabled for topic #{batch.topic}" }
+
+                  begin
+                    Datadog::DataStreams.set_consume_checkpoint(
+                      type: 'kafka',
+                      source: batch.topic,
+                      auto_instrumentation: true
+                    )
+                  rescue => e
+                    Datadog.logger.debug("Error setting DSM checkpoint: #{e.class}: #{e}")
+                  end
+
+                  yield(batch) if block
+                end
+
+                super(**kwargs, &wrapped_block)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/kafka/instrumentation/producer.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Kafka
+        module Instrumentation
+          # Instrumentation for Kafka::Producer
+          module Producer
+            def self.prepended(base)
+              base.prepend(InstanceMethods)
+            end
+
+            module InstanceMethods
+              def deliver_messages(**kwargs)
+                if Datadog::DataStreams.enabled?
+                  begin
+                    pending_messages = instance_variable_get(:@pending_message_queue)
+
+                    if pending_messages && !pending_messages.empty?
+                      pending_messages.each do |message|
+                        message.headers ||= {}
+                        Datadog::DataStreams.set_produce_checkpoint(
+                          type: 'kafka',
+                          destination: message.topic,
+                          auto_instrumentation: true
+                        ) do |key, value|
+                          message.headers[key] = value
+                        end
+                      end
+                    end
+                  rescue => e
+                    Datadog.logger.debug("Error setting DSM checkpoint: #{e.class}: #{e}")
+                  end
+                end
+
+                super
+              end
+
+              def send_messages(messages, **kwargs)
+                if Datadog::DataStreams.enabled?
+                  begin
+                    messages.each do |message|
+                      message[:headers] ||= {}
+                      Datadog::DataStreams.set_produce_checkpoint(
+                        type: 'kafka',
+                        destination: message[:topic],
+                        auto_instrumentation: true
+                      ) do |key, value|
+                        message[:headers][key] = value
+                      end
+                    end
+                  rescue => e
+                    Datadog.logger.debug("Error setting DSM checkpoint: #{e.class}: #{e}")
+                  end
+                end
+
+                super
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require_relative '../integration'
+require_relative 'configuration/settings'
+require_relative 'patcher'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Kafka
+        # Description of Kafka integration
+        class Integration
+          include Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new('0.7.10')
+
+          # @public_api Changing the integration name or integration options can cause breaking changes
+          register_as :kafka, auto_patch: false
+          def self.gem_name
+            'ruby-kafka'
+          end
+
+          def self.version
+            Gem.loaded_specs['ruby-kafka']&.version
+          end
+
+          def self.loaded?
+            !defined?(::Kafka).nil? \
+              && !defined?(::ActiveSupport::Notifications).nil?
+          end
+
+          def self.compatible?
+            super && version >= MINIMUM_VERSION
+          end
+
+          def new_configuration
+            Configuration::Settings.new
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative '../patcher'
+require_relative 'ext'
+require_relative 'events'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Kafka
+        # Patcher enables patching of 'kafka' module.
+        module Patcher
+          include Contrib::Patcher
+
+          module_function
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            # Subscribe to Kafka events
+            Events.subscribe!
+
+            # Apply monkey patches for additional instrumentation (e.g., DSM)
+            patch_producer if defined?(::Kafka::Producer)
+            patch_consumer if defined?(::Kafka::Consumer)
+          end
+
+          def patch_producer
+            require_relative 'instrumentation/producer'
+            ::Kafka::Producer.prepend(Instrumentation::Producer)
+          end
+
+          def patch_consumer
+            require_relative 'instrumentation/consumer'
+            ::Kafka::Consumer.prepend(Instrumentation::Consumer)
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require_relative '../../configuration/settings'
+require_relative '../ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        module Configuration
+          # @public_api
+          class Settings < Contrib::Configuration::Settings
+            option :enabled do |o|
+              o.type :bool
+              o.env Ext::ENV_ENABLED
+              o.default true
+            end
+
+            option :service_name
+
+            option :distributed_tracing, default: false, type: :bool
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require_relative '../../../distributed/fetcher'
+require_relative '../../../distributed/propagation'
+require_relative '../../../distributed/b3_multi'
+require_relative '../../../distributed/b3_single'
+require_relative '../../../distributed/datadog'
+require_relative '../../../distributed/none'
+require_relative '../../../distributed/trace_context'
+require_relative '../../../configuration/ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        module Distributed
+          # Extracts and injects propagation through Kafka message headers.
+          class Propagation < Tracing::Distributed::Propagation
+            def initialize(
+              propagation_style_inject:,
+              propagation_style_extract:,
+              propagation_extract_first:
+            )
+              super(
+                propagation_styles: {
+                  Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_B3_MULTI_HEADER =>
+                    Tracing::Distributed::B3Multi.new(fetcher: Tracing::Distributed::Fetcher),
+                  Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_B3_SINGLE_HEADER =>
+                    Tracing::Distributed::B3Single.new(fetcher: Tracing::Distributed::Fetcher),
+                  Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_DATADOG =>
+                    Tracing::Distributed::Datadog.new(fetcher: Tracing::Distributed::Fetcher),
+                  Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_TRACE_CONTEXT =>
+                    Tracing::Distributed::TraceContext.new(fetcher: Tracing::Distributed::Fetcher),
+                  Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_BAGGAGE =>
+                    Tracing::Distributed::Baggage.new(fetcher: Tracing::Distributed::Fetcher),
+                  Tracing::Configuration::Ext::Distributed::PROPAGATION_STYLE_NONE => Tracing::Distributed::None.new
+                },
+                propagation_style_inject: propagation_style_inject,
+                propagation_style_extract: propagation_style_extract,
+                propagation_extract_first: propagation_extract_first
+              )
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        module Ext
+          ENV_ENABLED = 'DD_TRACE_KARAFKA_ENABLED'
+
+          SPAN_MESSAGE_CONSUME = 'karafka.consume'
+          SPAN_WORKER_PROCESS = 'worker.process'
+
+          TAG_CONSUMER = 'kafka.consumer'
+          TAG_TOPIC = 'kafka.topic'
+          TAG_PARTITION = 'kafka.partition'
+          TAG_OFFSET = 'kafka.offset'
+          TAG_OFFSET_LAG = 'kafka.offset_lag'
+          TAG_MESSAGE_COUNT = 'kafka.message_count'
+          TAG_MESSAGE_KEY = 'kafka.message_key'
+          TAG_SYSTEM = 'kafka'
+
+          TAG_OPERATION_PROCESS_BATCH = 'consumer.process_batch'
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/karafka/integration.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require_relative '../integration'
+require_relative 'configuration/settings'
+require_relative 'patcher'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        # Description of Kafka integration
+        class Integration
+          include Contrib::Integration
+
+          # Minimum version of the Karafka library that we support
+          # https://karafka.io/docs/Versions-Lifecycle-and-EOL/#versioning-strategy
+          MINIMUM_VERSION = Gem::Version.new('2.3.0')
+
+          # @public_api Changing the integration name or integration options can cause breaking changes
+          register_as :karafka, auto_patch: false
+
+          def self.version
+            Gem.loaded_specs['karafka']&.version
+          end
+
+          def self.loaded?
+            !defined?(::Karafka).nil?
+          end
+
+          def self.compatible?
+            super && version >= MINIMUM_VERSION
+          end
+
+          def new_configuration
+            Configuration::Settings.new
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/karafka/monitor.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+require_relative 'ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        # Custom monitor for Karafka.
+        # Creating a custom monitor, instead of subscribing to an event
+        # (e.g. `Karafka.monitor.subscribe 'worker.processed'`),
+        # is required because event subscriptions cannot wrap the event execution (`yield`).
+        module Monitor
+          TRACEABLE_EVENTS = %w[
+            worker.processed
+          ].freeze
+
+          def instrument(event_id, payload = {}, &block)
+            return super unless TRACEABLE_EVENTS.include?(event_id)
+
+            Datadog::Tracing.trace(Ext::SPAN_WORKER_PROCESS) do |span|
+              job = payload[:job]
+              job_type = fetch_job_type(job.class)
+              consumer = job.executor.topic.consumer
+
+              action = case job_type
+              when 'Periodic', 'PeriodicNonBlocking'
+                'tick'
+              when 'Shutdown'
+                'shutdown'
+              when 'Revoked', 'RevokedNonBlocking'
+                'revoked'
+              when 'Idle'
+                'idle'
+              when 'Eofed', 'EofedNonBlocking'
+                'eofed'
+              else
+                'consume'
+              end
+
+              span.resource = "#{consumer}##{action}"
+
+              if action == 'consume'
+                span.set_tag(Ext::TAG_MESSAGE_COUNT, job.messages.count)
+                span.set_tag(Ext::TAG_PARTITION, job.executor.partition)
+                span.set_tag(Ext::TAG_OFFSET, job.messages.first.metadata.offset)
+                span.set_tag(Ext::TAG_CONSUMER, consumer)
+                span.set_tag(Contrib::Ext::Messaging::TAG_DESTINATION, job.executor.topic.name)
+                span.set_tag(Contrib::Ext::Messaging::TAG_SYSTEM, Ext::TAG_SYSTEM)
+
+                # DSM: Track consumer offset stats for batch processing
+                if Datadog.configuration.data_streams.enabled
+                  job.messages.each do |message|
+                    Datadog::DataStreams.track_kafka_consume(
+                      job.executor.topic.name,
+                      job.executor.partition,
+                      message.metadata.offset
+                    )
+                  end
+                end
+              end
+
+              super
+            end
+          end
+
+          private
+
+          def fetch_job_type(job_class)
+            @job_types_cache ||= {}
+            @job_types_cache[job_class] ||= job_class.to_s.split('::').last
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require_relative '../patcher'
+require_relative 'ext'
+require_relative 'distributed/propagation'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        # Patch to add tracing to Karafka::Messages::Messages
+        module MessagesPatch
+          def configuration
+            Datadog.configuration.tracing[:karafka]
+          end
+
+          def propagation
+            @propagation ||= Contrib::Karafka::Distributed::Propagation.new
+          end
+
+          # `each` is the most popular access point to Karafka messages,
+          # but not the only one
+          #  Other access patterns do not have a straightforward tracing avenue
+          # (e.g. `my_batch_operation messages.payloads`)
+          # @see https://github.com/karafka/karafka/blob/b06d1f7c17818e1605f80c2bb573454a33376b40/README.md?plain=1#L29-L35
+          def each(&block)
+            @messages_array.each do |message|
+              if configuration[:distributed_tracing]
+                headers = if message.metadata.respond_to?(:raw_headers)
+                  message.metadata.raw_headers
+                else
+                  message.metadata.headers
+                end
+                trace_digest = Karafka.extract(headers)
+                Datadog::Tracing.continue_trace!(trace_digest) if trace_digest
+              end
+
+              if Datadog::DataStreams.enabled?
+                begin
+                  headers = if message.metadata.respond_to?(:raw_headers)
+                    message.metadata.raw_headers
+                  else
+                    message.metadata.headers
+                  end
+
+                  Datadog::DataStreams.set_consume_checkpoint(
+                    type: 'kafka',
+                    source: message.topic,
+                    auto_instrumentation: true
+                  ) { |key| headers[key] }
+                rescue => e
+                  Datadog.logger.debug("Error setting DSM checkpoint: #{e.class}: #{e}")
+                end
+              end
+
+              Tracing.trace(Ext::SPAN_MESSAGE_CONSUME) do |span|
+                span.set_tag(Ext::TAG_OFFSET, message.metadata.offset)
+                span.set_tag(Contrib::Ext::Messaging::TAG_DESTINATION, message.topic)
+                span.set_tag(Contrib::Ext::Messaging::TAG_SYSTEM, Ext::TAG_SYSTEM)
+
+                span.resource = message.topic
+
+                yield message
+              end
+            end
+          end
+        end
+
+        # Patcher enables patching of 'karafka' module.
+        module Patcher
+          include Contrib::Patcher
+
+          module_function
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            require_relative 'monitor'
+
+            ::Karafka::Instrumentation::Monitor.prepend(Monitor)
+            ::Karafka::Messages::Messages.prepend(MessagesPatch)
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative 'component'
+require_relative 'karafka/integration'
+require_relative 'karafka/distributed/propagation'
+
+module Datadog
+  module Tracing
+    module Contrib
+      # `Karafka` integration public API
+      module Karafka
+        def self.inject(digest, data)
+          raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+          @propagation.inject!(digest, data)
+        end
+
+        def self.extract(data)
+          raise 'Please invoke Datadog.configure at least once before calling this method' unless @propagation
+
+          @propagation.extract(data)
+        end
+
+        Contrib::Component.register('karafka') do |config|
+          tracing = config.tracing
+          tracing.propagation_style
+
+          @propagation = Karafka::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/lograge/configuration/settings.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require_relative '../../configuration/settings'
+require_relative '../ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Lograge
+        module Configuration
+          # Custom settings for the Lograge integration
+          # @public_api
+          class Settings < Contrib::Configuration::Settings
+            option :enabled do |o|
+              o.type :bool
+              o.env Ext::ENV_ENABLED
+              o.default true
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Lograge
+        # Lograge integration constants
+        # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
+        module Ext
+          ENV_ENABLED = 'DD_TRACE_LOGRAGE_ENABLED'
+        end
+      end
+    end
+  end
+end