em-voldemort 0.1.5

data/lib/em-voldemort/cluster.rb

@@ -0,0 +1,247 @@
+module EM::Voldemort
+  # A client for a Voldemort cluster. The cluster is initialized by giving the hostname and port of
+  # one of its nodes, and the other nodes are autodiscovered.
+  #
+  # TODO if one node is down, a request should be retried on a replica.
+  class Cluster
+    include Protocol
+
+    attr_reader :bootstrap_host, :bootstrap_port, :logger, :cluster_name
+
+    RETRY_BOOTSTRAP_PERIOD = 10 # seconds
+    METADATA_STORE_NAME = 'metadata'.freeze
+    CLUSTER_INFO_KEY = 'cluster.xml'.freeze
+    STORES_INFO_KEY = 'stores.xml'.freeze
+
+    def initialize(options={})
+      @bootstrap_host = options[:host] or raise "#{self.class.name} requires :host"
+      @bootstrap_port = options[:port] or raise "#{self.class.name} requires :port"
+      @logger = options[:logger] || Logger.new($stdout)
+      @bootstrap_state = :not_started
+      @stores = {}
+    end
+
+    # Bootstraps the cluster (discovers all cluster nodes and metadata by connecting to the one node
+    # that was specified). Calling #connect is optional, since it also happens automatically when
+    # you start making requests.
+    def connect
+      @bootstrap_timer ||= setup_bootstrap_timer do
+        start_bootstrap if @bootstrap_state == :not_started || @bootstrap_state == :failed
+      end
+      start_bootstrap if @bootstrap_state == :not_started
+      @bootstrap
+    end
+
+    # Fetches the value associated with a particular key in a particular store. Returns a deferrable
+    # that succeeds with the value, or fails with an exception object.
+    def get(store_name, key, router=nil)
+      when_ready do |deferrable|
+        connections = choose_connections(key, router)
+        if connections.size == 0
+          deferrable.fail(ServerError.new('No connection can handle the request'))
+        elsif connections.first.health == :good
+          # Send the request to the preferred node, but fall back to others if it fails.
+          get_in_sequence(connections, store_name, key, deferrable)
+        else
+          # The request to the first node will probably fail, but we send it anyway, so that the
+          # connection can discover when the node comes back up. Make the request to other
+          # connections at the same time to avoid delaying the request.
+          get_in_parallel(connections, store_name, key, deferrable)
+        end
+      end
+    end
+
+    # Returns a {Store} object configured for accessing a particular store on the cluster.
+    def store(store_name)
+      @stores[store_name.to_s] ||= Store.new(self, store_name)
+    end
+
+    private
+
+    def setup_bootstrap_timer
+      EM.add_periodic_timer(RETRY_BOOTSTRAP_PERIOD) { yield }
+    end
+
+    def start_bootstrap
+      @bootstrap_state = :started
+      @bootstrap_conn = Connection.new(:host => bootstrap_host, :port => bootstrap_port, :logger => logger)
+      @bootstrap = EM::DefaultDeferrable.new
+
+      cluster_req = get_from_connection(@bootstrap_conn, METADATA_STORE_NAME, CLUSTER_INFO_KEY)
+
+      cluster_req.callback do |cluster_xml|
+        if parse_cluster_info(cluster_xml) == :cluster_info_ok
+          stores_req = get_from_connection(@bootstrap_conn, METADATA_STORE_NAME, STORES_INFO_KEY)
+          stores_req.callback do |stores_xml|
+            parse_stores_info(stores_xml)
+            finish_bootstrap
+          end
+          stores_req.errback do |error|
+            logger.warn "Could not load Voldemort's stores.xml: #{error}"
+            @bootstrap_state = :failed
+            finish_bootstrap
+          end
+        else
+          finish_bootstrap
+        end
+      end
+
+      cluster_req.errback do |error|
+        logger.warn "Could not load Voldemort's cluster.xml: #{error}"
+        @bootstrap_state = :failed
+        finish_bootstrap
+      end
+    end
+
+    def finish_bootstrap
+      @bootstrap_conn.close
+      @bootstrap_conn = nil
+      deferrable = @bootstrap
+      @bootstrap = nil
+      if @bootstrap_state == :complete
+        @bootstrap_timer.cancel
+        @bootstrap_timer = nil
+        deferrable.succeed
+      else
+        deferrable.fail
+      end
+    end
+
+    # Delays execution of the block until bootstrap has completed. Returns a new deferrable, and
+    # passes the same deferrable to the block when it is executed (it's up to the block to make the
+    # deferrable succeed or fail).
+    def when_ready(&block)
+      connect
+      request = EM::DefaultDeferrable.new
+      case @bootstrap_state
+      when :started, :cluster_info_ok
+        @bootstrap.callback { yield request }
+        @bootstrap.errback { request.fail(ServerError.new('Could not bootstrap Voldemort cluster')) }
+      when :complete
+        yield request
+      when :failed
+        request.fail(ServerError.new('Could not bootstrap Voldemort cluster'))
+      else
+        raise "bad bootstrap_state: #{@bootstrap_state.inspect}"
+      end
+      request
+    end
+
+    # Parses Voldemort's cluster.xml configuration file, as obtained in the bootstrap process.
+    def parse_cluster_info(cluster_xml)
+      doc = Nokogiri::XML(cluster_xml)
+      @cluster_name = doc.xpath('/cluster/name').text
+      @node_by_id = {}
+      @node_by_partition = {}
+      doc.xpath('/cluster/server').each do |node|
+        node_id = node.xpath('id').text.to_i
+        connection = Connection.new(
+          :host => node.xpath('host').text,
+          :port => node.xpath('socket-port').text.to_i,
+          :node_id => node_id,
+          :logger => logger
+        )
+        @node_by_id[node_id] = connection
+        node.xpath('partitions').text.split(/\s*,\s*/).map(&:to_i).each do |partition|
+          raise "duplicate assignment of partition #{partition}" if @node_by_partition[partition]
+          @node_by_partition[partition] = connection
+        end
+      end
+      raise 'no partitions defined on cluster' if @node_by_partition.empty?
+      (0...@node_by_partition.size).each do |partition|
+        raise "missing node assignment for partition #{partition}" unless @node_by_partition[partition]
+      end
+      @bootstrap_state = :cluster_info_ok
+    rescue => e
+      logger.warn "Error processing cluster.xml: #{e}"
+      @bootstrap_state = :failed
+    end
+
+    def parse_stores_info(stores_xml)
+      doc = Nokogiri::XML(stores_xml)
+      doc.xpath('/stores/store').each do |store|
+        store_name = store.xpath('name').text
+        @stores[store_name] ||= Store.new(self, store_name)
+        @stores[store_name].load_config(store)
+      end
+      @bootstrap_state = :complete
+    rescue => e
+      logger.warn "Error processing stores.xml: #{e}"
+      @bootstrap_state = :failed
+    end
+
+    def choose_connections(key, router=nil)
+      if router
+        router.partitions(key, @node_by_partition).map do |partition|
+          @node_by_partition[partition]
+        end.compact
+      else
+        @node_by_id.values.sample(2) # choose two random connections
+      end
+    end
+
+    # Makes a 'get' request to the first connection in the given list of connections. If the request
+    # fails with a ServerError, retries the request on the next connection in the list, etc,
+    # eventually failing if none of the connections can successfully handle the request.
+    def get_in_sequence(connections, store_name, key, deferrable)
+      raise ArgumentError, 'connections must not be empty' if connections.empty?
+      request = get_from_connection(connections.first, store_name, key)
+      request.callback {|response| deferrable.succeed(response) }
+      request.errback do |error|
+        if error.is_a?(ServerError) && connections.size > 1
+          get_in_sequence(connections.drop(1), store_name, key, deferrable)
+        else
+          deferrable.fail(error)
+        end
+      end
+    end
+
+    # Makes a 'get' request to all given connections in parallel. Succeeds with the first successful
+    # response, or fails if all connections' requests fail.
+    def get_in_parallel(connections, store_name, key, deferrable)
+      raise ArgumentError, 'connections must not be empty' if connections.empty?
+      responses = 0
+      done = false
+      connections.each do |connection|
+        request = get_from_connection(connection, store_name, key)
+        request.callback do |response|
+          deferrable.succeed(response) unless done
+          done = true
+        end
+        request.errback do |error|
+          if error.is_a?(ClientError) && !done
+            deferrable.fail(error)
+            done = true
+          elsif !done
+            responses += 1
+            deferrable.fail(error) if responses == connections.size
+          end
+        end
+      end
+    end
+
+    # Makes a 'get' request for a particular key to a particular Voldemort store, using a particular
+    # connection. Returns a deferrable that succeeds with the value in the store if successful, or
+    # fails with an exception object if not.
+    def get_from_connection(connection, store_name, key, deferrable=EM::DefaultDeferrable.new)
+      request = connection.send_request(get_request(store_name, key))
+
+      request.callback do |response|
+        begin
+          parsed_response = get_response(response)
+        rescue ClientError => error
+          deferrable.fail(error)
+        rescue => error
+          message = "protocol error #{error.class}: #{error.message} while parsing response: #{response.inspect}"
+          logger.error(message)
+          deferrable.fail(ServerError.new(message))
+        else
+          deferrable.succeed(parsed_response)
+        end
+      end
+
+      request.errback {|response| deferrable.fail(response) }
+      deferrable
+    end
+  end
+end

