RaymondChou-thrift_client 0.9.3

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: 9e16ff9dcf19cf0dff6b0fc0c760fdfaeee7560b
+  data.tar.gz: b79f3d1cae0c4fb18727988b6e42f82fdf39e260
+SHA512:
+  metadata.gz: 5a2895690e6da4f99c1f18670fcda05b55c243a357848a99cb2d629bfeb8c5352332ecdbea2a9933224ef4e6d6d468f7cb08a530c842d8e3f63d73346ce6a084
+  data.tar.gz: 82e6c2a65ab738c07b0e2c2e9982fb9852ba8c872460a7446c634113f627f932d4346d59e3beeb505f3d6570bc2189bbee1dd6383e126b36ac57ce405b5d12ef

data/lib/thrift_client.rb

@@ -0,0 +1,35 @@
+require 'thrift'
+require 'thrift_client/thrift'
+require 'thrift_client/server'
+require 'thrift_client/abstract_thrift_client'
+
+class ThriftClient < AbstractThriftClient
+  class NoServersAvailable < StandardError; end
+
+=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::Socket</tt>.
+<tt>:transport_wrapper</tt>:: Which Thrift transport wrapper to use. Defaults to <tt>Thrift::FramedTransport</tt>.
+<tt>:exception_classes</tt>:: Which exceptions to catch and retry when sending a request. Defaults to <tt>[IOError, Thrift::Exception, Thrift::ApplicationException, Thrift::TransportException, NoServersAvailable]</tt>
+<tt>:exception_class_overrides</tt>:: For specifying children of classes in exception_classes for which you don't want to retry or reconnect.
+<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 0.
+<tt>:server_retry_period</tt>:: How many seconds to wait before trying to reconnect to a dead server. Defaults to <tt>1</tt>. Set to <tt>nil</tt> to disable.
+<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>:connect_timeout</tt>:: Specify the connection timeout in seconds. Defaults to <tt>0.1</tt>.
+<tt>:timeout_overrides</tt>:: Specify additional timeouts on a per-method basis, in seconds. Only works with <tt>Thrift::BufferedTransport</tt>.
+<tt>:cached_connections</tt>:: Cache connections between requests. Trades connect() costs for open sockets. Defaults to <tt>false</tt>.
+<tt>:defaults</tt>:: Specify default values to return on a per-method basis, if <tt>:raise</tt> is set to false.
+<tt>:multiplexed_protocol</tt>:: Thrift::MultiplexedProtocol.new(protocol, 'WechatService')
+
+=end rdoc
+  def initialize(client_class, servers, options = {})
+    super
+  end
+end

data/lib/thrift_client/abstract_thrift_client.rb

