datadog 2.21.0 → 2.38.0

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

@@ -36,6 +36,7 @@ module Datadog
           @propagation = GRPC::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

data/lib/datadog/tracing/contrib/hanami/action_tracer.rb

@@ -22,6 +22,7 @@ module Datadog
               service: configuration[:service_name],
               type: Tracing::Metadata::Ext::HTTP::TYPE_INBOUND
             ) do |span_op, trace_op|
+              span_op.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
               span_op.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span_op.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_ACTION)
 

data/lib/datadog/tracing/contrib/hanami/renderer_policy_tracing.rb

@@ -24,6 +24,7 @@ module Datadog
               resource: action_klass.to_s,
               type: Tracing::Metadata::Ext::HTTP::TYPE_INBOUND
             ) do |span_op, _trace_op|
+              span_op.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
               span_op.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span_op.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_RENDER)
 

data/lib/datadog/tracing/contrib/hanami/router_tracing.rb

@@ -17,6 +17,7 @@ module Datadog
               service: configuration[:service_name],
               type: Tracing::Metadata::Ext::HTTP::TYPE_INBOUND
             ) do |span_op, trace_op|
+              span_op.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
               span_op.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span_op.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_ROUTING)
 

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

@@ -32,8 +32,11 @@ module Datadog
               o.default 1.0
             end
 
-            option :distributed_tracing, default: true, type: :bool
-
+            option :distributed_tracing do |o|
+              o.type :bool
+              o.env Ext::ENV_DISTRIBUTED_TRACING
+              o.default true
+            end
             option :service_name do |o|
               o.type :string, nilable: true
               o.default do
@@ -46,9 +49,17 @@ module Datadog
 
             option :error_status_codes do |o|
               o.env Ext::ENV_ERROR_STATUS_CODES
-              o.default 400...600
-              o.setter do |v|
-                Tracing::Contrib::StatusRangeMatcher.new(v) if v
+              o.setter do |value|
+                if value.nil?
+                  # Fallback to global config, which is defaulted to client (400..499) + server (500..599)
+                  # DEV-3.0: `http` is a client library, this should fall back to `http_error_statuses.client` only.
+                  # We cannot change it without causing a breaking change.
+                  client_global_error_statuses = Datadog.configuration.tracing.http_error_statuses.client
+                  server_global_error_statuses = Datadog.configuration.tracing.http_error_statuses.server
+                  client_global_error_statuses + server_global_error_statuses
+                else
+                  Tracing::Contrib::StatusRangeMatcher.new(value)
+                end
               end
               o.env_parser do |v|
                 Tracing::Contrib::StatusRangeEnvParser.call(v) if v

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

@@ -19,6 +19,7 @@ module Datadog
             def initialize(
               propagation_style_inject:,
               propagation_style_extract:,
+              propagation_behavior_extract:,
               propagation_extract_first:
             )
               super(
@@ -37,6 +38,7 @@ module Datadog
                 },
                 propagation_style_inject: propagation_style_inject,
                 propagation_style_extract: propagation_style_extract,
+                propagation_behavior_extract: propagation_behavior_extract,
                 propagation_extract_first: propagation_extract_first
               )
             end

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

@@ -8,6 +8,7 @@ module Datadog
         # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
         module Ext
           ENV_ENABLED = 'DD_TRACE_HTTP_ENABLED'
+          ENV_DISTRIBUTED_TRACING = 'DD_TRACE_HTTP_DISTRIBUTED_TRACING'
           ENV_SERVICE_NAME = 'DD_TRACE_NET_HTTP_SERVICE_NAME'
           ENV_PEER_SERVICE = 'DD_TRACE_NET_HTTP_PEER_SERVICE'
           # @!visibility private

data/lib/datadog/tracing/contrib/http/instrumentation.rb

@@ -32,12 +32,13 @@ module Datadog
 
               Tracing.trace(Ext::SPAN_REQUEST) do |span, trace|
                 span.service = service_name(host, request_options, client_config)
+                span.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
                 span.type = Tracing::Metadata::Ext::HTTP::TYPE_OUTBOUND
                 span.resource = req.method
 
                 if Tracing::Distributed::PropagationPolicy.enabled?(
                   pin_config: client_config,
-                  global_config: Datadog.configuration.tracing[:http],
+                  global_config: request_options,
                   trace: trace
                 )
                   Contrib::HTTP.inject(trace, req)
@@ -63,11 +64,6 @@ module Datadog
                 )
               end
 
