finagle-thrift 1.2.0

data/lib/finagle-thrift.rb

@@ -0,0 +1,28 @@
+require 'rubygems'
+
+require 'thrift'
+
+require 'finagle-thrift/thrift/tracing_types'
+# require 'finagle-thrift/thrift/tracing_constants'
+
+require 'finagle-thrift/client'
+require 'finagle-thrift/thrift_client'
+require 'finagle-thrift/trace'
+require 'finagle-thrift/tracer'
+
+module FinagleThrift
+  extend self
+  def enable_tracing!(service, client_id = nil, service_name = nil)
+    raise ArgumentError, "client_id must be nil or of type FinagleThrift::ClientId" if client_id && !client_id.is_a?(FinagleThrift::ClientId)
+    class << service
+      include ::FinagleThrift::ThriftClient
+    end
+
+    client_class = service.client_class
+    client_class.class_eval do
+      include ::FinagleThrift::Client
+    end
+    client_class.send(:define_method, :client_id) { client_id }
+    client_class.send(:define_method, :trace_service_name) { service_name }
+  end
+end

data/lib/finagle-thrift/client.rb

@@ -0,0 +1,85 @@
+module FinagleThrift
+  module Client
+    CanTraceMethodName = "__can__finagle__trace__v3__"
+
+    include ::Thrift::Client
+
+    alias_method :_orig_send_message, :send_message
+    alias_method :_orig_receive_message, :receive_message
+
+    def initialize(iprot, oprot=nil)
+      super
+      @upgraded = false
+      @client_port = 0
+      attempt_upgrade!
+    end
+
+    def send_message(name, args_class, args = {})
+      if @upgraded
+        header = ::FinagleThrift::RequestHeader.new
+        header.trace_id = Trace.id.trace_id.to_i
+        header.parent_span_id = Trace.id.parent_id.to_i
+        header.span_id = Trace.id.span_id.to_i
+        header.sampled = Trace.id.sampled?
+        header.debug = false
+
+        header.client_id = client_id if client_id
+
+        header.write(@oprot)
+      end
+      Trace.record(Trace::Annotation.new(Trace::Annotation::CLIENT_SEND, self.endpoint))
+
+      _orig_send_message(name, args_class, args)
+    end
+
+    def receive_message(klass)
+      if @upgraded
+        response = ::FinagleThrift::ResponseHeader.new
+        response.read(@iprot)
+      end
+      result = _orig_receive_message(klass)
+      Trace.record(Trace::Annotation.new(Trace::Annotation::CLIENT_RECV, self.endpoint))
+      result
+    end
+
+    protected
+    def client_id
+      nil
+    end
+
+    def trace_service_name
+      nil
+    end
+
+    def endpoint
+      @endpoint ||= Trace.default_endpoint.with_port(@client_port).with_service_name(trace_service_name)
+    end
+
+    private
+    def attempt_upgrade!
+      _orig_send_message(CanTraceMethodName, ::FinagleThrift::ConnectionOptions)
+      begin
+        extract_client_port
+        _orig_receive_message(::FinagleThrift::UpgradeReply)
+        @upgraded = true
+      rescue ::Thrift::ApplicationException
+        @upgraded = false
+      end
+    rescue
+      @upgraded = false
+      raise
+    end
+
+    def extract_client_port
+      @client_port = begin
+        # im sorry
+        trans = @oprot.instance_variable_get("@trans")
+        transport = trans.instance_variable_get("@transport")
+        handle = transport.instance_variable_get("@handle")
+        Socket.unpack_sockaddr_in(handle.getsockname).first
+      rescue
+        0
+      end
+    end
+  end
+end

data/lib/finagle-thrift/thrift/tracing_constants.rb

@@ -0,0 +1,18 @@
+#
+# Autogenerated by Thrift
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+
+require 'tracing_types'
+
+  module FinagleThrift
+    CLIENT_SEND = %q"cs"
+
+    CLIENT_RECV = %q"cr"
+
+    SERVER_SEND = %q"ss"
+
+    SERVER_RECV = %q"sr"
+
+end