@@ -0,0 +1,226 @@
+class AbstractThriftClient
+  include ThriftHelpers
+
+  DISCONNECT_ERRORS = [
+    IOError,
+    Thrift::Exception,
+    Thrift::ApplicationException,
+    Thrift::TransportException
+  ]
+
+  DEFAULT_WRAPPED_ERRORS = [
+    Thrift::ApplicationException,
+    Thrift::TransportException,
+  ]
+
+  DEFAULTS = {
+    :protocol => Thrift::BinaryProtocol,
+    :protocol_extra_params => [],
+    :transport => Thrift::Socket,
+    :transport_wrapper => Thrift::FramedTransport,
+    :raise => true,
+    :defaults => {},
+    :exception_classes => DISCONNECT_ERRORS,
+    :exception_class_overrides => [],
+    :retries => 0,
+    :server_retry_period => 1,
+    :server_max_requests => nil,
+    :retry_overrides => {},
+    :wrapped_exception_classes => DEFAULT_WRAPPED_ERRORS,
+    :connect_timeout => 0.1,
+    :timeout => 1,
+    :timeout_overrides => {},
+    :cached_connections => false
+  }
+
+  attr_reader :last_client, :client, :client_class, :current_server, :server_list, :options, :client_methods
+
+  def self.create_wrapped_exception_classes(client_class, wrapped_exception_classes = DEFAULT_WRAPPED_ERRORS)
+    wrapped_exception_classes.map do |exception_klass|
+      name = exception_klass.to_s.split('::').last
+      begin
+        client_class.const_get(name)
+      rescue NameError
+        client_class.const_set(name, Class.new(exception_klass))
+      end
+    end
+  end
+
+  def initialize(client_class, servers, options = {})
+    @options = DEFAULTS.merge(options)
+    @options[:server_retry_period] ||= 0
+
+    @client_class = client_class
+    @server_list = Array(servers).collect do |s|
+      Server.new(s, @client_class, @options)
+    end.sort_by { rand }
+
+    @current_server = @server_list.first
+
+    @callbacks = {}
+    @client_methods = []
+    @client_class.instance_methods.each do |method_name|
+      if method_name != 'send_message' && method_name =~ /^send_(.*)$/
+        instance_eval("def #{$1}(*args); handled_proxy(:'#{$1}', *args); end", __FILE__, __LINE__)
+        @client_methods << $1
+      end
+    end
+    @request_count = 0
+    self.class.create_wrapped_exception_classes(@client_class, @options[:wrapped_exception_classes])
+  end
+
+  # Adds a callback that will be invoked at a certain time. The valid callback types are:
+  #   :post_connect  - should accept a single AbstractThriftClient argument, which is the client object to
+  #                    which the callback was added. Called after a connection to the remote thrift server
+  #                    is established.
+  #   :before_method - should accept a single method name argument. Called before a method is invoked on the
+  #                    thrift server.
+  #   :on_exception  - should accept 2 args: an Exception instance and a method name. Called right before the
+  #                    exception is raised.
+  def add_callback(callback_type, &block)
+    case callback_type
+    when :post_connect, :before_method, :on_exception
+      @callbacks[callback_type] ||= []
+      @callbacks[callback_type].push(block)
+      # Allow chaining
+      return self
+    else
+      return nil
+    end
+  end
+
+  def inspect
+    "<#{self.class}(#{client_class}) @current_server=#{@current_server}>"
+  end
+
+  # Force the client to connect to the server. Not necessary to be
+  # called as the connection will be made on the first RPC method
+  # call.
+  def connect!(method = nil)
+    start_time ||= Time.now
+    @current_server = next_live_server
+    @client = @current_server.client
+    @last_client = @client
+    do_callbacks(:post_connect, self)
+  rescue IOError, Thrift::TransportException
+    disconnect!(true)
+    timeout = timeout(method)
+    if timeout && Time.now - start_time > timeout
+      no_servers_available!
+    else
+      retry
+    end
+  end
+
+  def disconnect!(error = false)
+    if @current_server
+      @current_server.mark_down!(@options[:server_retry_period]) if error
+      @current_server.close
+    end
+
+    @client = nil
+    @current_server = nil
+    @request_count = 0
+  end
+
+  private
+
+  # Calls all callbacks of the specified type with the given args
+  def do_callbacks(callback_type_sym, *args)
+    return unless @callbacks[callback_type_sym]
+    @callbacks[callback_type_sym].each do |callback|
+      callback.call(*args)
+    end
+  end
+
+  def next_live_server
+    @server_index ||= 0
+    @server_list.length.times do |i|
+      cur = (1 + @server_index + i) % @server_list.length
+      if @server_list[cur].up?
+        @server_index = cur
+        return @server_list[cur]
+      end
+    end
+    no_servers_available!
+  end
+
+  def ensure_socket_alignment
+    incomplete = true
+    result = yield
+    incomplete = false
+    result
+  # Thrift exceptions get read off the wire. We can consider them complete requests
+  rescue Thrift::Exception => e
+    incomplete = false
+    raise e
+  ensure
+    disconnect! if incomplete
+  end
+
+  def handled_proxy(method_name, *args)
+    begin
+      connect!(method_name.to_sym) unless @client
+      if has_timeouts?
+        @client.timeout = timeout(method_name.to_sym)
+      end
+      @request_count += 1
+      do_callbacks(:before_method, method_name)
+      ensure_socket_alignment { @client.send(method_name, *args) }
+    rescue *@options[:exception_class_overrides] => e
+      raise_or_default(e, method_name)
+    rescue *@options[:exception_classes] => e
+      disconnect!(true)
+      tries ||= (@options[:retry_overrides][method_name.to_sym] || @options[:retries]) + 1
+      tries -= 1
+      if tries > 0
+        retry
+      else
+        raise_or_default(e, method_name)
+      end
+    rescue Exception => e
+      raise_or_default(e, method_name)
+    ensure
+      disconnect! if @options[:server_max_requests] && @request_count >= @options[:server_max_requests]
+    end
+  end
+
+  def raise_or_default(e, method_name)
+    if @options[:raise]
+      raise_wrapped_error(e, method_name)
+    else
+      @options[:defaults][method_name.to_sym]
+    end
+  end
+
+  def raise_wrapped_error(e, method_name)
+    do_callbacks(:on_exception, e, method_name)
+    if @options[:wrapped_exception_classes].include?(e.class)
+      raise @client_class.const_get(e.class.to_s.split('::').last), e.message, e.backtrace
+    else
+      raise e
+    end
+  end
+
+  def has_timeouts?
+    @has_timeouts ||= @options[:timeout_overrides].any? && transport_can_timeout?
+  end
+
+  def timeout(method = nil)
+    @options[:timeout_overrides][method] || @options[:timeout]
+  end
+
+  def transport_can_timeout?
+    if (@options[:transport_wrapper] || @options[:transport]).method_defined?(:timeout=)
+      true
+    else
+      warn "ThriftClient: Timeout overrides have no effect with with transport type #{(@options[:transport_wrapper] || @options[:transport])}"
+      false
+    end
+  end
+
+  def no_servers_available!
+    servers = @server_list.map { |s| s.to_s }.join(',')
+    raise ThriftClient::NoServersAvailable, "No live servers in [#{servers}]."
+  end
+end