-              # Tag original global service name if not used
-              if span.service != Datadog.configuration.service
-                span.set_tag(Tracing::Contrib::Ext::Metadata::TAG_BASE_SERVICE, Datadog.configuration.service)
-              end
-
               span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_CLIENT)
 
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
@@ -93,7 +89,7 @@ module Datadog
 
               Contrib::SpanAttributeSchema.set_peer_service!(span, Ext::PEER_SERVICE_SOURCES)
             rescue => e
-              Datadog.logger.error("error preparing span from http request: #{e}")
+              Datadog.logger.error("error preparing span from http request: #{e.class}: #{e.message}")
               Datadog::Core::Telemetry::Logger.report(e)
             end
 
@@ -108,7 +104,7 @@ module Datadog
                 Datadog.configuration.tracing.header_tags.response_tags(response)
               )
             rescue => e
-              Datadog.logger.error("error preparing span from http response: #{e}")
+              Datadog.logger.error("error preparing span from http response: #{e.class}: #{e.message}")
               Datadog::Core::Telemetry::Logger.report(e)
             end
 

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

@@ -18,8 +18,6 @@ module Datadog
         class Integration
           include Contrib::Integration
 
-          MINIMUM_VERSION = Datadog::VERSION::MINIMUM_RUBY_VERSION
-
           # @public_api Changing the integration name or integration options can cause breaking changes
           register_as :http, auto_patch: true
           def self.gem_name

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

@@ -36,6 +36,7 @@ module Datadog
           @propagation = HTTP::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

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

@@ -32,8 +32,11 @@ module Datadog
               o.default 1.0
             end
 
-            option :distributed_tracing, default: true, type: :bool
-
+            option :distributed_tracing do |o|
+              o.type :bool
+              o.env Ext::ENV_DISTRIBUTED_TRACING
+              o.default true
+            end
             option :service_name do |o|
               o.default do
                 Contrib::SpanAttributeSchema.fetch_service_name(
@@ -45,9 +48,17 @@ module Datadog
 
             option :error_status_codes do |o|
               o.env Ext::ENV_ERROR_STATUS_CODES
-              o.default 400...600
-              o.setter do |v|
-                Tracing::Contrib::StatusRangeMatcher.new(v) if v
+              o.setter do |value|
+                if value.nil?
+                  # Fallback to global config, which is defaulted to client (400..499) + server (500..599)
+                  # DEV-3.0: `httpclient` is a client library, this should fall back to `http_error_statuses.client` only.
+                  # We cannot change it without causing a breaking change.
+                  client_global_error_statuses = Datadog.configuration.tracing.http_error_statuses.client
+                  server_global_error_statuses = Datadog.configuration.tracing.http_error_statuses.server
+                  client_global_error_statuses + server_global_error_statuses
+                else
+                  Tracing::Contrib::StatusRangeMatcher.new(value)
+                end
               end
               o.env_parser do |v|
                 Tracing::Contrib::StatusRangeEnvParser.call(v) if v

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

@@ -8,6 +8,7 @@ module Datadog
         # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
         module Ext
           ENV_ENABLED = 'DD_TRACE_HTTPCLIENT_ENABLED'
+          ENV_DISTRIBUTED_TRACING = 'DD_TRACE_HTTPCLIENT_DISTRIBUTED_TRACING'
           ENV_SERVICE_NAME = 'DD_TRACE_HTTPCLIENT_SERVICE_NAME'
           ENV_PEER_SERVICE = 'DD_TRACE_HTTPCLIENT_PEER_SERVICE'
 

data/lib/datadog/tracing/contrib/httpclient/instrumentation.rb

@@ -28,6 +28,7 @@ module Datadog
               Tracing.trace(Ext::SPAN_REQUEST) do |span, trace|
                 begin
                   span.service = service_name(host, request_options, client_config)
+                  span.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
                   span.type = Tracing::Metadata::Ext::HTTP::TYPE_OUTBOUND
 
                   if Tracing::Distributed::PropagationPolicy.enabled?(
@@ -41,7 +42,9 @@ module Datadog
                   # Add additional request specific tags to the span.
                   annotate_span_with_request!(span, req, request_options)
                 rescue => e
-                  Datadog.logger.error("error preparing span for httpclient request: #{e}, Source: #{e.backtrace}")
+                  Datadog.logger.error(
+                    "error preparing span for httpclient request: #{e.class}: #{e.message}, Source: #{e.backtrace}"
+                  )
                   Datadog::Core::Telemetry::Logger.report(e)
                 ensure
                   res = super
@@ -66,11 +69,6 @@ module Datadog
                 )
               end
 
-              # Tag original global service name if not used
-              if span.service != Datadog.configuration.service
-                span.set_tag(Tracing::Contrib::Ext::Metadata::TAG_BASE_SERVICE, Datadog.configuration.service)
-              end
-
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_REQUEST)
 
