fl-thrift_client 0.4.2

data/lib/thrift_client/event_machine.rb

@@ -0,0 +1,145 @@
+raise RuntimeError, "The eventmachine transport requires Ruby 1.9.x" if RUBY_VERSION < '1.9.0'
+
+require 'eventmachine'
+require 'fiber'
+
+# EventMachine-ready Thrift connection
+# Should not be used with a transport wrapper since it already performs buffering in Ruby.
+module Thrift
+  class EventMachineTransport < BaseTransport
+    def initialize(host, port=9090, timeout=5)
+      @host, @port, @timeout = host, port, timeout
+      @connection = nil
+    end
+
+    def open?
+      @connection && @connection.connected?
+    end
+
+    def open
+      fiber = Fiber.current
+      @connection = EventMachineConnection.connect(@host, @port, @timeout)
+      @connection.callback do
+        fiber.resume
+      end
+      @connection.errback do
+        fiber.resume
+      end
+      Fiber.yield
+      @connection
+    end
+
+    def close
+      @connection.close
+    end
+
+    def read(sz)
+      @connection.blocking_read(sz)
+    end
+
+    def write(buf)
+      @connection.send_data(buf)
+    end
+  end
+
+  module EventMachineConnection
+    GARBAGE_BUFFER_SIZE = 4096 # 4kB
+
+    include EM::Deferrable
+
+    def self.connect(host='localhost', port=9090, timeout=5, &block)
+      EM.connect(host, port, self, host, port) do |conn|
+        conn.pending_connect_timeout = timeout
+      end
+    end
+
+    def trap
+      begin
+        yield
+      rescue Exception => ex
+        puts ex.message
+        puts ex.backtrace.join("\n")
+      end
+    end
+
+    def initialize(host, port=9090)
+      @host, @port = host, port
+      @index = 0
+      @connected = false
+      @buf = ''
+    end
+
+    def close
+      trap do
+        @connected = false
+        close_connection(true)
+      end
+    end
+
+    def blocking_read(size)
+      raise IOError, "lost connection to #{@host}:#{@port}" unless @connected
+      trap do
+        if can_read?(size)
+          yank(size)
+        else
+          raise ArgumentError, "Unexpected state" if @size or @callback
+
+          fiber = Fiber.current
+          @size = size
+          @callback = proc { |data|
+            fiber.resume(data)
+          }
+          Fiber.yield
+        end
+      end
+    end
+
+    def receive_data(data)
+      trap do
+        (@buf) << data
+
+        if @callback and can_read?(@size)
+          callback = @callback
+          data = yank(@size)
+          @callback = @size = nil
+          callback.call(data)
+        end
+      end
+    end
+
+    def connected?
+      @connected
+    end
+
+    def connection_completed
+      @connected = true
+      succeed
+    end
+
+    def unbind
+      if @connected
+        @connected = false
+      else
+        fail
+      end
+    end
+
+    def can_read?(size)
+      @buf.size >= @index + size
+    end
+
+    private
+
+    def yank(len)      
+      data = @buf.slice(@index, len)
+      @index += len
+      @index = @buf.size if @index > @buf.size
+      if @index >= GARBAGE_BUFFER_SIZE
+        @buf = @buf.slice(@index..-1)
+        @index = 0
+      end
+      data
+    end
+
+  end
+end

data/lib/thrift_client/simple.rb

