Add meta_struct support on traces and spans

lib/datadog/appsec/context.rb

@@ -60,10 +60,13 @@ def extract_schema
       end
 
       def export_metrics
-        return if @span.nil?
+        # Required to satisfy steep, as @span is defined as nilable.
+        # According to soutaro, this is because instance variable can be changed by other threads.
+        span = @span
+        return if span.nil?
 
-        Metrics::Exporter.export_waf_metrics(@metrics.waf, @span)
-        Metrics::Exporter.export_rasp_metrics(@metrics.rasp, @span)
+        Metrics::Exporter.export_waf_metrics(@metrics.waf, span)
+        Metrics::Exporter.export_rasp_metrics(@metrics.rasp, span)
       end
 
       def finalize

lib/datadog/tracing/metadata.rb

@@ -2,6 +2,7 @@
 
 require_relative 'metadata/analytics'
 require_relative 'metadata/tagging'
+require_relative 'metadata/meta_struct'
 require_relative 'metadata/errors'
 
 module Datadog
@@ -11,6 +12,7 @@ module Metadata
       def self.included(base)
         base.include(Metadata::Tagging)
         base.include(Metadata::Errors)
+        base.include(Metadata::MetaStruct)
 
         # Additional extensions
         base.prepend(Metadata::Analytics)

lib/datadog/tracing/metadata/meta_struct.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Metadata
+      # Adds complex structures tagging behavior through meta_struct
+      # @public_api
+      module MetaStruct
+        def set_meta_struct(meta_struct)
+          self.meta_struct.merge!(meta_struct)
+        end
+
+        protected
+
+        def meta_struct
+          @meta_struct ||= {}
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/span.rb

@@ -21,6 +21,7 @@ class Span
         :end_time,
         :id,
         :meta,
+        :meta_struct,
         :metrics,
         :name,
         :parent_id,
