consistent_probability_tracestate.rb in opentelemetry-sdk-experimental-0.1.1

- old
+ new

@@ -15,10 +15,35 @@
         module ConsistentProbabilityTraceState
           DECIMAL = /\A\d+\z/.freeze
           MAX_LIST_LENGTH = 256 # Defined by https://www.w3.org/TR/trace-context/
           private_constant(:DECIMAL, :MAX_LIST_LENGTH)
 
+          private
+
+          # sanitized_tracestate returns an OpenTelemetry Tracestate object with the
+          # tracestate sanitized according to the Context invariants defined in the
+          # tracestate probability sampling spec.
+          #
+          # @param span_context [OpenTelemetry::Trace::SpanContext] the parent span context
+          # @return [OpenTelemetry::Trace::Tracestate] the sanitized tracestate
+          def sanitized_tracestate(span_context)
+            sampled = span_context.trace_flags.sampled?
+            tracestate = span_context.tracestate
+            parse_ot_vendor_tag(tracestate) do |p, r, rest|
+              if !r.nil? && r > 62
+                p = r = nil
+              elsif !p.nil? && p > 63
+                p = nil
+              elsif !p.nil? && !r.nil? && !invariant(p, r, sampled)
+                p = nil
+              else
+                return tracestate
+              end
+              update_tracestate(tracestate, p, r, rest)
+            end
+          end
+
           # parse_ot_vendor_tag parses the 'ot' vendor tag of the tracestate.
           # It yields the parsed probability fields and the remaining tracestate.
           # It returns the result of the block.
           def parse_ot_vendor_tag(tracestate)
             return yield(nil, nil, nil) if tracestate.empty?
@@ -42,19 +67,26 @@
             rest = nil if rest.empty?
             yield(p, r, rest)
           end
 
           def update_tracestate(tracestate, p, r, rest) # rubocop:disable Naming/UncommunicativeMethodParamName
-            # This could be more efficient and allocate less, however it *should* only be used for root spans and sanitizing invalid tracestate
-            # in the most common configuration of parent_consistent_probability_based(root: consistent_probability_based(p)).
-            ps = "p:#{p}" unless p.nil?
-            rs = "r:#{r}" unless r.nil?
-            ot = [ps, rs, rest].compact.join(';')
-            if ot.empty?
+            if p.nil? && r.nil? && rest.nil?
               tracestate.delete('ot')
+            elsif p.nil? && r.nil?
+              tracestate.set_value('ot', rest)
+            elsif p.nil? && rest.nil?
+              tracestate.set_value('ot', "r:#{r}")
+            elsif r.nil? && rest.nil?
+              tracestate.set_value('ot', "p:#{p}")
+            elsif p.nil?
+              tracestate.set_value('ot', "r:#{r};#{rest}")
+            elsif r.nil?
+              tracestate.set_value('ot', "p:#{p};#{rest}")
+            elsif rest.nil?
+              tracestate.set_value('ot', "p:#{p};r:#{r}")
             else
-              tracestate.set_value('ot', ot)
+              tracestate.set_value('ot', "p:#{p};r:#{r};#{rest}")
             end
           end
 
           def new_tracestate(p: nil, r: nil) # rubocop:disable Naming/UncommunicativeMethodParamName
             if p.nil? && r.nil?
@@ -72,9 +104,15 @@
             ((p <= r) == sampled) || (sampled && (p == 63))
           end
 
           def decimal(str)
             str.to_i if !str.nil? && DECIMAL.match?(str)
+          end
+
+          def generate_r(trace_id)
+            x = trace_id[8, 8].unpack1('Q>') | 0x3
+            clz = 64 - x.bit_length
+            clz
           end
         end
       end
     end
   end