Add sampling info to Tracestate headers

CHANGELOG.asciidoc

@@ -35,6 +35,12 @@ endif::[]
 [[release-notes-3.x]]
 === Ruby Agent version 3.x
 
+[float]
+[[unreleased]]
+==== Unreleased
+
+- Add and read sampling info from Tracestate headers {pull}858[#858]
+
 [[release-notes-3.10.1]]
 ==== 3.10.1 (2020-08-26)
 

lib/elastic_apm/instrumenter.rb

@@ -119,7 +119,13 @@ def start_transaction(
           "Already inside #{transaction.inspect}"
       end
 
-      sampled = trace_context ? trace_context.recorded? : random_sample?(config)
+      if trace_context
+        samled = trace_context.recorded?
+        sample_rate = trace_context.tracestate.sample_rate
+      else
+        sampled = random_sample?(config)
+        sample_rate = config.transaction_sample_rate
+      end
 
       transaction =
         Transaction.new(
@@ -128,6 +134,7 @@ def start_transaction(
           context: context,
           trace_context: trace_context,
           sampled: sampled,
+          sample_rate: sample_rate,
           config: config
         )
 

lib/elastic_apm/span.rb

@@ -38,7 +38,8 @@ def initialize(
       action: nil,
       context: nil,
       stacktrace_builder: nil,
-      sync: nil
+      sync: nil,
+      sample_rate: nil
     )
       @name = name
 
@@ -53,6 +54,7 @@ def initialize(
       @transaction = transaction
       @parent = parent
       @trace_context = trace_context || parent.trace_context.child
+      @sample_rate = transaction.sample_rate
 
       @context = context || Span::Context.new(sync: sync)
       @stacktrace_builder = stacktrace_builder
@@ -73,6 +75,7 @@ def initialize(
       :context,
       :duration,
       :parent,
+      :sample_rate,
       :self_time,
       :stacktrace,
       :timestamp,

lib/elastic_apm/trace_context.rb

@@ -33,7 +33,7 @@ def initialize(
       **legacy_traceparent_attrs
     )
       @traceparent = traceparent || Traceparent.new(**legacy_traceparent_attrs)
-      @tracestate = tracestate
+      @tracestate = tracestate || Tracestate.new
     end
 
     attr_accessor :traceparent, :tracestate

lib/elastic_apm/trace_context/traceparent.rb

@@ -33,20 +33,18 @@ def initialize(
         trace_id: nil,
         span_id: nil,
         id: nil,
-        recorded: true,
-        tracestate: nil
+        recorded: true
       )
         @version = version
         @trace_id = trace_id || hex(TRACE_ID_LENGTH)
         # TODO: rename span_id kw arg to parent_id with next major version bump
         @parent_id = span_id
         @id = id || hex(ID_LENGTH)
         @recorded = recorded
-        @tracestate = tracestate
       end
       # rubocop:enable Metrics/ParameterLists
 
-      attr_accessor :version, :id, :trace_id, :parent_id, :recorded, :tracestate
+      attr_accessor :version, :id, :trace_id, :parent_id, :recorded
 
       alias :recorded? :recorded
 

lib/elastic_apm/trace_context/tracestate.rb

@@ -21,22 +21,127 @@ module ElasticAPM
   class TraceContext
     # @api private
     class Tracestate
-      def initialize(values = [])
-        @values = values
+      # @api private
+      class Entry
+        def initialize(key, value)
+          @key = key
+          @value = value
+        end
+
+        attr_reader :key, :value
+
+        def to_s
+          "#{key}=#{value}"
+        end
+      end
+
+      class EsEntry
+        ASSIGN = ':'
+        SPLIT = ';'
+
+        SHORT_TO_LONG = { 's' => 'sample_rate' }
+        LONG_TO_SHORT = { 'sample_rate' => 's' }
+
+        def initialize(values = nil)
+          parse(values)
+        end
+
+        attr_reader :sample_rate
+
+        def key
+          'es'
+        end
+
+        def value
+          LONG_TO_SHORT.map do |l, s|
+            "#{s}#{ASSIGN}#{send(l)}"
+          end.join(SPLIT)
+        end
+
+        def empty?
+          !sample_rate
+        end
+
+        def sample_rate=(val)
+          float = Float(val).round(3)
+
+          return nil unless (0.0..1.0).include?(float)
+
+          @sample_rate = float
+        rescue ArgumentError => e
+          nil
+        end
+
+        def to_s
+          return nil if empty?
+
+          "es=#{value}"
+        end
+
+        private
+
+        def parse(values)
+          return unless values
+
+          values.split(SPLIT).map do |kv|
+            k, v = kv.split(ASSIGN)
+            next unless SHORT_TO_LONG.keys.include?(k)
+            send("#{SHORT_TO_LONG[k]}=", v)
+          end
+        end
+      end
+
+      extend Forwardable
+
+      def initialize(entries: {}, sample_rate: nil)
+        @entries = entries
+
+        self.sample_rate = sample_rate if sample_rate
       end
 
-      attr_accessor :values
+      attr_accessor :entries
+
+      def_delegators :es_entry, :sample_rate, :sample_rate=
 
       def self.parse(header)
-        # HTTP allows multiple headers with the same name, eg. multiple
-        # Set-Cookie headers per response.
-        # Rack handles this by joining the headers under the same key, separated
-        # by newlines, see https://www.rubydoc.info/github/rack/rack/file/SPEC
-        new(String(header).split("\n"))
+        entries =
+          split_by_nl_and_comma(header)
+          .each_with_object({}) do |entry, hsh|
+            k, v = entry.split('=')
+
+            hsh[k] =
+              case k
+              when 'es' then EsEntry.new(v)
+              else Entry.new(k, v)
+              end
+          end
+
+        new(entries: entries)
       end
 
       def to_header
-        values.join(',')
+        return "" unless entries.any?
+
+        entries.values.map(&:to_s).join(',')
+      end
+
+      private
+
+      def es_entry
+        # lazy generate this so we only add it if necessary
+        entries['es'] ||= EsEntry.new
+      end
+
+      class << self
+        private
+
+        def split_by_nl_and_comma(str)
+          # HTTP allows multiple headers with the same name, eg. multiple
+          # Set-Cookie headers per response.
+          # Rack handles this by joining the headers under the same key, separated
+          # by newlines, see https://www.rubydoc.info/github/rack/rack/file/SPEC
+          String(str).split("\n").map { |s| s.split(',') }.flatten
+        end
       end
     end
   end

lib/elastic_apm/transaction.rb

@@ -34,6 +34,7 @@ def initialize(
       name = nil,
       type = nil,
       sampled: true,
+      sample_rate: 1,
       context: nil,
       config:,
       trace_context: nil
@@ -52,13 +53,19 @@ def initialize(
       @default_labels = config.default_labels
 
       @sampled = sampled
+      @sample_rate = sample_rate
 
       @context = context || Context.new # TODO: Lazy generate this?
       if @default_labels
         Util.reverse_merge!(@context.labels, @default_labels)
       end
 
-      @trace_context = trace_context || TraceContext.new(recorded: sampled)
+      unless (@trace_context = trace_context)
+        @trace_context = TraceContext.new(
+          traceparent: TraceContext::Traceparent.new(recorded: sampled),
+          tracestate: TraceContext::Tracestate.new(sample_rate: sampled ? sample_rate : 0)
+        )
+      end
 
       @started_spans = 0
       @dropped_spans = 0
@@ -69,10 +76,22 @@ def initialize(
 
     attr_accessor :name, :type, :result
 
-    attr_reader :context, :duration, :started_spans, :dropped_spans,
-      :timestamp, :trace_context, :notifications, :self_time,
-      :span_frames_min_duration, :collect_metrics, :breakdown_metrics,
-      :framework_name, :transaction_max_spans
+    attr_reader(
+      :breakdown_metrics,
+      :collect_metrics,
+      :context,
+      :dropped_spans,
+      :duration,
+      :framework_name,
+      :notifications,
+      :self_time,
+      :sample_rate,
+      :span_frames_min_duration,
+      :started_spans,
+      :timestamp,
+      :trace_context,
+      :transaction_max_spans 
+    )
 
     def sampled?
       @sampled

lib/elastic_apm/transport/serializers/span_serializer.rb

@@ -42,7 +42,8 @@ def build(span)
               context: context_serializer.build(span.context),
               stacktrace: span.stacktrace.to_a,
               timestamp: span.timestamp,
-              trace_id: span.trace_id
+              trace_id: span.trace_id,
+              sample_rate: span.sample_rate
             }
           }
         end

lib/elastic_apm/transport/serializers/transaction_serializer.rb

@@ -38,6 +38,7 @@ def build(transaction)
               duration: ms(transaction.duration),
               timestamp: transaction.timestamp,
               sampled: transaction.sampled?,
+              sample_rate: transaction.sample_rate,
               context: context_serializer.build(transaction.context),
               span_count: {
                 started: transaction.started_spans,

spec/elastic_apm/middleware_spec.rb

@@ -74,6 +74,7 @@ class MiddlewareTestError < StandardError; end
       trace_context = @intercepted.transactions.first.trace_context
       expect(trace_context).to_not be_nil
       expect(trace_context).to be_recorded
+      expect(trace_context.tracestate.sample_rate).to_not be nil
     end
 
     describe 'Distributed Tracing' do
@@ -108,14 +109,14 @@ class MiddlewareTestError < StandardError; end
                 '/',
                 'HTTP_ELASTIC_APM_TRACEPARENT' =>
                 '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00',
-                'HTTP_TRACESTATE' => 'thing=value'
+                'HTTP_TRACESTATE' => 'es=s:0.75,abc=123'
               )
             )
           end
 
           trace_context = @intercepted.transactions.first.trace_context
           expect(trace_context.tracestate).to be_a(TraceContext::Tracestate)
-          expect(trace_context.tracestate.values).to match(['thing=value'])
+          expect(trace_context.tracestate.to_header).to match('es=s:0.75,abc=123')
         end
       end
 

spec/elastic_apm/span_spec.rb

@@ -48,6 +48,7 @@ module ElasticAPM
       its(:trace_id) { should eq trace_context.trace_id }
       its(:id) { should eq trace_context.id }
       its(:parent_id) { should eq trace_context.parent_id }
+      its(:sample_rate) { is_expected.to eq transaction.sample_rate }
 
       context 'with a dot-separated type' do
         it 'splits type' do
@@ -73,7 +74,7 @@ module ElasticAPM
       subject do
         described_class.new(
           name: 'Spannest name',
-          transaction: transaction.id,
+          transaction: transaction,
           parent: transaction,
           trace_context: trace_context
         )
@@ -93,7 +94,7 @@ module ElasticAPM
       subject do
         described_class.new(
           name: 'Spannest name',
-          transaction: transaction.id,
+          transaction: transaction,
           parent: transaction,
           trace_context: trace_context
         )