Support different karafka/waterdrop configurations per topic

Author: Drowze

docs/GettingStarted.md

@@ -1154,6 +1154,10 @@ require 'datadog'
 
 Datadog.configure do |c|
   c.tracing.instrument :karafka, **options
+
+  # different options for different kafka topics
+  c.tracing.instrument :karafka, describes: "some-topic", distributed_tracing: false
+  c.tracing.instrument :karafka, describes: /^topic-regex-.*/, enabled: false
 end
 
 ```
@@ -1176,6 +1180,10 @@ require 'datadog'
 
 Datadog.configure do |c|
   c.tracing.instrument :waterdrop, **options
+
+  # different options for different kafka topics
+  c.tracing.instrument :waterdrop, describes: "some-topic", distributed_tracing: false
+  c.tracing.instrument :waterdrop, describes: /^topic-regex-.*/, enabled: false
 end
 
 ```

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

@@ -9,18 +9,23 @@ module Karafka
         # - instrument parts of the framework when needed
         module Framework
           def self.setup
+            karafka_configurations = Datadog.configuration.tracing.fetch_integration(:karafka).configurations
+
             Datadog.configure do |datadog_config|
-              karafka_config = datadog_config.tracing[:karafka]
-              activate_waterdrop!(datadog_config, karafka_config)
+              karafka_configurations.each do |config_name, karafka_config|
+                activate_waterdrop!(datadog_config, config_name, karafka_config)
+              end
             end
           end
 
           # Apply relevant configuration from Karafka to WaterDrop
-          def self.activate_waterdrop!(datadog_config, karafka_config)
+          def self.activate_waterdrop!(datadog_config, config_name, karafka_config)
             datadog_config.tracing.instrument(
               :waterdrop,
+              enabled: karafka_config[:enabled],
               service_name: karafka_config[:service_name],
               distributed_tracing: karafka_config[:distributed_tracing],
+              describes: config_name,
             )
           end
         end

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

@@ -38,6 +38,10 @@ def new_configuration
           def patcher
             Patcher
           end
+
+          def resolver
+            @resolver ||= Contrib::Configuration::Resolvers::PatternResolver.new
+          end
         end
       end
     end

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

@@ -10,21 +10,14 @@ 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|
+              configuration = datadog_configuration(message.topic)
               trace_digest = if configuration[:distributed_tracing]
                 headers = if message.metadata.respond_to?(:raw_headers)
                   message.metadata.raw_headers
@@ -63,6 +56,12 @@ def each(&block)
               end
             end
           end
+
+          private
+
+          def datadog_configuration(topic)
+            Datadog.configuration.tracing[:karafka, topic]
+          end
         end
 
         module AppPatch

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

@@ -36,6 +36,10 @@ def new_configuration
           def patcher
             Patcher
           end
+
+          def resolver
+            @resolver ||= Contrib::Configuration::Resolvers::PatternResolver.new
+          end
         end
       end
     end

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

@@ -13,7 +13,7 @@ def call(message)
               trace_op = Datadog::Tracing.active_trace
 
               if trace_op && Datadog::Tracing::Distributed::PropagationPolicy.enabled?(
-                global_config: configuration,
+                global_config: datadog_configuration(message[:topic]),
                 trace: trace_op
               )
                 WaterDrop.inject(trace_op.to_digest, message[:headers] ||= {})
@@ -35,8 +35,8 @@ def call(message)
 
             private
 
-            def configuration
-              Datadog.configuration.tracing[:waterdrop]
+            def datadog_configuration(topic)
+              Datadog.configuration.tracing[:waterdrop, topic]
             end
           end
         end

spec/datadog/tracing/contrib/karafka/patcher_spec.rb

@@ -17,6 +17,7 @@
   before do
     Datadog.configure do |c|
       c.tracing.instrument :karafka, configuration_options
+      c.tracing.instrument :karafka, describes: /special_/, distributed_tracing: false
     end
   end
 