data/lib/em-voldemort/compressor.rb

@@ -0,0 +1,39 @@
+module EM::Voldemort
+  # Compression/decompression codec for keys and values in a store
+  class Compressor
+
+    attr_reader :type, :options
+
+    def initialize(xml)
+      @type = xml && xml.xpath('type').text
+      @options = xml && xml.xpath('options').text
+      if type != nil && type != 'gzip'
+        raise "Unsupported compression codec: #{type}"
+      end
+    end
+
+    def encode(data)
+      case type
+      when nil
+        data
+      when 'gzip'
+        buffer = StringIO.new
+        buffer.set_encoding(Encoding::BINARY)
+        gz = Zlib::GzipWriter.new(buffer)
+        gz.write(data)
+        gz.close
+        buffer.rewind
+        buffer.string
+      end
+    end
+
+    def decode(data)
+      case type
+      when nil
+        data
+      when 'gzip'
+        Zlib::GzipReader.new(StringIO.new(data)).read.force_encoding(Encoding::BINARY)
+      end
+    end
+  end
+end

data/lib/em-voldemort/connection.rb

@@ -0,0 +1,234 @@
+module EM::Voldemort
+  # TCP connection to one Voldemort node. The connection can be used to access multiple stores.
+  # Automatically reconnects if the connection is lost, but does not automatically retry failed
+  # requests (that is the cluster's job).
+  class Connection
+    attr_reader :host, :port, :node_id, :protocol, :logger, :health
+
+    DEFAULT_PROTOCOL = 'pb0' # Voldemort's protobuf-based protocol
+    STATUS_CHECK_PERIOD = 5 # Every 5 seconds, check on the health of the connection
+    REQUEST_TIMEOUT = 5 # If a request takes longer than 5 seconds, close the connection
+
+    def initialize(options={})
+      @host = options[:host] or raise ArgumentError, "#{self.class.name} requires :host"
+      @port = options[:port] or raise ArgumentError, "#{self.class.name} requires :port"
+      @node_id = options[:node_id]
+      @protocol = options[:protocol] || DEFAULT_PROTOCOL
+      @logger = options[:logger] || Logger.new($stdout)
+      @health = :good
+    end
+
+    # Establishes a connection to the node. Calling #connect is optional, since it also happens
+    # automatically when you start making requests.
+    def connect
+      force_connect unless @handler
+    end
+
+    # Sends a request to the node, given as a binary string (not including the request size prefix).
+    # Establishes a connection if necessary. If a request is already in progress, this request is
+    # queued up. Returns a deferrable that succeeds with the node's response (again without the size
+    # prefix), or fails if there was a network-level error.
+    def send_request(request)
+      connect
+      @handler.enqueue_request(request)
+    end
+
+    # Waits for the outstanding request (if any) to complete, then gracefully shuts down the
+    # connection.  Returns a deferrable that succeeds once the connection is closed (never fails).
+    def close
+      return @closing_deferrable if @closing_deferrable
+      @closing_deferrable = EM::DefaultDeferrable.new
+      @timer.cancel
+
+      if @handler
+        @handler.close_gracefully
+      else
+        @closing_deferrable.succeed
+      end
+
+      @handler = FailHandler.new(self)
+      @health = :bad
+      @closing_deferrable
+    end
+
+    # Called by the connection handler when the connection is closed for any reason (closed by us,
+    # closed by peer, rejected, timeout etc). Do not call from application code.
+    def connection_closed(handler, reason=nil)
+      logger.info ["Connection to Voldemort node #{host}:#{port} closed", reason].compact.join(': ')
+      @handler = FailHandler.new(self) if handler.equal? @handler
+      @health = :bad
+      @closing_deferrable.succeed if @closing_deferrable
+    end
+
+    private
+
+    def setup_status_check_timer
+      EM.add_periodic_timer(STATUS_CHECK_PERIOD) { yield }
+    end
+
+    def status_check
+      if @closing_deferrable
+        # Do nothing (don't reconnect once we've been asked to shut down).
+      elsif !@handler || @handler.is_a?(FailHandler)
+        # Connect for the first time, or reconnect after failure.
+        force_connect
+      elsif @handler.in_flight && Time.now - @handler.last_request >= REQUEST_TIMEOUT
+        # Request timed out. Pronounce the connection dead, and reconnect.
+        @handler.close_connection
+        force_connect
+      end
+    end
+
+    def force_connect
+      @timer ||= setup_status_check_timer(&method(:status_check))
+      @handler = EM.connect(host, port, Handler, self)
+      @handler.in_flight.callback { @health = :good } # restore health if protocol negotiation succeeded
+    rescue EventMachine::ConnectionError => e
+      # A synchronous exception is typically thrown on DNS resolution failure
+      logger.warn "Cannot connect to Voldemort node: #{e.class.name}: #{e.message}"
+      connection_closed(@handler)
+      @handler = FailHandler.new(self)
+    end
+
+
+    # EventMachine handler for a Voldemort node connection
+    module Handler
+      # The EM::Voldemort::Connection object for which we're handling the connection
+      attr_reader :connection
+
+      # State machine. One of :connecting, :protocol_proposal, :idle, :request, :disconnected
+      attr_reader :state
+
+      # If a request is currently in flight, this is a deferrable that will succeed or fail when the
+      # request completes. The protocol requires that only one request can be in flight at once.
+      attr_reader :in_flight
+
+      # The time at which the request currently in flight was sent
+      attr_reader :last_request
+
+      # Array of [request_data, deferrable] pairs, containing requests that have not yet been sent
+      attr_reader :request_queue
+
+      def initialize(connection)
+        @connection = connection
+        @state = :connecting
+        @in_flight = EM::DefaultDeferrable.new
+        @last_request = Time.now
+        @request_queue = []
+      end
+
+      def enqueue_request(request)
+        EM::DefaultDeferrable.new.tap do |deferrable|
+          request_queue << [request, deferrable]
+          send_next_request unless in_flight
+        end
+      end
+
+      # First action when the connection is established: client tells the server which version of
+      # the Voldemort protocol it wants to use
+      def send_protocol_proposal(protocol)
+        raise ArgumentError, 'protocol must be 3 bytes long' if protocol.bytesize != 3
+        raise "unexpected state before protocol proposal: #{@state.inspect}" unless @state == :connecting
+        send_data(protocol)
+        @state = :protocol_proposal
+      end
+
+      # Takes the request at the front of the queue and sends it to the Voldemort node
+      def send_next_request
+        return if request_queue.empty?
+        raise "cannot make a request while in #{@state.inspect} state" unless @state == :idle
+        request, @in_flight = request_queue.shift
+        send_data([request.size, request].pack('NA*'))
+        @recv_buf = ''.force_encoding('BINARY')
+        @last_request = Time.now
+        @state = :request
+      end
+
+      # Connection established (called by EventMachine)
+      def post_init
+        connection.logger.info "Connected to Voldemort node at #{connection.host}:#{connection.port}"
+        send_protocol_proposal(connection.protocol)
+        in_flight.errback do |response|
+          connection.logger.warn "Voldemort protocol #{connection.protocol} not accepted: #{response.inspect}"
+        end
+      end
+
+      # The Voldemort node is talking to us (called by EventMachine)
+      def receive_data(data)
+        case @state
+        when :protocol_proposal
+          deferrable = @in_flight
+          @state = :idle
+          @in_flight = nil
+          if data == 'ok'
+            send_next_request
+            deferrable.succeed
+          else
+            close_connection
+            deferrable.fail("server response: #{data.inspect}")
+          end
+
+        when :request
+          @recv_buf << data
+          response_size = @recv_buf.unpack('N').first
+          if response_size && @recv_buf.bytesize >= response_size + 4
+            response = @recv_buf[4, response_size]
+            deferrable = @in_flight
+            @state = :idle
+            @in_flight = @recv_buf = nil
+            send_next_request
+            deferrable.succeed(response)
+          end
+
+        else
+          raise "Received data in unexpected state: #{@state.inspect}"
+        end
+      end
+
+      # Connection is asking us to shut down. Wait for the currently in-flight request to complete,
+      # but fail any unsent requests in the queue.
+      def close_gracefully
+        @request_queue.each {|request, deferrable| deferrable.fail(ServerError.new('shutdown requested')) }
+        @request_queue = []
+        if in_flight
+          in_flight.callback { close_connection }
+          in_flight.errback  { close_connection }
+        else
+          close_connection
+        end
+      end
+
+      # Connection closed (called by EventMachine)
+      def unbind(reason=nil)
+        @state = :disconnected
+        deferrable = @in_flight
+        @in_flight = nil
+        deferrable.fail(ServerError.new('connection closed')) if deferrable
+        @request_queue.each {|request, deferrable| deferrable.fail(ServerError.new('connection closed')) }
+        @request_queue = []
+        connection.connection_closed(self, reason)
+      end
+    end
+
+
+    # Quacks like a EM::Voldemort::Connection::Handler, but fails all requests.
+    # Useful for representing a connection in an error state.
+    class FailHandler
+      attr_reader :in_flight
+
+      def initialize(connection)
+        @connection = connection
+      end
+
+      def enqueue_request(request)
+        EM::DefaultDeferrable.new.tap do |deferrable|
+          deferrable.fail(ServerError.new('Connection to Voldemort node closed'))
+        end
+      end
+
+      def close_gracefully
+        @connection.connection_closed(self)
+      end
+    end
+  end
+end