Add setting a tag at the trace-level

lib/datadog/appsec/contrib/rack/gateway/watcher.rb

@@ -27,7 +27,6 @@ def self.watch
                   # TODO: factor out
                   if defined?(Datadog::Tracing) && Datadog::Tracing.respond_to?(:active_span)
                     active_trace = Datadog::Tracing.active_trace
-                    root_span = active_trace.send(:root_span) if active_trace
                     active_span = Datadog::Tracing.active_span
 
                     Datadog.logger.debug { "active span: #{active_span.span_id}" } if active_span
@@ -45,7 +44,6 @@ def self.watch
                       event = {
                         waf_result: result,
                         trace: active_trace,
-                        root_span: root_span,
                         span: active_span,
                         request: request,
                         action: action
@@ -77,7 +75,6 @@ def self.watch
                   # TODO: factor out
                   if defined?(Datadog::Tracing) && Datadog::Tracing.respond_to?(:active_span)
                     active_trace = Datadog::Tracing.active_trace
-                    root_span = active_trace.send(:root_span) if active_trace
                     active_span = Datadog::Tracing.active_span
 
                     Datadog.logger.debug { "active span: #{active_span.span_id}" } if active_span
@@ -95,7 +92,6 @@ def self.watch
                       event = {
                         waf_result: result,
                         trace: active_trace,
-                        root_span: root_span,
                         span: active_span,
                         response: response,
                         action: action

lib/datadog/appsec/event.rb

@@ -43,26 +43,20 @@ def self.record(*events)
       end
 
       # rubocop:disable Metrics/AbcSize
-      # rubocop:disable Metrics/MethodLength
       def self.record_via_span(*events)
-        events.group_by { |e| e[:root_span] }.each do |root_span, event_group|
-          unless root_span
-            Datadog.logger.debug { "{ error: 'no root span: cannot record', event_group: #{event_group.inspect}}" }
+        events.group_by { |e| e[:trace] }.each do |trace, event_group|
+          unless trace
+            Datadog.logger.debug { "{ error: 'no trace: cannot record', event_group: #{event_group.inspect}}" }
             next
           end
 
-          # TODO: this is a hack but there is no API to do that
-          root_span_tags = root_span.send(:meta).keys
+          trace.keep!
 
           # prepare and gather tags to apply
-          root_tags = event_group.each_with_object({}) do |event, tags|
+          trace_tags = event_group.each_with_object({}) do |event, tags|
             span = event[:span]
-            trace = event[:trace]
 
-            if span
-              span.set_tag('appsec.event', 'true')
-              trace.keep!
-            end
+            span.set_tag('appsec.event', 'true') if span
 
             request = event[:request]
             response = event[:response]
@@ -98,16 +92,15 @@ def self.record_via_span(*events)
           # apply tags to root span
 
           # complex types are unsupported, we need to serialize to a string
-          triggers = root_tags.delete('_dd.appsec.triggers')
-          root_span.set_tag('_dd.appsec.json', JSON.dump({ triggers: triggers }))
+          triggers = trace_tags.delete('_dd.appsec.triggers')
+          trace.send(:set_tag, '_dd.appsec.json', JSON.dump({ triggers: triggers }))
 
-          root_tags.each do |key, value|
-            root_span.set_tag(key, value.is_a?(String) ? value.encode('UTf-8') : value) unless root_span_tags.include?(key)
+          trace_tags.each do |key, value|
+            trace.send(:set_tag, key, value.is_a?(String) ? value.encode('UTF-8') : value)
           end
         end
       end
       # rubocop:enable Metrics/AbcSize
-      # rubocop:enable Metrics/MethodLength
     end
   end
 end

lib/datadog/tracing/trace_operation.rb

@@ -63,7 +63,8 @@ def initialize(
         sample_rate: nil,
         sampled: nil,
         sampling_priority: nil,
-        service: nil
+        service: nil,
+        tags: nil
       )
         # Attributes
         @events = events || Events.new
@@ -84,6 +85,9 @@ def initialize(
         @sampling_priority = sampling_priority
         @service = service
 
+        # Generic tags
+        @tags = set_tags(tags)
+
         # State
         @root_span = nil
         @active_span = nil
@@ -239,7 +243,7 @@ def to_digest
           trace_resource: @resource,
           trace_runtime_id: Core::Environment::Identity.id,
           trace_sampling_priority: @sampling_priority,
-          trace_service: @service
+          trace_service: @service,
         ).freeze
       end
 
@@ -261,7 +265,8 @@ def fork_clone
           sample_rate: @sample_rate,
           sampled: @sampled,
           sampling_priority: @sampling_priority,
-          service: (@service && @service.dup)
+          service: (@service && @service.dup),
+          tags: @tags.dup
         )
       end
 
@@ -318,6 +323,22 @@ def subscribe_deactivate_trace(&block)
         :events,
         :root_span
 
+      def tags
+        @tags ||= {}
+      end
+
+      def get_tag(key)
+        tags[key]
+      end
+
+      def set_tag(key, value = nil)
+        tags[key] = value
+      end
+
+      def set_tags(tags)
+        (tags || {}).each { |k, v| set_tag(k, v) }
+      end
+
       def activate_span!(span_op)
         parent = @active_span
 
@@ -404,6 +425,7 @@ def build_trace(spans, partial = false)
           name: @name,
           resource: @resource,
           service: @service,
+          tags: !partial ? @tags : nil,
           root_span_id: !partial ? @root_span && @root_span.id : nil
         )
       end

