Simple produce checkpoint

Author: ericfirth

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

@@ -19,6 +19,17 @@ def call(message)
                 WaterDrop.inject(trace_op.to_digest, message[:headers] ||= {})
               end
 
+              if Datadog::DataStreams.enabled?
+                Datadog::DataStreams.set_produce_checkpoint(
+                  type: 'kafka',
+                  destination: message[:topic],
+                  auto_instrumentation: true
+                ) do |key, value|
+                  message[:headers] ||= {}
+                  message[:headers][key] = value
+                end
+              end
+
               message
             end
 

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

@@ -28,6 +28,13 @@ def patch
 
               included_middlewares = producer.middleware.instance_variable_get(:@steps)
               producer.middleware.append(Middleware) unless included_middlewares.include?(Middleware)
+
+              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

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

@@ -60,5 +60,65 @@
         )
       end
     end
+
+    context 'when DataStreams is enabled' do
+      before do
+        allow(Datadog::DataStreams).to receive(:enabled?).and_return(true)
+        allow(Datadog::DataStreams).to receive(:set_produce_checkpoint) do |**_kwargs, &block|
+          block.call('data_streams_key', 'data_streams_value')
+        end
+      end
+
+      it 'calls set_produce_checkpoint and injects headers' do
+        message = {topic: 'some_topic', payload: 'hello'}
+
+        middleware.call(message)
+
+        expect(Datadog::DataStreams).to have_received(:set_produce_checkpoint).with(
+          type: 'kafka',
+          destination: 'some_topic',
+          auto_instrumentation: true
+        )
+        expect(message[:headers]).to include('data_streams_key' => 'data_streams_value')
+      end
+
+      it 'initializes headers if not present' do
+        message = {topic: 'some_topic', payload: 'hello'}
+
+        middleware.call(message)
+
+        expect(Datadog::DataStreams).to have_received(:set_produce_checkpoint).with(
+          type: 'kafka',
+          destination: 'some_topic',
+          auto_instrumentation: true
+        )
+      end
+
+      it 'preserves existing headers' do
+        message = {topic: 'some_topic', payload: 'hello', headers: {'existing' => 'header'}}
+
+        middleware.call(message)
+
+        expect(message[:headers]).to include(
+          'data_streams_key' => 'data_streams_value',
+          'existing' => 'header'
+        )
+      end
+    end
+
+    context 'when DataStreams is disabled' do
+      before do
+        allow(Datadog::DataStreams).to receive(:enabled?).and_return(false)
+        allow(Datadog::DataStreams).to receive(:set_produce_checkpoint)
+      end
+
+      it 'does not call set_produce_checkpoint' do
+        message = {topic: 'some_topic', payload: 'hello'}
+
+        middleware.call(message)
+
+        expect(Datadog::DataStreams).not_to have_received(:set_produce_checkpoint)
+      end
+    end
   end
 end

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

@@ -72,5 +72,19 @@
         )
       end
     end
+
+    context 'when DataStreams is enabled' do
+      before do
+        allow(Datadog::DataStreams).to receive(:enabled?).and_return(true)
+      end
+
+      it 'patches without errors' do
+        expect do
+          WaterDrop::Producer.new do |config|
+            config.client_class = WaterDrop::Clients::Buffered
+          end
+        end.not_to raise_error
+      end
+    end
   end
 end