@@ -107,7 +105,9 @@ module Datadog
                 Datadog.configuration.tracing.header_tags.response_tags(response.header)
               )
             rescue => e
-              Datadog.logger.error("error preparing span from httpclient response: #{e}, Source: #{e.backtrace}")
+              Datadog.logger.error(
+                "error preparing span from httpclient response: #{e.class}: #{e.message}, Source: #{e.backtrace}"
+              )
               Datadog::Core::Telemetry::Logger.report(e)
             end
 

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

@@ -32,8 +32,11 @@ module Datadog
               o.default 1.0
             end
 
-            option :distributed_tracing, default: true, type: :bool
-
+            option :distributed_tracing do |o|
+              o.type :bool
+              o.env Ext::ENV_DISTRIBUTED_TRACING
+              o.default true
+            end
             option :service_name do |o|
               o.default do
                 Contrib::SpanAttributeSchema.fetch_service_name(
@@ -45,9 +48,17 @@ module Datadog
 
             option :error_status_codes do |o|
               o.env Ext::ENV_ERROR_STATUS_CODES
-              o.default 400...600
-              o.setter do |v|
-                Tracing::Contrib::StatusRangeMatcher.new(v) if v
+              o.setter do |value|
+                if value.nil?
+                  # Fallback to global config, which is defaulted to client (400..499) + server (500..599)
+                  # DEV-3.0: `httprb` is a client library, this should fall back to `http_error_statuses.client` only.
+                  # We cannot change it without causing a breaking change.
+                  client_global_error_statuses = Datadog.configuration.tracing.http_error_statuses.client
+                  server_global_error_statuses = Datadog.configuration.tracing.http_error_statuses.server
+                  client_global_error_statuses + server_global_error_statuses
+                else
+                  Tracing::Contrib::StatusRangeMatcher.new(value)
+                end
               end
               o.env_parser do |v|
                 Tracing::Contrib::StatusRangeEnvParser.call(v) if v

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

@@ -8,6 +8,7 @@ module Datadog
         # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
         module Ext
           ENV_ENABLED = 'DD_TRACE_HTTPRB_ENABLED'
+          ENV_DISTRIBUTED_TRACING = 'DD_TRACE_HTTPRB_DISTRIBUTED_TRACING'
           ENV_SERVICE_NAME = 'DD_TRACE_HTTPRB_SERVICE_NAME'
           ENV_PEER_SERVICE = 'DD_TRACE_HTTPRB_PEER_SERVICE'
           # @!visibility private

data/lib/datadog/tracing/contrib/httprb/instrumentation.rb

@@ -28,6 +28,7 @@ module Datadog
               Tracing.trace(Ext::SPAN_REQUEST) do |span, trace|
                 begin
                   span.service = service_name(host, request_options, client_config)
+                  span.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
                   span.type = Tracing::Metadata::Ext::HTTP::TYPE_OUTBOUND
 
                   if Tracing::Distributed::PropagationPolicy.enabled?(
@@ -35,13 +36,13 @@ module Datadog
                     global_config: Datadog.configuration.tracing[:httprb],
                     trace: trace
                   )
-                    Contrib::HTTP.inject(trace, req)
+                    Contrib::HTTP.inject(trace, req.headers)
                   end
 
                   # Add additional request specific tags to the span.
                   annotate_span_with_request!(span, req, request_options)
                 rescue => e
-                  logger.error("error preparing span for http.rb request: #{e}, Source: #{e.backtrace}")
+                  logger.error("error preparing span for http.rb request: #{e.class}: #{e.message}, Source: #{e.backtrace}")
                   Datadog::Core::Telemetry::Logger.report(e)
                 ensure
                   res = super(req, options)
@@ -66,11 +67,6 @@ module Datadog
                 )
               end
 
-              # Tag original global service name if not used
-              if span.service != Datadog.configuration.service
-                span.set_tag(Tracing::Contrib::Ext::Metadata::TAG_BASE_SERVICE, Datadog.configuration.service)
-              end
-
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_REQUEST)
 