lib/datadog/tracing/trace_segment.rb

@@ -49,12 +49,13 @@ def initialize(
         runtime_id: nil,
         sample_rate: nil,
         sampling_priority: nil,
-        service: nil
+        service: nil,
+        tags: nil
       )
         @id = id
         @root_span_id = root_span_id
         @spans = spans || []
-        @tags = {}
+        @tags = tags || {}
 
         # Set well-known tags
         self.agent_sample_rate = agent_sample_rate

lib/ddtrace/transport/trace_formatter.rb

@@ -32,6 +32,11 @@ def format!
         # trace metadata, we must put our trace
         # metadata on the root span. This "root span"
         # is needed by the agent/API to ingest the trace.
+
+        # Apply generic trace tags. Any more specific value will be overridden
+        # by the subsequent calls below.
+        set_trace_tags!
+
         set_resource!
 
         tag_agent_sample_rate!
@@ -59,6 +64,10 @@ def set_resource!
         root_span.resource = trace.resource
       end
 
+      def set_trace_tags!
+        root_span.set_tags(trace.tags)
+      end
+
       def tag_agent_sample_rate!
         return unless trace.agent_sample_rate
 

spec/datadog/tracing/trace_operation_spec.rb

@@ -775,6 +775,87 @@
     end
   end
 