@@ -0,0 +1,263 @@
+require 'socket'
+require 'getoptlong'
+
+class ThriftClient
+
+  # This is a simplified form of thrift, useful for clients only, and not
+  # making any attempt to have good performance. It's intended to be used by
+  # small command-line tools that don't want to install a dozen ruby files.
+  module Simple
+    VERSION_1 = 0x8001
+
+    # message types
+    CALL, REPLY, EXCEPTION = (1..3).to_a
+
+    # value types
+    STOP, VOID, BOOL, BYTE, DOUBLE, _, I16, _, I32, _, I64, STRING, STRUCT, MAP, SET, LIST = (0..15).to_a
+
+    FORMATS = {
+      BYTE => "c",
+      DOUBLE => "G",
+      I16 => "n",
+      I32 => "N",
+    }
+
+    SIZES = {
+      BYTE => 1,
+      DOUBLE => 8,
+      I16 => 2,
+      I32 => 4,
+    }
+
+    module ComplexType
+      module Extends
+        def type_id=(n)
+          @type_id = n
+        end
+
+        def type_id
+          @type_id
+        end
+      end
+
+      module Includes
+        def to_i
+          self.class.type_id
+        end
+
+        def to_s
+          args = self.values.map { |v| self.class.type_id == STRUCT ? v.name : v.to_s }.join(", ")
+          "#{self.class.name}.new(#{args})"
+        end
+      end
+    end
+
+    def self.make_type(type_id, name, *args)
+      klass = Struct.new("STT_#{name}", *args)
+      klass.send(:extend, ComplexType::Extends)
+      klass.send(:include, ComplexType::Includes)
+      klass.type_id = type_id
+      klass
+    end
+
+    ListType = make_type(LIST, "ListType", :element_type)
+    MapType = make_type(MAP, "MapType", :key_type, :value_type)
+    SetType = make_type(SET, "SetType", :element_type)
+    StructType = make_type(STRUCT, "StructType", :struct_class)
+
+    class << self
+      def pack_value(type, value)
+        case type
+        when BOOL
+          [ value ? 1 : 0 ].pack("c")
+        when STRING
+          [ value.size, value ].pack("Na*")
+        when I64
+          [ value >> 32, value & 0xffffffff ].pack("NN")
+        when ListType
+          [ type.element_type.to_i, value.size ].pack("cN") + value.map { |item| pack_value(type.element_type, item) }.join("")
+        when MapType
+          [ type.key_type.to_i, type.value_type.to_i, value.size ].pack("ccN") + value.map { |k, v| pack_value(type.key_type, k) + pack_value(type.value_type, v) }.join("")
+        when SetType
+          [ type.element_type.to_i, value.size ].pack("cN") + value.map { |item| pack_value(type.element_type, item) }.join("")
+        when StructType
+          value._pack
+        else
+          [ value ].pack(FORMATS[type])
+        end
+      end
+
+      def pack_request(method_name, arg_struct, request_id=0)
+        [ VERSION_1, CALL, method_name.to_s.size, method_name.to_s, request_id, arg_struct._pack ].pack("nnNa*Na*")
+      end
+
+      def read_value(s, type)
+        case type
+        when BOOL
+          s.read(1).unpack("c").first != 0
+        when STRING
+          len = s.read(4).unpack("N").first
+          s.read(len)
+        when I64
+          hi, lo = s.read(8).unpack("NN")
+          (hi << 32) | lo
+        when LIST
+          read_list(s)
+        when MAP
+          read_map(s)
+        when STRUCT
+          read_struct(s)
+        when ListType
+          read_list(s, type.element_type)
+        when MapType
+          read_map(s, type.key_type, type.value_type)
+        when StructType
+          read_struct(s, type.struct_class)
+        else
+          s.read(SIZES[type]).unpack(FORMATS[type]).first
+        end
+      end
+
+      def read_list(s, element_type=nil)
+        etype, len = s.read(5).unpack("cN")
+        expected_type = (element_type and element_type.to_i == etype.to_i) ? element_type : etype
+        rv = []
+        len.times do
+          rv << read_value(s, expected_type)
+        end
+        rv
+      end
+
+      def read_map(s, key_type=nil, value_type=nil)
+        ktype, vtype, len = s.read(6).unpack("ccN")
+        rv = {}
+        expected_key_type, expected_value_type = if key_type and value_type and key_type.to_i == ktype and value_type.to_i == vtype
+          [ key_type, value_type ]
+        else
+          [ ktype, vtype ]
+        end
+        len.times do
+          key = read_value(s, expected_key_type)
+          value = read_value(s, expected_value_type)
+          rv[key] = value
+        end
+        rv
+      end
+
+      def read_struct(s, struct_class=nil)
+        rv = struct_class.new()
+        while true
+          type = s.read(1).unpack("c").first
+          return rv if type == STOP
+          fid = s.read(2).unpack("n").first
+          field = struct_class ? struct_class._fields.find { |f| (f.fid == fid) and (f.type.to_i == type) } : nil
+          value = read_value(s, field ? field.type : type)
+          rv[field.name] = value if field
+        end
+      end
+
+      def read_response(s, rv_class)
+        version, message_type, method_name_len = s.read(8).unpack("nnN")
+        method_name = s.read(method_name_len)
+        seq_id = s.read(4).unpack("N").first
+        [ method_name, seq_id, read_struct(s, rv_class).rv ]
+      end
+    end
+
+    ## ----------------------------------------
+
+    class Field
+      attr_accessor :name, :type, :fid
+
+      def initialize(name, type, fid)
+        @name = name
+        @type = type
+        @fid = fid
+      end
+
+      def pack(value)
+        value.nil? ? "" : [ type.to_i, fid, ThriftClient::Simple.pack_value(type, value) ].pack("cna*")
+      end
+    end
+
+    class ThriftException < RuntimeError
+      def initialize(reason)
+        @reason = reason
+      end
+
+      def to_s
+        "ThriftException(#{@reason.inspect})"
+      end
+    end
+
+    module ThriftStruct
+      module Include
+        def _pack
+          self.class._fields.map { |f| f.pack(self[f.name]) }.join + [ STOP ].pack("c")
+        end
+      end
+
+      module Extend
+        def _fields
+          @fields
+        end
+
+        def _fields=(f)
+          @fields = f
+        end
+      end
+    end
+
+    def self.make_struct(name, *fields)
+      st_name = "ST_#{name}"
+      if Struct.constants.include?(st_name)
+        warn "#{caller[0]}: Struct::#{st_name} is already defined; returning original class."
+        Struct.const_get(st_name)
+      else
+        names = fields.map { |f| f.name.to_sym }
+        klass = Struct.new(st_name, *names)
+        klass.send(:include, ThriftStruct::Include)
+        klass.send(:extend, ThriftStruct::Extend)
+        klass._fields = fields
+        klass
+      end
+    end
+
+    class ThriftService
+      def initialize(sock)
+        @sock = sock
+      end
+
+      def self._arg_structs
+        @_arg_structs = {} if @_arg_structs.nil?
+        @_arg_structs
+      end
+
+      def self.thrift_method(name, rtype, *args)
+        arg_struct = ThriftClient::Simple.make_struct("Args__#{self.name}__#{name}", *args)
+        rv_struct = ThriftClient::Simple.make_struct("Retval__#{self.name}__#{name}", ThriftClient::Simple::Field.new(:rv, rtype, 0))
+        _arg_structs[name.to_sym] = [ arg_struct, rv_struct ]
+
+        arg_names = args.map { |a| a.name.to_s }.join(", ")
+        class_eval "def #{name}(#{arg_names}); _proxy(:#{name}#{args.size > 0 ? ', ' : ''}#{arg_names}); end"
+      end
+
+      def _proxy(method_name, *args)
+        cls = self.class.ancestors.find { |cls| cls.respond_to?(:_arg_structs) and cls._arg_structs[method_name.to_sym] }
+        arg_class, rv_class = cls._arg_structs[method_name.to_sym]
+        arg_struct = arg_class.new(*args)
+        @sock.write(ThriftClient::Simple.pack_request(method_name, arg_struct))
+        rv = ThriftClient::Simple.read_response(@sock, rv_class)
+        rv[2]
+      end
+
+      # convenience. robey is lazy.
+      [[ :field, "Field.new" ], [ :struct, "StructType.new" ],
+       [ :list, "ListType.new" ], [ :map, "MapType.new" ]].each do |new_name, old_name|
+        class_eval "def self.#{new_name}(*args); ThriftClient::Simple::#{old_name}(*args); end"
+      end
+
+      [ :void, :bool, :byte, :double, :i16, :i32, :i64, :string ].each { |sym| class_eval "def self.#{sym}; ThriftClient::Simple::#{sym.to_s.upcase}; end" }
+    end
+  end
+end

