thrift_client-adamd 0.9.3

checksums.yaml

@@ -0,0 +1,15 @@
+---
+!binary "U0hBMQ==":
+  metadata.gz: !binary |-
+    OWViNmI3NWQ1NjRmYzZiMmY2MzliMTY5Y2EzNWUwZDY2MTJlYTQ2ZA==
+  data.tar.gz: !binary |-
+    YTZiNzZmOWRmYjI3OTVjYmRmMmU4YWI3NTdiNjNiZTFiZWEzZTFlNQ==
+SHA512:
+  metadata.gz: !binary |-
+    M2YxNjg2YzFhNWQ0MWYxNTFmMjNhOWUyOTQzNjliMDZmMWQyYmI2YjVkMzY3
+    MWRkZDRmMmIxZjM5NzYxOWYwMzE5ZWZjNTExYmFkNzVhMzZjNjdmYjBkMjk3
+    ZWQ2M2Q4YzlkMWVhMjkxNzgyNzY1NWRkZTYwOTJjMTExMTE0MWM=
+  data.tar.gz: !binary |-
+    OWViYThkYWRjOWJkNzBlZWJiOWM3NGJlMjc5ZjE4MThmNGE1ZGQ4ODNlYTA4
+    Y2QxYWYyM2E1MGNjOTUwZDU4NTRkOThkY2I4MTMzMTE5Y2FkZWRjN2JlNzA1
+    ZDQ0YjAyYmUyZTJjY2NmNmZiZTY1NzQ4Zjg0NmQ4YmI3OGExOTU=

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/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/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/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

data/lib/thrift_client/server.rb

@@ -0,0 +1,108 @@
+require 'thrift_client/connection'
+
+module ThriftHelpers
+  class Server
+    class ServerMarkedDown < StandardError; end
+
+    def initialize(connection_string, client_class, options = {})
+      @connection_string = connection_string
+      @client_class = client_class
+      @options = options
+
+      @cached = @options.has_key?(:cached_connections) ? @options[:cached_connections] : true
+
+      @marked_down_til = nil
+    end
+
+    def mark_down!(til)
+      close(true)
+      @marked_down_til = Time.now + til
+    end
+
+    def up?
+      !down?
+    end
+
+    def down?
+      @marked_down_til && @marked_down_til > Time.now
+    end
+
+    def to_s
+      @connection_string
+    end
+
+    def connection
+      @connection ||= Connection::Factory.create(
+        @options[:transport], @options[:transport_wrapper],
+        @connection_string, @options[:connect_timeout])
+    end
+
+    def connect!
+      return if open?
+
+      self.timeout = @options[:connect_timeout]
+      connection.connect!
+      self.timeout = @options[:timeout]
+    end
+
+    def client
+      @client ||= begin
+        connect!
+
+        @client_class.new(
+          @options[:protocol].new(self, *@options[:protocol_extra_params]))
+      end
+    end
+
+    def open?
+      connection.open?
+    end
+
+    def close(teardown = false)
+      if teardown || !@cached
+        connection.close if open?
+        @client = nil
+      end
+    end
+
+    def transport
+      connection.transport
+    end
+
+    module TransportInterface
+      def read(sz)
+        transport.read(sz)
+      end
+
+      def read_byte
+        transport.read_byte
+      end
+
+      def read_into_buffer(buffer, size)
+        transport.read_into_buffer(buffer, size)
+      end
+
+      def read_all(sz)
+        transport.read_all(sz)
+      end
+
+      def write(buf)
+        transport.write(buf)
+      end
+      alias_method :<<, :write
+
+      def flush
+        transport.flush
+      end
+
+      def timeout=(timeout)
+        transport.timeout = timeout if transport.respond_to?(:timeout=)
+      end
+
+      def timeout
+        transport.timeout
+      end
+    end
+    include TransportInterface
+  end
+end