+  describe '#get_tag' do
+    before do
+      trace_op.send(:set_tag, 'foo', 'bar')
+    end
+
+    it 'gets tag set on trace' do
+      expect(trace_op.send(:get_tag, 'foo')).to eq('bar')
+    end
+
+    it 'gets unset tag as nil' do
+      expect(trace_op.send(:get_tag, 'unset')).to be_nil
+    end
+  end
+
+  describe '#set_tag' do
+    it 'sets tag on trace before a measurement' do
+      trace_op.send(:set_tag, 'foo', 'bar')
+      trace_op.measure('top') {}
+
+      trace = trace_op.flush!
+
+      expect(trace.tags).to include('foo' => 'bar')
+    end
+
+    it 'sets tag on trace after a measurement' do
+      trace_op.measure('top') {}
+      trace_op.send(:set_tag, 'foo', 'bar')
+
+      trace = trace_op.flush!
+
+      expect(trace.tags).to include('foo' => 'bar')
+    end
+
+    it 'sets tag on trace from a measurement' do
+      trace_op.measure('top') do
+        trace_op.send(:set_tag, 'foo', 'bar')
+      end
+
+      trace = trace_op.flush!
+
+      expect(trace.tags).to include('foo' => 'bar')
+    end
+
+    it 'sets tag on trace from a nested measurement' do
+      trace_op.measure('grandparent') do
+        trace_op.measure('parent') do
+          trace_op.send(:set_tag, 'foo', 'bar')
+        end
+      end
+
+      trace = trace_op.flush!
+
+      expect(trace.spans).to have(2).items
+      expect(trace.spans.map(&:name)).to include('parent')
+      expect(trace.tags).to include('foo' => 'bar')
+    end
+
+    context 'with partial flushing' do
+      subject(:flush!) { trace_op.flush! }
+      let(:trace) { flush! }
+
+      it 'sets tag on trace from a nested measurement' do
+        trace_op.measure('grandparent') do
+          trace_op.measure('parent') do
+            trace_op.send(:set_tag, 'foo', 'bar')
+          end
+          flush!
+        end
+
+        expect(trace.spans).to have(1).items
+        expect(trace.spans.map(&:name)).to include('parent')
+        expect(trace.tags).to_not include('foo' => 'bar')
+
+        final_flush = trace_op.flush!
+        expect(final_flush.spans).to have(1).items
+        expect(final_flush.spans.map(&:name)).to include('grandparent')
+        expect(final_flush.tags).to include('foo' => 'bar')
+      end
+    end
+  end
+
   describe '#build_span' do
     subject(:build_span) { trace_op.build_span(span_name, **span_options) }
     let(:span_name) { 'web.request' }

spec/ddtrace/transport/trace_formatter_spec.rb

@@ -16,6 +16,10 @@
   let(:trace_options) { {} }
 
   shared_context 'trace metadata' do
+    let(:trace_tags) do
+      nil
+    end
+
     let(:trace_options) do
       {
         resource: resource,
@@ -28,7 +32,8 @@
         rule_sample_rate: rule_sample_rate,
         runtime_id: runtime_id,
         sample_rate: sample_rate,
-        sampling_priority: sampling_priority
+        sampling_priority: sampling_priority,
+        tags: trace_tags
       }
     end
 
@@ -45,6 +50,17 @@
     let(:sampling_priority) { Datadog::Tracing::Sampling::Ext::Priority::USER_KEEP }
   end
 
+  shared_context 'trace metadata with tags' do
+    include_context 'trace metadata'
+
+    let(:trace_tags) do
+      {
+        'foo' => 'bar',
+        'baz' => 42
+      }
+    end
+  end
+
   shared_context 'no root span' do
     let(:trace) { Datadog::Tracing::TraceSegment.new(spans, **trace_options) }
     let(:spans) { Array.new(3) { Datadog::Tracing::Span.new('my.job') } }
@@ -139,6 +155,16 @@
         end
       end
 
+      shared_examples 'root span with generic tags' do
+        context 'metrics' do
+          it { expect(root_span.metrics).to include({ 'baz' => 42 }) }
+        end
+
+        context 'meta' do
+          it { expect(root_span.meta).to include({ 'foo' => 'bar' }) }
+        end
+      end
+
       context 'with no root span' do
         include_context 'no root span'
 
@@ -159,6 +185,16 @@
 
           it_behaves_like 'root span with tags'
         end
+
+        context 'when trace has metadata set with generic tags' do
+          include_context 'trace metadata with tags'
+
+          it { is_expected.to be(trace) }
+          it { expect(root_span.resource).to eq('my.job') }
+
+          it_behaves_like 'root span with tags'
+          it_behaves_like 'root span with generic tags'
+        end
       end
 
       context 'with missing root span' do
@@ -181,6 +217,16 @@
 
           it_behaves_like 'root span with tags'
         end
+
+        context 'when trace has metadata set with generic tags' do
+          include_context 'trace metadata with tags'
+
+          it { is_expected.to be(trace) }
+          it { expect(root_span.resource).to eq('my.job') }
+
+          it_behaves_like 'root span with tags'
+          it_behaves_like 'root span with generic tags'
+        end
       end
 
       context 'with a root span' do