@@ -53,6 +54,7 @@ def initialize(
         end_time: nil,
         id: nil,
         meta: nil,
+        meta_struct: nil,
         metrics: nil,
         parent_id: 0,
         resource: name,
@@ -75,6 +77,7 @@ def initialize(
         @trace_id = trace_id || Tracing::Utils.next_id
 
         @meta = meta || {}
+        @meta_struct = meta_struct || {}
         @metrics = metrics || {}
         @status = status || 0
 
@@ -144,6 +147,7 @@ def to_hash
           error: @status,
           meta: @meta,
           metrics: @metrics,
+          meta_struct: @meta_struct,
           name: @name,
           parent_id: @parent_id,
           resource: @resource,
@@ -185,12 +189,16 @@ def pretty_print(q)
               q.text "#{key} => #{value}"
             end
           end
-          q.group(2, 'Metrics: [', ']') do
+          q.group(2, 'Metrics: [', "]\n") do
             q.breakable
             q.seplist @metrics.each do |key, value|
               q.text "#{key} => #{value}"
             end
           end
+          q.group(2, 'Meta-Struct: [', ']') do
+            q.breakable
+            q.pp meta_struct
+          end
         end
       end
 

lib/datadog/tracing/span_operation.rb

@@ -289,6 +289,7 @@ def to_hash
           id: @id,
           meta: meta,
           metrics: metrics,
+          meta_struct: meta_struct,
           name: @name,
           parent_id: @parent_id,
           resource: @resource,
@@ -328,12 +329,16 @@ def pretty_print(q)
               q.text "#{key} => #{value}"
             end
           end
-          q.group(2, 'Metrics: [', ']') do
+          q.group(2, 'Metrics: [', "]\n") do
             q.breakable
             q.seplist metrics.each do |key, value|
               q.text "#{key} => #{value}"
             end
           end
+          q.group(2, 'Meta-Struct: [', ']') do
+            q.breakable
+            q.pp meta_struct
+          end
         end
       end
 
@@ -456,6 +461,7 @@ def build_span
           id: @id,
           meta: Core::Utils::SafeDup.frozen_or_dup(meta),
           metrics: Core::Utils::SafeDup.frozen_or_dup(metrics),
+          meta_struct: Core::Utils::SafeDup.frozen_or_dup(meta_struct),
           parent_id: @parent_id,
           resource: @resource,
           service: @service,

lib/datadog/tracing/trace_operation.rb

@@ -4,7 +4,7 @@
 require_relative '../core/utils'
 require_relative 'tracer'
 require_relative 'event'
-require_relative 'metadata/tagging'
+require_relative 'metadata'
 require_relative 'sampling/ext'
 require_relative 'span_operation'
 require_relative 'trace_digest'
@@ -25,7 +25,7 @@ module Tracing
     #
     # @public_api
     class TraceOperation
-      include Metadata::Tagging
+      include Metadata
 
       DEFAULT_MAX_LENGTH = 100_000
 
@@ -73,6 +73,7 @@ def initialize(
         profiling_enabled: nil,
         tags: nil,
         metrics: nil,
+        meta_struct: nil,
         trace_state: nil,
         trace_state_unknown_fields: nil,
         remote_parent: false,
@@ -105,6 +106,7 @@ def initialize(
         # Generic tags
         set_tags(tags) if tags
         set_tags(metrics) if metrics
+        set_meta_struct(meta_struct) if meta_struct
 
         # State
         @root_span = nil
@@ -369,6 +371,7 @@ def fork_clone
           trace_state_unknown_fields: (@trace_state_unknown_fields && @trace_state_unknown_fields.dup),
           tags: meta.dup,
           metrics: metrics.dup,
+          meta_struct: meta_struct.dup,
           remote_parent: @remote_parent
         )
       end
@@ -508,6 +511,7 @@ def build_trace(spans, partial = false)
           service: service,
           tags: meta,
           metrics: metrics,
+          meta_struct: meta_struct,
           root_span_id: !partial ? root_span && root_span.id : nil,
           profiling_enabled: @profiling_enabled,
         )

lib/datadog/tracing/trace_segment.rb

@@ -58,6 +58,7 @@ def initialize(
         service: nil,
         tags: nil,
         metrics: nil,
+        meta_struct: nil,
         profiling_enabled: nil
       )
         @id = id
@@ -68,6 +69,7 @@ def initialize(
         # The caller is expected to have done that
         @meta = (tags && tags.dup) || {}
         @metrics = (metrics && metrics.dup) || {}
+        @meta_struct = (meta_struct && meta_struct.dup) || {}
 
         # Set well-known tags, defaulting to getting the values from tags
         @agent_sample_rate = agent_sample_rate || agent_sample_rate_tag
@@ -146,7 +148,8 @@ def high_order_tid
       attr_reader \
         :root_span_id,
         :meta,
-        :metrics
+        :metrics,
+        :meta_struct
 
       private
 

lib/datadog/tracing/transport/serializable_trace.rb

@@ -65,7 +65,7 @@ def initialize(span, native_events_supported)
         def to_msgpack(packer = nil)
           packer ||= MessagePack::Packer.new
 
-          number_of_elements_to_write = 11
+          number_of_elements_to_write = 12
 
           number_of_elements_to_write += 1 if span.events.any? && @native_events_supported
 
@@ -113,6 +113,9 @@ def to_msgpack(packer = nil)
           packer.write(span.meta)
           packer.write('metrics')
           packer.write(span.metrics)
+          packer.write('meta_struct')
+          # We encapsulate the resulting msgpack in a binary msgpack
+          packer.write(span.meta_struct.transform_values(&:to_msgpack))
           packer.write('span_links')
           packer.write(span.links.map(&:to_hash))
           packer.write('error')

lib/datadog/tracing/transport/trace_formatter.rb

@@ -43,6 +43,7 @@ def format!
           # Apply generic trace tags. Any more specific value will be overridden
           # by the subsequent calls below.
           set_trace_tags!
+          set_meta_struct!
 
           set_resource!
 
@@ -89,6 +90,12 @@ def set_trace_tags!
           root_span.set_tags(trace.send(:metrics))
         end
 
+        def set_meta_struct!
+          return if partial?
+
+          root_span.set_meta_struct(trace.send(:meta_struct))
+        end
+
         def tag_agent_sample_rate!
           return unless trace.agent_sample_rate
 

lib/datadog/tracing/transport/traces.rb

@@ -125,6 +125,7 @@ def initialize(apis, default_api)
           end
 
           def send_traces(traces)
+            # object that extends Datadog::Core::Encoding::Encoder (MsgpackEncoder or JSONEncoder)
             encoder = current_api.encoder
             chunker = Datadog::Tracing::Transport::Traces::Chunker.new(encoder)
 

sig/datadog/appsec/context.rbs

@@ -3,9 +3,9 @@ module Datadog
     class Context
       type input_data = SecurityEngine::Runner::input_data
 
-      @trace: Tracing::TraceOperation
+      @trace: Tracing::TraceOperation?
 
-      @span: Tracing::SpanOperation
+      @span: Tracing::SpanOperation?
 
       @events: ::Array[untyped]
 
@@ -17,17 +17,17 @@ module Datadog
 
       ActiveContextError: ::StandardError
 
-      attr_reader trace: Tracing::TraceOperation
+      attr_reader trace: Tracing::TraceOperation?
 
-      attr_reader span: Tracing::SpanOperation
+      attr_reader span: Tracing::SpanOperation?
 
       attr_reader events: ::Array[untyped]
 
       def self.activate: (Context context) -> Context
 
       def self.deactivate: () -> void
 
-      def self.active: () -> Context
+      def self.active: () -> Context?
 
       def initialize: (Tracing::TraceOperation trace, Tracing::SpanOperation span, AppSec::Processor security_engine) -> void
 

sig/datadog/appsec/event.rbs

@@ -24,7 +24,7 @@ module Datadog
 
       def self.gzip: (untyped value) -> untyped
 
-      def self.add_distributed_tags: (Tracing::TraceOperation trace) -> void
+      def self.add_distributed_tags: (Tracing::TraceOperation? trace) -> void
     end
   end
 end

sig/datadog/tracing/metadata.rbs

@@ -3,6 +3,7 @@ module Datadog
     module Metadata
       include Metadata::Tagging
       include Metadata::Errors
+      include Metadata::MetaStruct
       prepend Metadata::Analytics
     end
   end

sig/datadog/tracing/metadata/meta_struct.rbs

@@ -0,0 +1,11 @@
+module Datadog
+  module Tracing
+    module Metadata
+      module MetaStruct
+        def set_meta_struct: (Hash[String, untyped] meta_struct) -> void
+
+        def meta_struct: () -> Hash[String, untyped]
+      end
+    end
+  end
+end

sig/datadog/tracing/span.rbs

@@ -3,44 +3,49 @@ module Datadog
     class Span
       include Metadata
 
-      attr_accessor end_time: (Time | nil)
+      attr_accessor end_time: Time?
       attr_accessor id: Integer
       attr_accessor meta: Hash[String, String]
+      attr_accessor meta_struct: Hash[String, untyped]
       attr_accessor metrics: Hash[String, Float]
       attr_accessor name: String
       attr_accessor parent_id: Integer
       attr_accessor resource: String
-      attr_accessor service: (String | nil)
-      attr_accessor links: Array[untyped]
-      attr_accessor events: Array[untyped]
-      attr_accessor type: (String | nil)
-      attr_accessor start_time: (Time | nil)
+      attr_accessor service: String?
+      attr_accessor links: Array[Datadog::Tracing::SpanLink]
+      attr_accessor events: Array[Datadog::Tracing::SpanEvent]
+      attr_accessor type: String?
+      attr_accessor start_time: Time?
       attr_accessor status: Integer
       attr_accessor trace_id: Integer
-      attr_writer duration: (Float | nil)
+      attr_writer duration: Float?
 
       def initialize: (
           String name,
-          ?duration: (Float | nil),
-          ?end_time: (Time | nil),
-          ?id: (Integer | nil),
-          ?meta: (Hash[String, String] | nil),
-          ?metrics: (Hash[String, Float] | nil),
+          ?duration: Float?,
+          ?end_time: Time?,
+          ?id: Integer?,
+          ?meta: Hash[String, String]?,
+          ?meta_struct: Hash[String, untyped]?,
+          ?metrics: Hash[String, Float]?,
           ?parent_id: Integer,
           ?resource: String,
-          ?service: (String | nil),
-          ?start_time: (Time | nil),
+          ?service: String?,
+          ?start_time: Time?,
           ?status: Integer,
-          ?type: (String | nil),
-          ?trace_id: (Integer | nil),
-          ?service_entry: (bool | nil),
-          ?links: (Array[untyped] | nil),
-          ?events: (Array[untyped] | nil)
+          ?type: String?,
+          ?trace_id: Integer?,
+          ?service_entry: bool?,
+          ?links: Array[Datadog::Tracing::SpanLink]?,
+          ?events: Array[Datadog::Tracing::SpanEvent]?
         ) -> void
 
       def started?: -> bool
       def stopped?: -> bool
-      def duration: -> (Float | nil)
+
+      alias finished? stopped?
+
+      def duration: -> Float?
       def set_error: (Exception e) -> void
       def ==: (Span other) -> bool
       def to_s: -> String