trace.rb

require 'zipkin-tracer/zipkin_tracer_base'
require 'zipkin-tracer/trace_container'
# Most of this code is copied from Finagle
# https://github.com/twitter/finagle/blob/finagle-6.39.0/finagle-thrift/src/main/ruby/lib/finagle-thrift/trace.rb
# But moved and improved here.
module Trace
  # These methods and attr_accessor below are used as global configuration of this gem
  # Most of these are set by the config class and then used around.
  # TODO: Move this out of the Trace module , take out that extend self and be happier
  extend self
  attr_accessor :trace_id_128bit

  # This method is deprecated, please use TraceGenerator.current
  # Note that this method will always return a trace, it will
  # generate a new one if none was available.
  def id
    ZipkinTracer::TraceGenerator.new.current
  end

  def self.tracer
    @tracer
  end

  def self.sample_rate
    @sample_rate
  end

  def self.tracer=(tracer)
    @tracer = tracer
  end

  def self.sample_rate=(sample_rate)
    if sample_rate > 1 || sample_rate < 0
      raise ArgumentError.new("sample rate must be [0,1]")
    end
    @sample_rate = sample_rate
  end

  def default_endpoint=(endpoint)
    @default_endpoint = endpoint
  end

  def default_endpoint
    @default_endpoint
  end

  # These classes all come from Finagle-thrift + some needed modifications (.to_h)
  # Moved here as a first step, eventually move them out of the Trace module

  class Annotation
    CLIENT_SEND = "cs"
    CLIENT_RECV = "cr"
    SERVER_SEND = "ss"
    SERVER_RECV = "sr"

    attr_reader :value, :host, :timestamp
    def initialize(value, host)
      @timestamp = (Time.now.to_f * 1000 * 1000).to_i # micros
      @value = value
      @host = host
    end

    def to_h
      {
        value: @value,
        timestamp: @timestamp,
        endpoint: host.to_h
      }
    end
  end

  class BinaryAnnotation
    SERVER_ADDRESS = 'sa'.freeze
    URI = 'http.url'.freeze
    METHOD = 'http.method'.freeze
    PATH = 'http.path'.freeze
    STATUS = 'http.status'.freeze
    LOCAL_COMPONENT = 'lc'.freeze
    ERROR = 'error'.freeze

    module Type
      BOOL = "BOOL"
      STRING = "STRING"
    end
    attr_reader :key, :value, :host

    def initialize(key, value, annotation_type, host)
      @key = key
      @value = value
      @annotation_type = annotation_type
      @host = host
    end

    def to_h
      {
        key: @key,
        value: @value,
        endpoint: host.to_h
      }
    end
  end

  class Flags
    # no flags set
    EMPTY = 0
    # the debug flag is used to ensure we pass all the sampling layers and that the trace is stored
    DEBUG = 1
  end

  class SpanId
    HEX_REGEX = /^[a-f0-9]{16,32}$/i
    MAX_SIGNED_I64 = 9223372036854775807
    MASK = (2 ** 64) - 1

    def self.from_value(v)
      if v.is_a?(String) && v =~ HEX_REGEX
        # drops any bits higher than 64 by selecting right-most 16 characters
        new(v.length > 16 ? v[v.length - 16, 16].hex : v.hex)
      elsif v.is_a?(Numeric)
        new(v)
      elsif v.is_a?(SpanId)
        v
      end
    end

    def initialize(value)
      @value = value
      @i64 = if @value > MAX_SIGNED_I64
        -1 * ((@value ^ MASK) + 1)
      else
        @value
      end
    end

    def to_s; "%016x" % @value; end
    def to_i; @i64; end
  end

  # A TraceId contains all the information of a given trace id
  class TraceId
    attr_reader :trace_id, :parent_id, :span_id, :sampled, :flags

    def initialize(trace_id, parent_id, span_id, sampled, flags)
      @trace_id = Trace.trace_id_128bit ? TraceId128Bit.from_value(trace_id) : SpanId.from_value(trace_id)
      @parent_id = parent_id.nil? ? nil : SpanId.from_value(parent_id)
      @span_id = SpanId.from_value(span_id)
      @sampled = sampled
      @flags = flags
    end

    def next_id
      TraceId.new(@trace_id, @span_id, ZipkinTracer::TraceGenerator.new.generate_id, @sampled, @flags)
    end

    # the debug flag is used to ensure the trace passes ALL samplers
    def debug?
      @flags & Flags::DEBUG == Flags::DEBUG
    end

    def sampled?
      debug? || ['1', 'true'].include?(@sampled)
    end

    def to_s
      "TraceId(trace_id = #{@trace_id.to_s}, parent_id = #{@parent_id.to_s}, span_id = #{@span_id.to_s}, sampled = #{@sampled.to_s}, flags = #{@flags.to_s})"
    end
  end

  class TraceId128Bit < SpanId
    HEX_REGEX_16 = /^[a-f0-9]{16}$/i
    HEX_REGEX_32 = /^[a-f0-9]{32}$/i
    MAX_SIGNED_I128 = (2 ** 128 / 2) -1
    MASK = (2 ** 128) - 1

    def self.from_value(v)
      if v.is_a?(String) && v =~ HEX_REGEX_16
        SpanId.new(v.hex)
      elsif v.is_a?(String) && v =~ HEX_REGEX_32
        new(v.hex)
      elsif v.is_a?(Numeric)
        new(v)
      elsif v.is_a?(SpanId)
        v
      end
    end

    def initialize(value)
      @value = value
      @i128 = if @value > MAX_SIGNED_I128
        -1 * ((@value ^ MASK) + 1)
      else
        @value
      end
    end

    def to_s; '%032x' % @value; end
    def to_i; @i128; end
  end

  # A span may contain many annotations
  class Span
    attr_accessor :name, :annotations, :binary_annotations, :debug
    def initialize(name, span_id)
      @name = name
      @span_id = span_id
      @annotations = []
      @binary_annotations = []
      @debug = span_id.debug?
      @timestamp = to_microseconds(Time.now)
      @duration = UNKNOWN_DURATION
    end

    def close
      @duration = to_microseconds(Time.now) - @timestamp
    end

    def to_h
      h = {
        name: @name,
        traceId: @span_id.trace_id.to_s,
        id: @span_id.span_id.to_s,
        annotations: @annotations.map(&:to_h),
        binaryAnnotations: @binary_annotations.map(&:to_h),
        timestamp: @timestamp,
        duration: @duration,
        debug: @debug
      }
      h[:parentId] = @span_id.parent_id.to_s unless @span_id.parent_id.nil?
      h
    end

    # We record information into spans, then we send these spans to zipkin
    def record(value, endpoint = Trace.default_endpoint)
      annotations << Trace::Annotation.new(value.to_s, endpoint)
    end

    def record_tag(key, value, type = Trace::BinaryAnnotation::Type::STRING, endpoint = Trace.default_endpoint)
      value = value.to_s if type == Trace::BinaryAnnotation::Type::STRING
      binary_annotations << Trace::BinaryAnnotation.new(key, value, type, endpoint)
    end

    def record_local_component(value)
      record_tag(BinaryAnnotation::LOCAL_COMPONENT, value)
    end

    def has_parent_span?
      !@span_id.parent_id.nil?
    end

    private

    UNKNOWN_DURATION = 0 # mark duration was not set

    def to_microseconds(time)
      (time.to_f * 1_000_000).to_i
    end
  end

  class Endpoint < Struct.new(:ipv4, :port, :service_name, :ip_format)
    MAX_I32 = ((2 ** 31) - 1)
    MASK = (2 ** 32) - 1
    UNKNOWN_URL = 'unknown'.freeze

    def self.host_to_i32(host)
      unsigned_i32 = Socket.getaddrinfo(host, nil)[0][3].split(".").map do |i|
        i.to_i
      end.inject(0) { |a,e| (a << 8) + e }

      signed_i32 = if unsigned_i32 > MAX_I32
        -1 * ((unsigned_i32 ^ MASK) + 1)
      else
        unsigned_i32
      end

      signed_i32
    end

    def self.local_endpoint(service_name, ip_format)
      hostname = Socket.gethostname
      Endpoint.new(hostname, nil, service_name, ip_format)
    end

    def self.remote_endpoint(url, remote_service_name, ip_format)
      service_name = remote_service_name || url.host.split('.').first || UNKNOWN_URL # default to url-derived service name
      Endpoint.new(url.host, url.port, service_name, ip_format)
    end

    def to_h
      hsh = {
        ipv4: ipv4,
        serviceName: service_name
      }
      hsh[:port] = port if port
      hsh
    end

  end
end