data/lib/thrift_client/connection.rb

@@ -0,0 +1,4 @@
+require "thrift_client/connection/base"
+require "thrift_client/connection/socket"
+require "thrift_client/connection/http"
+require "thrift_client/connection/factory"

data/lib/thrift_client/connection/base.rb

@@ -0,0 +1,25 @@
+module ThriftHelpers
+  module Connection
+    class Base
+      attr_accessor :transport, :server
+
+      def initialize(transport, transport_wrapper, server, timeout)
+        @transport = transport
+        @transport_wrapper = transport_wrapper
+        @server = server
+        @timeout = timeout
+      end
+
+      def connect!
+        raise NotImplementedError
+      end
+
+      def open?
+        @transport.open?
+      end
+
+      def close
+      end
+    end
+  end
+end

data/lib/thrift_client/connection/factory.rb

@@ -0,0 +1,13 @@
+module ThriftHelpers
+  module Connection
+    class Factory
+      def self.create(transport, transport_wrapper, server, timeout)
+        if transport == Thrift::HTTPClientTransport
+          Connection::HTTP.new(transport, transport_wrapper, server, timeout)
+        else
+          Connection::Socket.new(transport, transport_wrapper, server, timeout)
+        end
+      end
+    end
+  end
+end

data/lib/thrift_client/connection/http.rb

@@ -0,0 +1,29 @@
+module ThriftHelpers
+  module Connection
+    class HTTP < Base
+      def initialize(*args)
+        super *args
+
+        uri = parse_server(@server)
+        @transport = Thrift::HTTPClientTransport.new(@server)
+      end
+
+      def connect!
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = uri.scheme == "https"
+        http.get(uri.path)
+      end
+
+      def open?
+        true
+      end
+
+      private
+      def parse_server(server)
+        uri = URI.parse(server)
+        raise ArgumentError, 'Servers must start with http' unless uri.scheme =~ /^http/
+        uri
+      end
+    end
+  end
+end

data/lib/thrift_client/connection/socket.rb

@@ -0,0 +1,29 @@
+module ThriftHelpers
+  module Connection
+    class Socket < Base
+      def initialize(*args)
+        super *args
+
+        host, port = parse_server(@server)
+        @transport = @transport.new(host, port.to_i, @timeout)
+        @transport = @transport_wrapper.new(@transport) if @transport_wrapper
+      end
+
+      def close
+        @transport.close
+      end
+
+      def connect!
+        @transport.open
+      end
+
+      private
+
+      def parse_server(server)
+        host, port = server.to_s.split(":")
+        raise ArgumentError, 'Servers must be in the form "host:port"' unless host and port
+        [host, port]
+      end
+    end
+  end
+end

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
+
+      raise Thrift::TransportException, "Unable to connect to #{@host}:#{@port}" unless @connection.connected?
+      @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
+      @disconnected = 'not connected'
+      @buf = ''
+    end
+
+    def close
+      trap do
+        @disconnected = 'closed'
+        close_connection(true)
+      end
+    end
+
+    def blocking_read(size)
+      raise IOError, "lost connection to #{@host}:#{@port}: #{@disconnected}" if @disconnected
+      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
+
+    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?
+      !@disconnected
+    end
+
+    def connection_completed
+      @disconnected = nil
+      succeed
+    end
+
+    def unbind
+      if !@disconnected
+        @disconnected = 'unbound'
+      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