data/lib/thrift_client/thrift.rb

@@ -0,0 +1,21 @@
+module Thrift
+  class BufferedTransport
+    def timeout=(timeout)
+      @transport.timeout = timeout
+    end
+
+    def timeout
+      @transport.timeout
+    end
+  end
+
+  module Client
+    def timeout=(timeout)
+      @iprot.trans.timeout = timeout
+    end
+
+    def timeout
+      @iprot.trans.timeout
+    end
+  end
+end

data/lib/thrift_client.rb

@@ -0,0 +1,45 @@
+if ENV["ANCIENT_THRIFT"]
+  $LOAD_PATH.unshift("/Users/eweaver/p/twitter/rails/vendor/gems/thrift-751142/lib")
+  $LOAD_PATH.unshift("/Users/eweaver/p/twitter/rails/vendor/gems/thrift-751142/ext")
+  require 'thrift'
+else
+  require 'rubygems'
+  require 'thrift'
+end
+
+require 'rubygems'
+require 'thrift_client/thrift'
+require 'thrift_client/connection'
+require 'thrift_client/abstract_thrift_client'
+
+class ThriftClient < AbstractThriftClient
+  # This error is for backwards compatibility only. If defined in
+  # RetryingThriftClient instead, causes the test suite will break.
+  class NoServersAvailable < StandardError; end
+  include RetryingThriftClient
+  include TimingOutThriftClient
+
+=begin rdoc
+Create a new ThriftClient instance. Accepts an internal Thrift client class (such as CassandraRb::Client), a list of servers with ports, and optional parameters.
+
+Valid optional parameters are:
+
+<tt>:protocol</tt>:: Which Thrift protocol to use. Defaults to <tt>Thrift::BinaryProtocol</tt>.
+<tt>:protocol_extra_params</tt>:: An array of additional parameters to pass to the protocol initialization call. Defaults to <tt>[]</tt>.
+<tt>:transport</tt>:: Which Thrift transport to use. Defaults to <tt>Thrift::FramedTransport</tt>.
+<tt>:randomize_server_list</tt>:: Whether to connect to the servers randomly, instead of in order. Defaults to <tt>true</tt>.
+<tt>:exception_classes</tt>:: Which exceptions to catch and retry when sending a request. Defaults to <tt>[IOError, Thrift::Exception, Thrift::ProtocolException, Thrift::ApplicationException, Thrift::TransportException, NoServersAvailable]</tt>
+<tt>:raise</tt>:: Whether to reraise errors if no responsive servers are found. Defaults to <tt>true</tt>.
+<tt>:retries</tt>:: How many times to retry a request. Defaults to the number of servers defined.
+<tt>:server_retry_period</tt>:: How many seconds to wait before trying to reconnect after marking all servers as down. Defaults to <tt>1</tt>. Set to <tt>nil</tt> to retry endlessly.
+<tt>:server_max_requests</tt>:: How many requests to perform before moving on to the next server in the pool, regardless of error status. Defaults to <tt>nil</tt> (no limit).
+<tt>:timeout</tt>:: Specify the default timeout in seconds. Defaults to <tt>1</tt>.
+<tt>:timeout_overrides</tt>:: Specify additional timeouts on a per-method basis, in seconds. Only works with <tt>Thrift::BufferedTransport</tt>.
+<tt>:defaults</tt>:: Specify default values to return on a per-method basis, if <tt>:raise</tt> is set to false.
+
+=end rdoc
+
+  def initialize(client_class, servers, options = {})
+    super
+  end
+end