@@ -31,16 +32,12 @@
     let(:span_name) { Datadog::Tracing::Contrib::Karafka::Ext::SPAN_MESSAGE_CONSUME }
 
     it 'is expected to send a span' do
-      metadata = ::Karafka::Messages::Metadata.new
-      metadata['offset'] = 412
+      metadata = ::Karafka::Messages::Metadata.new(offset: 412, timestamp: Time.now, topic: 'topic_a')
       raw_payload = rand.to_s
 
       message = ::Karafka::Messages::Message.new(raw_payload, metadata)
-      allow(message).to receive(:timestamp).and_return(Time.now)
-      allow(message).to receive(:topic).and_return('topic_a')
-
-      topic = ::Karafka::Routing::Topic.new('topic_a', double(id: 0))
 
+      topic = ::Karafka::Routing::Topic.new(message.topic, double(id: 0))
       messages = ::Karafka::Messages::Builders::Messages.call([message], topic, 0, Time.now)
 
       expect(messages).to all(be_a(::Karafka::Messages::Message))
@@ -55,20 +52,21 @@
     end
 
     context 'when the message has tracing headers' do
+      let(:topic_name) { 'topic_a' }
       let(:message) do
         headers = {}
         Datadog::Tracing.trace('producer') do |span, trace|
           Datadog::Tracing::Contrib::Karafka.inject(trace.to_digest, headers)
         end
-        metadata = ::Karafka::Messages::Metadata.new
-        metadata['offset'] = 412
-        metadata[headers_accessor] = headers
+        metadata = ::Karafka::Messages::Metadata.new(
+          :offset => 412,
+          headers_accessor => headers,
+          :topic => topic_name,
+          :timestamp => Time.now
+        )
         raw_payload = rand.to_s
 
-        message = ::Karafka::Messages::Message.new(raw_payload, metadata)
-        allow(message).to receive(:timestamp).and_return(Time.now)
-        allow(message).to receive(:topic).and_return('topic_a')
-        message
+        ::Karafka::Messages::Message.new(raw_payload, metadata)
       end
       let(:headers_accessor) do
         ::Karafka::Messages::Metadata.members.include?(:raw_headers) ? 'raw_headers' : 'headers'
@@ -85,7 +83,7 @@
             consumer_span = Datadog::Tracing.active_span
             consumer_trace = Datadog::Tracing.active_trace
 
-            topic = ::Karafka::Routing::Topic.new('topic_a', double(id: 0))
+            topic = ::Karafka::Routing::Topic.new(topic_name, double(id: 0))
             messages = ::Karafka::Messages::Builders::Messages.call([message], topic, 0, Time.now)
             # NOTE: The following will iterate through the messages and create a new span representing
             #       the individual message processing (and `span` will refer to that particular span)
@@ -108,6 +106,34 @@
         end
       end
 
+      context 'when distributed tracing is disabled for the topic in particular' do
+        let(:topic_name) { 'special_topic' }
+
+        it 'does not continue the span that produced the message' do
+          consumer_span = nil
+          consumer_trace = nil
+
+          Datadog::Tracing.trace('consumer') do
+            consumer_span = Datadog::Tracing.active_span
+            consumer_trace = Datadog::Tracing.active_trace
+
+            topic = ::Karafka::Routing::Topic.new(topic_name, double(id: 0))
+            messages = ::Karafka::Messages::Builders::Messages.call([message], topic, 0, Time.now)
+            expect(messages).to all(be_a(::Karafka::Messages::Message))
+
+            # assert that the current trace re-set to the original trace after iterating the messages
+            expect(Datadog::Tracing.active_trace).to eq(consumer_trace)
+            expect(Datadog::Tracing.active_span).to eq(consumer_span)
+          end
+
+          expect(spans).to have(3).items
+
+          # assert that the message span is not continuation of the producer span
+          expect(span.parent_id).to eq(consumer_span.id)
+          expect(span.trace_id).to eq(consumer_trace.id)
+        end
+      end
+
       context 'when distributed tracing is not enabled' do
         let(:configuration_options) { {distributed_tracing: false} }
 
