Support different waterdrop configurations per topic

Author: Drowze

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/waterdrop/integration.rb

@@ -37,6 +37,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/monitor.rb

@@ -18,10 +18,6 @@ module Monitor
             message.produced_sync
           ].freeze
 
-          def configuration
-            Datadog.configuration.tracing[:waterdrop]
-          end
-
           def instrument(event_id, payload = {}, &block)
             return super unless TRACEABLE_EVENTS.include?(event_id)
 
@@ -40,15 +36,15 @@ def instrument(event_id, payload = {}, &block)
 
                 span.set_tag(Contrib::Karafka::Ext::TAG_MESSAGE_COUNT, payload[:messages].size)
 
-                payload[:messages].each { |message| inject(trace_digest, message) } if configuration[:distributed_tracing]
+                payload[:messages].each { |message| inject(trace_digest, message) }
               else
                 action = event_id.sub('message.produced', 'produce')
 
                 span.set_tag(Contrib::Ext::Messaging::TAG_DESTINATION, payload[:message][:topic])
                 span.set_tag(Contrib::Karafka::Ext::TAG_PARTITION, payload[:message][:partition])
                 span.set_tag(Contrib::Karafka::Ext::TAG_MESSAGE_COUNT, 1)
 
-                inject(trace_digest, payload[:message]) if configuration[:distributed_tracing]
+                inject(trace_digest, payload[:message])
               end
 
               span.resource = "waterdrop.#{action}"
@@ -63,9 +59,17 @@ def instrument(event_id, payload = {}, &block)
           private
 
           def inject(trace_digest, message)
+            return unless datadog_configuration(message[:topic])[:distributed_tracing]
+
             message[:headers] ||= {}
             WaterDrop.inject(trace_digest, message[:headers])
           end
+
+          # cache the configuration resolution per topic to avoid repeated lookups in message batches
+          def datadog_configuration(topic)
+            @datadog_configuration ||= {}
+            @datadog_configuration[topic] ||= Datadog.configuration.tracing[:waterdrop, topic]
+          end
         end
       end
     end

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

@@ -199,4 +199,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/monitor_spec.rb

@@ -8,8 +8,6 @@
 end
 require 'datadog'
 
-puts "waterdrop version: #{WaterDrop::VERSION}"
-
 RSpec.describe 'Waterdrop monitor' do
   before do
     Datadog.configure do |c|