data/test/greeter/greeter.rb

@@ -0,0 +1,77 @@
+#
+# Autogenerated by Thrift
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+
+require 'thrift'
+
+module Greeter
+  class Client
+    include ::Thrift::Client
+
+    def greeting(name)
+      send_greeting(name)
+      return recv_greeting()
+    end
+
+    def send_greeting(name)
+      send_message('greeting', Greeting_args, :name => name)
+    end
+
+    def recv_greeting()
+      result = receive_message(Greeting_result)
+      return result.success unless result.success.nil?
+      raise ::Thrift::ApplicationException.new(::Thrift::ApplicationException::MISSING_RESULT, 'greeting failed: unknown result')
+    end
+
+  end
+
+  class Processor
+    include ::Thrift::Processor
+
+    def process_greeting(seqid, iprot, oprot)
+      args = read_args(iprot, Greeting_args)
+      result = Greeting_result.new()
+      result.success = @handler.greeting(args.name)
+      write_result(result, oprot, 'greeting', seqid)
+    end
+
+  end
+
+  # HELPER FUNCTIONS AND STRUCTURES
+
+  class Greeting_args
+    include ::Thrift::Struct
+    NAME = 1
+
+    ::Thrift::Struct.field_accessor self, :name
+    FIELDS = {
+      NAME => {:type => ::Thrift::Types::STRING, :name => 'name'}
+    }
+
+    def struct_fields; FIELDS; end
+
+    def validate
+    end
+
+  end
+
+  class Greeting_result
+    include ::Thrift::Struct
+    SUCCESS = 0
+
+    ::Thrift::Struct.field_accessor self, :success
+    FIELDS = {
+      SUCCESS => {:type => ::Thrift::Types::STRING, :name => 'success'}
+    }
+
+    def struct_fields; FIELDS; end
+
+    def validate
+    end
+
+  end
+
+end
+

data/test/greeter/greeter.thrift

@@ -0,0 +1,3 @@
+service Greeter {
+  string greeting(1:string name)
+}

data/test/greeter/server.rb