@@ -119,7 +145,8 @@
             consumer_span = Datadog::Tracing.active_span
             consumer_trace = Datadog::Tracing.active_trace
 
-            topic = ::Karafka::Routing::Topic.new('topic_a', double(id: 0))
+            topic = ::Karafka::Routing::Topic.new(topic_name, double(id: 0))
+
             messages = ::Karafka::Messages::Builders::Messages.call([message], topic, 0, Time.now)
             # NOTE: The following will iterate through the messages and create a new span representing
             #       the individual message processing (and `span` will refer to that particular span)
@@ -148,12 +175,11 @@
     let(:span_name) { Datadog::Tracing::Contrib::Karafka::Ext::SPAN_WORKER_PROCESS }
 
     it 'is expected to send a span' do
-      metadata = ::Karafka::Messages::Metadata.new
-      metadata['offset'] = 412
+      metadata = ::Karafka::Messages::Metadata.new(offset: 412, topic: 'topic_a')
       raw_payload = rand.to_s
 
       message = ::Karafka::Messages::Message.new(raw_payload, metadata)
-      job = double(executor: double(topic: double(name: 'topic_a', consumer: 'ABC'), partition: 0), messages: [message])
+      job = double(executor: double(topic: double(name: message.topic, consumer: 'ABC'), partition: 0), messages: [message])
 
       Karafka.monitor.instrument('worker.processed', {job: job}) do
         # Noop
@@ -169,4 +195,33 @@
       expect(span.resource).to eq 'ABC#consume'
     end
   end
+
+  describe 'framework auto-instrumentation' do
+    around do |example|
+      # Reset before and after each example; don't allow global state to linger.
+      Datadog.registry[:waterdrop].reset_configuration!
+      example.run
+      Datadog.registry[:waterdrop].reset_configuration!
+
+      # reset Karafka internal state as well
+      Karafka::App.config.internal.status.reset!
+      Karafka.refresh!
+    end
+
+    it 'automatically enables waterdrop instrumentation' do
+      Karafka::App.setup do |c|
+        c.kafka = {"bootstrap.servers": '127.0.0.1:9092'}
+      end
+
+      expect(Datadog.configuration.tracing[:karafka][:enabled]).to be true
+      expect(Datadog.configuration.tracing[:karafka][:distributed_tracing]).to be true
+      expect(Datadog.configuration.tracing[:karafka, 'special_topic'][:enabled]).to be true
+      expect(Datadog.configuration.tracing[:karafka, 'special_topic'][:distributed_tracing]).to be false
+
+      expect(Datadog.configuration.tracing[:waterdrop][:enabled]).to be true
+      expect(Datadog.configuration.tracing[:waterdrop][:distributed_tracing]).to be true
+      expect(Datadog.configuration.tracing[:waterdrop, 'special_topic'][:enabled]).to be true
+      expect(Datadog.configuration.tracing[:waterdrop, 'special_topic'][:distributed_tracing]).to be false
+    end
+  end
 end

spec/datadog/tracing/contrib/waterdrop/middleware_spec.rb

@@ -8,6 +8,7 @@
   before do
     Datadog.configure do |c|
       c.tracing.instrument :waterdrop, tracing_options
+      c.tracing.instrument :waterdrop, describes: /special_/, distributed_tracing: false
     end
   end
 
@@ -57,6 +58,17 @@
       end
     end
 
+    context 'when distributed tracing is disabled for the topic in particular' do
+      it 'does not propagate trace context in message headers' do
+        message_1 = {topic: 'special_topic', payload: 'foo'}
+        Datadog::Tracing.trace('test.span') do
+          middleware.call(message_1)
+        end
+
+        expect(message_1[:headers]).to be_nil
+      end
+    end
+
     context 'when DataStreams is enabled' do
       before do
         allow(Datadog::DataStreams).to receive(:enabled?).and_return(true)