data/lib/finagle-thrift/thrift/tracing_types.rb

@@ -0,0 +1,210 @@
+#
+# Autogenerated by Thrift
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+
+
+module FinagleThrift
+    module AnnotationType
+      BOOL = 0
+      BYTES = 1
+      I16 = 2
+      I32 = 3
+      I64 = 4
+      DOUBLE = 5
+      STRING = 6
+      VALUE_MAP = {0 => "BOOL", 1 => "BYTES", 2 => "I16", 3 => "I32", 4 => "I64", 5 => "DOUBLE", 6 => "STRING"}
+      VALID_VALUES = Set.new([BOOL, BYTES, I16, I32, I64, DOUBLE, STRING]).freeze
+    end
+
+    class Endpoint
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      IPV4 = 1
+      PORT = 2
+      SERVICE_NAME = 3
+
+      FIELDS = {
+        IPV4 => {:type => ::Thrift::Types::I32, :name => 'ipv4'},
+        PORT => {:type => ::Thrift::Types::I16, :name => 'port'},
+        SERVICE_NAME => {:type => ::Thrift::Types::STRING, :name => 'service_name'}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class Annotation
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      TIMESTAMP = 1
+      VALUE = 2
+      HOST = 3
+
+      FIELDS = {
+        TIMESTAMP => {:type => ::Thrift::Types::I64, :name => 'timestamp'},
+        VALUE => {:type => ::Thrift::Types::STRING, :name => 'value'},
+        HOST => {:type => ::Thrift::Types::STRUCT, :name => 'host', :class => FinagleThrift::Endpoint, :optional => true}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class BinaryAnnotation
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      KEY = 1
+      VALUE = 2
+      ANNOTATION_TYPE = 3
+      HOST = 4
+
+      FIELDS = {
+        KEY => {:type => ::Thrift::Types::STRING, :name => 'key'},
+        VALUE => {:type => ::Thrift::Types::STRING, :name => 'value', :binary => true},
+        ANNOTATION_TYPE => {:type => ::Thrift::Types::I32, :name => 'annotation_type', :enum_class => FinagleThrift::AnnotationType},
+        HOST => {:type => ::Thrift::Types::STRUCT, :name => 'host', :class => FinagleThrift::Endpoint, :optional => true}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+        unless @annotation_type.nil? || FinagleThrift::AnnotationType::VALID_VALUES.include?(@annotation_type)
+          raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Invalid value of field annotation_type!')
+        end
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class Span
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      TRACE_ID = 1
+      NAME = 3
+      ID = 4
+      PARENT_ID = 5
+      ANNOTATIONS = 6
+      BINARY_ANNOTATIONS = 8
+
+      FIELDS = {
+        TRACE_ID => {:type => ::Thrift::Types::I64, :name => 'trace_id'},
+        NAME => {:type => ::Thrift::Types::STRING, :name => 'name'},
+        ID => {:type => ::Thrift::Types::I64, :name => 'id'},
+        PARENT_ID => {:type => ::Thrift::Types::I64, :name => 'parent_id', :optional => true},
+        ANNOTATIONS => {:type => ::Thrift::Types::LIST, :name => 'annotations', :element => {:type => ::Thrift::Types::STRUCT, :class => FinagleThrift::Annotation}},
+        BINARY_ANNOTATIONS => {:type => ::Thrift::Types::LIST, :name => 'binary_annotations', :element => {:type => ::Thrift::Types::STRUCT, :class => FinagleThrift::BinaryAnnotation}}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    # At connection time, we can let the server know who we are so
+    # they can book keep and optionally reject unknown clients.
+    class ClientId
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      NAME = 1
+
+      FIELDS = {
+        NAME => {:type => ::Thrift::Types::STRING, :name => 'name'}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    # RequestHeader defines headers for the request. These carry the span data, and
+    # a flag indicating whether the request is to be debugged.
+    class RequestHeader
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      TRACE_ID = 1
+      SPAN_ID = 2
+      PARENT_SPAN_ID = 3
+      DEBUG = 4
+      SAMPLED = 5
+      CLIENT_ID = 6
+
+      FIELDS = {
+        TRACE_ID => {:type => ::Thrift::Types::I64, :name => 'trace_id'},
+        SPAN_ID => {:type => ::Thrift::Types::I64, :name => 'span_id'},
+        PARENT_SPAN_ID => {:type => ::Thrift::Types::I64, :name => 'parent_span_id', :optional => true},
+        DEBUG => {:type => ::Thrift::Types::BOOL, :name => 'debug'},
+        SAMPLED => {:type => ::Thrift::Types::BOOL, :name => 'sampled', :optional => true},
+        CLIENT_ID => {:type => ::Thrift::Types::STRUCT, :name => 'client_id', :class => FinagleThrift::ClientId, :optional => true}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    # The Response carries a reply header for tracing. These are
+    # empty unless the request is being debugged, in which case a
+    # transcript is copied.
+    class ResponseHeader
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      SPANS = 1
+
+      FIELDS = {
+        SPANS => {:type => ::Thrift::Types::LIST, :name => 'spans', :element => {:type => ::Thrift::Types::STRUCT, :class => FinagleThrift::Span}}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    # These are connection-level options negotiated during protocol
+    # upgrade.
+    class ConnectionOptions
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+
+      FIELDS = {
+
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    # This is the struct that a successful upgrade will reply with.
+    class UpgradeReply
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+
+      FIELDS = {
+
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+  end

data/lib/finagle-thrift/thrift_client.rb

@@ -0,0 +1,10 @@
+module FinagleThrift
+  module ThriftClient
+    def handled_proxy(method_name, *args)
+      Trace.push(Trace.id.next_id) do
+        Trace.set_rpc_name(method_name)
+        super(method_name, *args)
+      end
+    end
+  end
+end

data/lib/finagle-thrift/trace.rb

@@ -0,0 +1,135 @@
+module Trace
+  extend self
+  DEFAULT_SAMPLE_RATE = 0.001
+  TRACE_ID_UPPER_BOUND = 2 ** 64
+
+  def id
+    if stack.empty?
+      span_id = generate_id
+      trace_id = TraceId.new(span_id, nil, span_id, should_sample?)
+      stack.push(trace_id)
+    end
+    stack.last
+  end
+
+  def push(trace_id)
+    stack.push(trace_id)
+    if block_given?
+      begin
+        yield
+      ensure
+        pop
+      end
+    end
+  end
+
+  def pop
+    stack.pop
+  end
+
+  def unwind
+    if block_given?
+      begin
+        saved_stack = stack.dup
+        yield
+      ensure
+        @stack = saved
+      end
+    end
+  end
+
+  def record(annotation)
+    tracer.record(id, annotation) unless stack.empty?
+  end
+
+  def set_rpc_name(name)
+    tracer.set_rpc_name(id, name) unless stack.empty?
+  end
+
+  def 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 tracer=(tracer)
+    @tracer = tracer
+  end
+
+  class TraceId
+    attr_reader :trace_id, :parent_id, :span_id, :sampled
+    alias :sampled? :sampled
+    def initialize(trace_id, parent_id, span_id, sampled)
+      @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
+    end
+
+    def next_id
+      TraceId.new(@trace_id, @span_id, Trace.generate_id, @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})"
+    end
+  end
+
+  class SpanId
+    HEX_REGEX = /^[a-f0-9]{16}$/i
+    MAX_SIGNED_I64 = 9223372036854775807
+    MASK = (2 ** 64) - 1
+
+    def self.from_value(v)
+      if v.is_a?(String) && v =~ HEX_REGEX
+        new(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
+
+  def generate_id
+    rand(TRACE_ID_UPPER_BOUND)
+  end
+
+  def should_sample?
+    rand < (@sample_rate || DEFAULT_SAMPLE_RATE)
+  end
+
+  def default_endpoint=(endpoint)
+    @default_endpoint = endpoint
+  end
+
+  def default_endpoint
+    @default_endpoint ||= begin
+      Endpoint.new(Endpoint.host_to_i32(Socket.gethostname), 0, "finagle-ruby")
+    end
+  end
+
+  private
+
+  def stack
+    @stack ||= []
+  end
+
+  def tracer
+    @tracer ||= NullTracer.new
+  end
+
+end

data/lib/finagle-thrift/tracer.rb

@@ -0,0 +1,245 @@
+module Trace
+  class Tracer
+    def record(id, annotation)
+      raise "not implemented"
+    end
+
+    def set_rpc_name(id, name)
+      raise "not implemented"
+    end
+  end
+
+  class NullTracer < Tracer
+    def record(id, annotation)
+    end
+
+    def set_rpc_name(id, name)
+    end
+  end
+
+  class FanoutTracer < Tracer
+    def initialize(tracers)
+      @tracers = tracers
+    end
+
+    def record(id, annotation)
+      @tracers.each { |tracer| tracer.record(id, annotation) }
+    end
+
+    def set_rpc_name(id, name)
+      @tracers.each { |tracer| tracer.set_rpc_name(id, name) }
+    end
+  end
+
+  class FileTracer < Tracer
+    def initialize
+      @last_trace_id = nil
+      @file = nil
+    end
+
+    def record(id, annotation)
+      return unless id.sampled?
+      file = get_file_for_id(id)
+      file.puts("#{id.to_s}: #{annotation.to_s}")
+      file.flush if (annotation.is_a?(Annotation) && annotation.value == Annotation::SERVER_SEND)
+    end
+
+    def set_rpc_name(id, name)
+      return unless id.sampled?
+      get_file_for_id(id).puts("#{id.to_s}: name = #{name}")
+    end
+
+    private
+    def get_file_for_id(id)
+      id = id.trace_id.to_i
+      if @last_trace_id != id
+        @last_trace_id = id
+        @file.close if @file
+        @file = File.open("/tmp/traces/#{id}.log", "a")
+      end
+      @file
+    end
+  end
+
+  class ZipkinTracer < Tracer
+    TRACER_CATEGORY = "zipkin"
+    def initialize(scribe, max_buffer)
+      @scribe = scribe
+      @max_buffer = max_buffer
+      reset
+    end
+
+    def record(id, annotation)
+      return unless id.sampled?
+      span = get_span_for_id(id)
+
+      case annotation
+      when BinaryAnnotation
+        span.binary_annotations << annotation
+      when Annotation
+        span.annotations << annotation
+      end
+
+      @count += 1
+      if @count >= @max_buffer || (annotation.is_a?(Annotation) && annotation.value == Annotation::SERVER_SEND)
+        flush!
+      end
+    end
+
+    def set_rpc_name(id, name)
+      return unless id.sampled?
+      span = get_span_for_id(id)
+      span.name = name.to_s
+    end
+
+    private
+    def get_span_for_id(id)
+      key = id.span_id.to_s
+      @spans[key] ||= begin
+        Span.new("", id)
+      end
+    end
+
+    def reset
+      @count = 0
+      @spans = {}
+    end
+
+    def flush!
+      @scribe.batch do
+        messages = @spans.values.map do |span|
+          buf = ''
+          trans = Thrift::MemoryBufferTransport.new(buf)
+          oprot = Thrift::BinaryProtocol.new(trans)
+          span.to_thrift.write(oprot)
+          binary = Base64.encode64(buf).gsub("\n", "")
+          @scribe.log(binary, TRACER_CATEGORY)
+        end
+      end
+      reset
+    end
+  end
+
+  class Span
+    attr_accessor :name, :annotations, :binary_annotations
+    def initialize(name, span_id)
+      @name = name
+      @span_id = span_id
+      @annotations = []
+      @binary_annotations = []
+    end
+
+    def to_thrift
+      FinagleThrift::Span.new(
+        :name => @name,
+        :trace_id => @span_id.trace_id.to_i,
+        :id => @span_id.span_id.to_i,
+        :parent_id => @span_id.parent_id.nil? ? nil : @span_id.parent_id.to_i,
+        :annotations => @annotations.map { |a| a.to_thrift },
+        :binary_annotations => @binary_annotations.map { |a| a.to_thrift }
+      )
+    end
+  end
+
+  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_thrift
+      thrift_host = host ? host.to_thrift : nil
+      FinagleThrift::Annotation.new(
+        :value => @value,
+        :host => thrift_host,
+        :timestamp => @timestamp
+      )
+    end
+
+    def to_s
+      "#{@value} at #{@timestamp} for (#{@host.to_s})"
+    end
+  end
+
+  class BinaryAnnotation
+    module Type
+      BOOL = "BOOL"
+      BYTES = "BYTES"
+      I16 = "I16"
+      I32 = "I32"
+      I64 = "I64"
+      DOUBLE = "DOUBLE"
+      STRING = "STRING"
+
+      def self.to_thrift(v)
+        FinagleThrift::AnnotationType::VALUE_MAP.index(v)
+      end
+    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_thrift
+      thrift_host = host ? host.to_thrift : nil
+      FinagleThrift::BinaryAnnotation.new(
+        :key => @key,
+        :value => @value,
+        :annotation_type => Type.to_thrift(@annotation_type),
+        :host => thrift_host
+      )
+    end
+
+  end
+
+  class Endpoint < Struct.new(:ipv4, :port, :service_name)
+    MAX_I32 = ((2 ** 31) - 1)
+    MASK = (2 ** 32) - 1
+
+    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 with_port(port)
+      Endpoint.new(self.ipv4, port, self.service_name)
+    end
+
+    def with_service_name(service_name)
+      Endpoint.new(self.ipv4, self.port, service_name)
+    end
+
+    def to_thrift
+      FinagleThrift::Endpoint.new(
+        :ipv4 => self.ipv4,
+        :port => self.port,
+        :service_name => self.service_name
+      )
+    end
+
+    def to_s
+      "#{service_name}@#{ipv4}:#{port}"
+    end
+  end
+
+end

data/lib/finagle-thrift/version.rb

@@ -0,0 +1,3 @@
+module FinagleThrift
+  VERSION = "1.2.0"
+end

metadata

@@ -0,0 +1,75 @@
+--- !ruby/object:Gem::Specification 
+name: finagle-thrift
+version: !ruby/object:Gem::Version 
+  hash: 31
+  prerelease: 
+  segments: 
+  - 1
+  - 2
+  - 0
+  version: 1.2.0
+platform: ruby
+authors: 
+- Arya Asemanfar
+autorequire: 
+bindir: bin
+cert_chain: []
+
+date: 2012-05-25 00:00:00 Z
+dependencies: []
+
+description: A Ruby client library for integrating into finagle's thrift tracing protocol
+email: 
+- arya@twitter.com
+executables: []
+
+extensions: []
+
+extra_rdoc_files: []
+
+files: 
+- lib/finagle-thrift/client.rb
+- lib/finagle-thrift/thrift/tracing_constants.rb
+- lib/finagle-thrift/thrift/tracing_types.rb
+- lib/finagle-thrift/thrift_client.rb
+- lib/finagle-thrift/trace.rb
+- lib/finagle-thrift/tracer.rb
+- lib/finagle-thrift/version.rb
+- lib/finagle-thrift.rb
+homepage: 
+licenses: []
+
+post_install_message: 
+rdoc_options: []
+
+require_paths: 
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement 
+  none: false
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      hash: 3
+      segments: 
+      - 0
+      version: "0"
+required_rubygems_version: !ruby/object:Gem::Requirement 
+  none: false
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      hash: 17
+      segments: 
+      - 1
+      - 3
+      - 5
+      version: 1.3.5
+requirements: []
+
+rubyforge_project: 
+rubygems_version: 1.8.15
+signing_key: 
+specification_version: 3
+summary: ""
+test_files: []
+