@@ -0,0 +1,40 @@
+module Greeter
+  class Handler
+    def greeting(name)
+      "hello there #{name}!"
+    end
+  end
+
+  class Server
+    def initialize(port)
+      @port = port
+      handler = Greeter::Handler.new
+      processor = Greeter::Processor.new(handler)
+      transport = Thrift::ServerSocket.new("127.0.0.1", port)
+      transportFactory = Thrift::FramedTransportFactory.new()
+      @server = Thrift::SimpleServer.new(processor, transport, transportFactory)
+    end
+
+    def serve
+      @server.serve()
+    end
+  end
+
+  # client:
+  # trans = Thrift::HTTPClientTransport.new("http://127.0.0.1:9292/greeter")
+  # prot = Thrift::BinaryProtocol.new(trans)
+  # c = Greeter::Client.new(prot)
+  class HTTPServer
+    def initialize(uri)
+      uri = URI.parse(uri)
+      handler = Greeter::Handler.new
+      processor = Greeter::Processor.new(handler)
+      path = uri.path[1..-1]
+      @server = Thrift::MongrelHTTPServer.new(processor, :port => uri.port, :ip => uri.host, :path => path)
+    end
+
+    def serve
+      @server.serve()
+    end
+  end
+end

data/test/multiple_working_servers_test.rb

@@ -0,0 +1,112 @@
+require "#{File.dirname(__FILE__)}/test_helper"
+
+class MultipleWorkingServersTest < Test::Unit::TestCase
+  def setup
+    @servers = ["127.0.0.1:1461", "127.0.0.1:1462", "127.0.0.1:1463"]
+    @socket = 1461
+    @timeout = 0.2
+    @options = {:protocol_extra_params => [false]}
+    @pids = []
+    @servers.each do |s|
+      @pids << Process.fork do
+        Signal.trap("INT") { exit }
+        Greeter::Server.new(s.split(':').last).serve
+      end
+    end
+    # Need to give the child process a moment to open the listening socket or
+    # we get occasional "could not connect" errors in tests.
+    sleep 0.05
+  end
+
+  def teardown
+    @pids.each do |pid|
+      Process.kill("INT", pid)
+      Process.wait(pid)
+    end
+  end
+
+  def test_server_creates_new_client_that_can_talk_to_all_servers_after_disconnect
+    client = ThriftClient.new(Greeter::Client, @servers, @options)
+    client.greeting("someone")
+    internal_client = client.client
+    client.greeting("someone")
+    assert_equal internal_client, client.client # Sanity check
+
+    client.disconnect!
+    client.greeting("someone")
+    internal_client = client.client
+    client.greeting("someone")
+    assert_equal internal_client, client.client
+    internal_client = client.client
+    client.greeting("someone")
+    assert_equal internal_client, client.client
+
+    # Moves on to the second server
+    assert_nothing_raised {
+      client.greeting("someone")
+      client.greeting("someone")
+    }
+  end
+
+  def test_server_doesnt_max_out_after_explicit_disconnect
+    client = ThriftClient.new(Greeter::Client, @servers, @options.merge(:server_max_requests => 2))
+    client.greeting("someone")
+    internal_client = client.client
+    client.greeting("someone")
+    assert_equal internal_client, client.client # Sanity check
+
+    client.disconnect!
+
+    client.greeting("someone")
+    internal_client = client.client
+    client.greeting("someone")
+    assert_equal internal_client, client.client, "ThriftClient should not have reset the internal client if the counter was reset on disconnect"
+  end
+
+  def test_server_disconnect_doesnt_drop_servers_with_retry_period
+    client = ThriftClient.new(Greeter::Client, @servers, @options.merge(:server_max_requests => 2, :retry_period => 1))
+    3.times {
+      client.greeting("someone")
+      internal_client = client.client
+      client.greeting("someone")
+      assert_equal internal_client, client.client # Sanity check
+
+      client.disconnect!
+
+      client.greeting("someone")
+      internal_client = client.client
+      client.greeting("someone")
+      assert_equal internal_client, client.client, "ThriftClient should not have reset the internal client if the counter was reset on disconnect"
+    }
+  end
+
+
+  def test_server_max_requests
+    client = ThriftClient.new(Greeter::Client, @servers, @options.merge(:server_max_requests => 2))
+
+    client.greeting("someone")
+    internal_client = client.client
+
+    client.greeting("someone")
+    assert_equal internal_client, client.client
+
+    # This next call maxes out the requests for that "client" object
+    # and moves on to the next.
+    client.greeting("someone")
+    assert_not_equal internal_client, new_client = client.client
+
+    # And here we should still have the same client as the last one...
+    client.greeting("someone")
+    assert_equal new_client, client.client
+
+    # Until we max it out, too.
+    client.greeting("someone")
+    assert_not_equal new_client, client.client
+    assert_not_nil client.client
+
+    new_new_client = client.client
+    # And we should still have one server left
+    client.greeting("someone")
+    assert_equal new_new_client, client.client
+  end
+end