@@ -117,7 +113,7 @@ module Datadog
                 Datadog.configuration.tracing.header_tags.response_tags(response.headers)
               )
             rescue => e
-              logger.error("error preparing span from http.rb response: #{e}, Source: #{e.backtrace}")
+              logger.error("error preparing span from http.rb response: #{e.class}: #{e.message}, Source: #{e.backtrace}")
               Datadog::Core::Telemetry::Logger.report(e)
             end
 

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

@@ -30,6 +30,7 @@ module Datadog
             end
 
             def on_start(span, _event, _id, payload)
+              span.set_tag(Tracing::Metadata::Ext::TAG_SVC_SRC, Ext::TAG_COMPONENT)
               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/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.message}")
+                  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.message}")
+                  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.message}")
+                  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.message}")
+                  end
+                end
+
+                super
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -21,6 +21,20 @@ module Datadog
           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

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

@@ -18,7 +18,11 @@ module Datadog
 
             option :service_name
 
-            option :distributed_tracing, default: false, type: :bool
+            option :distributed_tracing do |o|
+              o.type :bool
+              o.env Ext::ENV_DISTRIBUTED_TRACING
+              o.default false
+            end
           end
         end
       end

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

@@ -19,6 +19,7 @@ module Datadog
             def initialize(
               propagation_style_inject:,
               propagation_style_extract:,
+              propagation_behavior_extract:,
               propagation_extract_first:
             )
               super(
@@ -37,6 +38,7 @@ module Datadog
                 },
                 propagation_style_inject: propagation_style_inject,
                 propagation_style_extract: propagation_style_extract,
+                propagation_behavior_extract: propagation_behavior_extract,
                 propagation_extract_first: propagation_extract_first
               )
             end

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

@@ -6,6 +6,7 @@ module Datadog
       module Karafka
         module Ext
           ENV_ENABLED = 'DD_TRACE_KARAFKA_ENABLED'
+          ENV_DISTRIBUTED_TRACING = 'DD_TRACE_KARAFKA_DISTRIBUTED_TRACING'
 
           SPAN_MESSAGE_CONSUME = 'karafka.consume'
           SPAN_WORKER_PROCESS = 'worker.process'

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

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Karafka
+        # Karafka framework code, used to essentially:
+        # - handle configuration entries which are specific to Datadog tracing
+        # - instrument parts of the framework when needed
+        module Framework
+          def self.setup
+            Datadog.configure do |datadog_config|
+              karafka_config = datadog_config.tracing[:karafka]
+              activate_waterdrop!(datadog_config, karafka_config)
+            end
+          end
+
+          # Apply relevant configuration from Karafka to WaterDrop
+          def self.activate_waterdrop!(datadog_config, karafka_config)
+            datadog_config.tracing.instrument(
+              :waterdrop,
+              service_name: karafka_config[:service_name],
+              distributed_tracing: karafka_config[:distributed_tracing],
+            )
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -47,6 +47,17 @@ module Datadog
                 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

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

@@ -25,17 +25,34 @@ module Datadog
           # @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]
+              trace_digest = 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
+                Karafka.extract(headers)
               end
 
-              Tracing.trace(Ext::SPAN_MESSAGE_CONSUME) do |span|
+              Tracing.trace(Ext::SPAN_MESSAGE_CONSUME, continue_from: trace_digest) do |span, trace|
+                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.message}")
+                  end
+                end
+
                 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)
@@ -52,6 +69,8 @@ module Datadog
         module Patcher
           include Contrib::Patcher
 
+          ACTIVATE_FRAMEWORK_ONLY_ONCE = Core::Utils::OnlyOnce.new
+
           module_function
 
           def target_version
@@ -60,9 +79,21 @@ module Datadog
 
           def patch
             require_relative 'monitor'
+            require_relative 'framework'
+            require_relative '../waterdrop'
 
             ::Karafka::Instrumentation::Monitor.prepend(Monitor)
             ::Karafka::Messages::Messages.prepend(MessagesPatch)
+
+            if Contrib::WaterDrop::Integration.compatible?
+              ::Karafka.monitor.subscribe('app.initialized') do |event|
+                ACTIVATE_FRAMEWORK_ONLY_ONCE.run do
+                  Contrib::Karafka::Framework.setup
+                end
+
+                Contrib::WaterDrop::Patcher.add_middleware(::Karafka.producer)
+              end
+            end
           end
         end
       end