krakow 0.2.2 → 0.3.0

data/lib/krakow/connection.rb

@@ -1,13 +1,29 @@
-require 'krakow/version'
+require 'krakow'
 require 'celluloid/io'
-require 'celluloid/autostart'
 
 module Krakow
+
+  # Provides TCP connection to NSQD
   class Connection
 
+    # Generate identifier for connection
+    #
+    # @param host [String]
+    # @param port [String, Integer]
+    # @param topic [String]
+    # @param channel [String]
+    # @return [String]
+    def self.identifier(host, port, topic, channel)
+      [host, port, topic, channel].compact.join('__')
+    end
+
     include Utils::Lazy
+    # @!parse include Krakow::Utils::Lazy::InstanceMethods
+    # @!parse extend Krakow::Utils::Lazy::ClassMethods
+
     include Celluloid::IO
 
+    # Available connection features
     FEATURES = [
       :max_rdy_count,
       :max_msg_timeout,
@@ -19,63 +35,125 @@ module Krakow
       :snappy,
       :sample_rate
     ]
+
+    # List of features that may not be enabled together
     EXCLUSIVE_FEATURES = [[:snappy, :deflate]]
+
+    # List of features that may be enabled by the client
     ENABLEABLE_FEATURES = [:tls_v1, :snappy, :deflate]
 
     finalizer :goodbye_my_love!
 
-    attr_reader :socket, :endpoint_settings
+    # @return [Hash] current configuration for endpoint
+    attr_reader :endpoint_settings
+    # @return [Socket-ish] underlying socket like instance
+    attr_reader :socket
+
+    attr_reader :connector, :reconnector, :reconnect_notifier, :responder, :running
+
+    # @!group Attributes
 
+    # @!macro [attach] attribute
+    #   @!method $1
+    #     @return [$2] the $1 $0
+    #   @!method $1?
+    #     @return [TrueClass, FalseClass] truthiness of the $1 $0
+    attribute :host, String, :required => true
+    attribute :port, [String,Integer], :required => true
+    attribute :topic, String
+    attribute :channel, String
+    attribute :version, String, :default => 'v2'
+    attribute :queue, Queue, :default => ->{ Queue.new }
+    attribute :callbacks, Hash, :default => ->{ Hash.new }
+    attribute :responses, Queue, :default => ->{ Queue.new }
+    attribute :notifier, Celluloid::Actor
+    attribute :features, Hash, :default => ->{ Hash.new }
+    attribute :response_wait, Numeric, :default => 1.0
+    attribute :response_interval, Numeric, :default => 0.03
+    attribute :error_wait, Numeric, :default => 0
+    attribute :enforce_features, [TrueClass,FalseClass], :default => true
+    attribute :features_args, Hash, :default => ->{ Hash.new }
+
+    # @!endgroup
+
+    # Create new instance
+    #
+    # @param args [Hash]
+    # @option args [String] :host (required) server host
+    # @option args [String, Numeric] :port (required) server port
+    # @option args [String] :version
+    # @option args [Queue] :queue received message queue
+    # @option args [Hash] :callbacks
+    # @option args [Queue] :responses received responses queue
+    # @option args [Celluloid::Actor] :notifier actor to notify on new message
+    # @option args [Hash] :features features to enable
+    # @option args [Numeric] :response_wait time to wait for response
+    # @option args [Numeric] :response_interval sleep interval for wait loop
+    # @option args [Numeric] :error_wait time to wait for error response
+    # @option args [TrueClass, FalseClass] :enforce_features fail if features are unavailable
+    # @option args [Hash] :feature_args options for connection features
     def initialize(args={})
       super
-      required! :host, :port
-      optional(
-        :version, :queue, :callback, :responses, :notifier,
-        :features, :response_wait, :error_wait, :enforce_features
-      )
-      arguments[:queue] ||= Queue.new
-      arguments[:responses] ||= Queue.new
-      arguments[:version] ||= 'v2'
-      arguments[:features] ||= {}
-      arguments[:response_wait] ||= 1
-      arguments[:error_wait] ||= 0.4
-      if(arguments[:enforce_features].nil?)
-        arguments[:enforce_features] = true
-      end
-      @socket = TCPSocket.new(host, port)
+      @connector = Mutex.new
+      @reconnector = Mutex.new
+      @responder = Mutex.new
+      @reconnect_notifier = Celluloid::Signals.new
+      @socket_retries = 0
+      @socket_max_retries = 10
+      @reconnect_pause = 0.5
       @endpoint_settings = {}
+      @running = false
     end
 
+    # @return [String] identifier for this connection
+    def identifier
+      self.class.identifier(host, port, topic, channel)
+    end
+
+    # @return [String] stringify object
     def to_s
       "<#{self.class.name}:#{object_id} {#{host}:#{port}}>"
     end
 
     # Initialize the connection
+    #
+    # @return [nil]
     def init!
-      debug 'Initializing connection'
-      socket.write version.rjust(4).upcase
-      identify_and_negotiate
-      async.process_to_queue!
-      info 'Connection initialized'
+      connector.synchronize do
+        connect!
+      end
+      nil
     end
 
-    # message:: Command instance to send
-    # Send the message
-    # TODO: Do we want to validate Command instance and abort if
-    # response is already set?
+    # Send message to remote server
+    #
+    # @param message [Krakow::Message] message to send
+    # @return [TrueClass, Krakow::FrameType] response if expected or true
     def transmit(message)
       output = message.to_line
-      debug ">>> #{output}"
-      socket.write output
       response_wait = wait_time_for(message)
-      responses.clear if response_wait
-      if(response_wait)
+      if(response_wait > 0)
+        transmit_with_response(message, response_wait)
+      else
+        debug ">>> #{output}"
+        safe_socket{|socket| socket.write output }
+        true
+      end
+    end
+
+    # Sends message and waits for response
+    #
+    # @param message [Krakow::Message] message to send
+    # @return [Krakow::FrameType] response
+    def transmit_with_response(message, wait_time)
+      responder.synchronize do
+        safe_socket{|socket| socket.write(message.to_line) }
+        responses.clear
         response = nil
-        (response_wait / 0.1).to_i.times do |i|
+        (wait_time / response_interval).to_i.times do |i|
           response = responses.pop unless responses.empty?
           break if response
-          debug "Response wait sleep for 0.1 seconds (#{i+1} time)"
-          sleep(0.1)
+          sleep(response_interval)
         end
         if(response)
           message.response = response
@@ -90,23 +168,32 @@ module Krakow
             abort Error::BadResponse::NoResponse.new "No response provided for message #{message}"
           end
         end
-      else
-        true
       end
     end
 
-    # Cleanup prior to destruction
+    # Destructor method for cleanup
+    #
+    # @return [nil]
     def goodbye_my_love!
       debug 'Tearing down connection'
       if(socket && !socket.closed?)
-        socket.write Command::Cls.new.to_line
-        socket.close
+        [lambda{ socket.write Command::Cls.new.to_line}, lambda{socket.close}].each do |action|
+          begin
+            action.call
+          rescue IOError, SystemCallError => e
+            warn "Socket error encountered during teardown: #{e.class}: #{e}"
+          end
+        end
       end
       @socket = nil
       info 'Connection torn down'
+      nil
     end
 
-    # Receive message and return proper FrameType instance
+    # Receive from server
+    #
+    # @return [Krakow::FrameType, nil] message or nothing if read was empty
+    # @raise [Error::ConnectionUnavailable] socket is closed
     def receive
       debug 'Read wait for frame start'
       buf = socket.recv(8)
@@ -115,7 +202,7 @@ module Krakow
         debug "<<< #{buf.inspect}"
         struct = FrameType.decode(buf)
         debug "Decoded structure: #{struct.inspect}"
-        struct[:data] = socket.recv(struct[:size])
+        struct[:data] = socket.read(struct[:size])
         debug "<<< #{struct[:data].inspect}"
         @receiving = false
         frame = FrameType.build(struct)
@@ -123,31 +210,43 @@ module Krakow
         frame
       else
         if(socket.closed?)
-          raise Error.new("#{self} encountered closed socket!")
+          abort Error::ConnectionUnavailable.new("#{self} encountered closed socket!")
         end
         nil
       end
     end
 
-    # Currently in the process of receiving a message
+    # @return [TrueClass, FalseClass] is connection currently receiving a message
     def receiving?
       !!@receiving
     end
 
-    # Pull message and queue
+    # Receive messages and place into queue
+    #
+    # @return [nil]
     def process_to_queue!
-      loop do
-        message = handle(receive)
-        if(message)
-          debug "Adding message to queue #{message}"
-          queue << message
-          notifier.signal(message) if notifier
+      @running = true
+      while(@running)
+        begin
+          message = handle(receive)
+          if(message)
+            debug "Adding message to queue #{message}"
+            queue << message
+            notifier.signal(message) if notifier
+          end
+        rescue Error::ConnectionUnavailable => e
+          warn "Failed to receive message: #{e.class} - #{e}"
+          @running = false
+          async.reconnect!
         end
       end
+      nil
     end
 
-    # message:: FrameType instance
-    # Handle message if not an actual message
+    # Handle non-message type Krakow::FrameType
+    #
+    # @param message [Krakow::FrameType] received message
+    # @return [Krakow::FrameType, nil]
     def handle(message)
       # Grab heartbeats upfront
       if(message.is_a?(FrameType::Response) && message.response == '_heartbeat_')
@@ -155,10 +254,7 @@ module Krakow
         transmit Command::Nop.new
         nil
       else
-        if(callback && callback[:actor] && callback[:method])
-          debug "Sending #{message} to callback `#{callback[:actor]}##{callback[:method]}`"
-          message = callback[:actor].send(callback[:method], message, current_actor)
-        end
+        message = callback_for(:handle, message)
         if(!message.is_a?(FrameType::Message))
           debug "Captured non-message type response: #{message}"
           responses << message
@@ -169,6 +265,28 @@ module Krakow
       end
     end
 
+    # Execute callback for given type
+    #
+    # @overload callback_for(type, arg, connection)
+    #   @param type [Symbol] type of callback
+    #   @param arg [Object] argument for callback (can be multiple)
+    #   @param connection [Krakow::Connection] current connection
+    # @return [Object] result of callback
+    def callback_for(type, *args)
+      callback = callbacks[type]
+      if(callback)
+        debug "Processing connection callback for #{type.inspect} (#{callback.inspect})"
+        callback[:actor].send(callback[:method], *(args + [current_actor]))
+      else
+        debug "No connection callback defined for #{type.inspect}"
+        args.size == 1 ? args.first : args
+      end
+    end
+
+    # Returns configured wait time for given message type
+    #
+    # @param message [Krakow::Command]
+    # @return [Numeric] seconds to wait
     def wait_time_for(message)
       case Command.response_for(message)
       when :required
@@ -178,6 +296,7 @@ module Krakow
       end
     end
 
+    # @return [Hash] default settings for IDENTIFY
     def identify_defaults
       unless(@identify_defaults)
         @identify_defaults = {
@@ -190,12 +309,15 @@ module Krakow
       @identify_defaults
     end
 
+    # IDENTIFY with server and negotiate features
+    #
+    # @return [TrueClass]
     def identify_and_negotiate
       expected_features = identify_defaults.merge(features)
       ident = Command::Identify.new(
         expected_features
       )
-      socket.write(ident.to_line)
+      safe_socket{|socket| socket.write(ident.to_line) }
       response = receive
       if(expected_features[:feature_negotiation])
         begin
@@ -219,25 +341,129 @@ module Krakow
       true
     end
 
+    # Enable snappy feature on underlying socket
+    #
+    # @return [TrueClass]
     def snappy
       info 'Loading support for snappy compression and converting connection'
-      @socket = ConnectionFeatures::SnappyFrames::Io.new(socket)
+      @socket = ConnectionFeatures::SnappyFrames::Io.new(socket, features_args)
       response = receive
       info "Snappy connection conversion complete. Response: #{response.inspect}"
+      true
     end
 
+    # Enable deflate feature on underlying socket
+    #
+    # @return [TrueClass]
     def deflate
       debug 'Loading support for deflate compression and converting connection'
-      @socket = ConnectionFeatures::Deflate::Io.new(socket)
+      @socket = ConnectionFeatures::Deflate::Io.new(socket, features_args)
       response = receive
       info "Deflate connection conversion complete. Response: #{response.inspect}"
+      true
     end
 
+    # Enable TLS feature on underlying socket
+    #
+    # @return [TrueClass]
     def tls_v1
       info 'Enabling TLS for connection'
-      @socket = ConnectionFeatures::Ssl::Io.new(socket)
+      @socket = ConnectionFeatures::Ssl::Io.new(socket, features_args)
       response = receive
       info "TLS enable complete. Response: #{response.inspect}"
+      true
+    end
+
+    # @return [TrueClass, FalseClass] underlying socket is connected
+    def connected?
+      socket && !socket.closed?
+    end
+
+    protected
+
+    # Destruct the underlying socket
+    #
+    # @return [nil]
+    def teardown_socket
+      if(socket && (socket.closed? || socket.eof?))
+        socket.close unless socket.closed?
+        @socket = nil
+        warn 'Existing socket instance has been destroyed from this connection'
+      end
+      nil
+    end
+
+    # Provides socket failure state handling around given block
+    #
+    # @yield [socket] execute within socket safety layer
+    # @yieldparam [socket] underlying socket
+    # @return [Object] result of executed block
+    def safe_socket(*args)
+      begin
+        if(socket.nil? || socket.closed?)
+          raise Error::ConnectionUnavailable.new 'Current connection is closed!'
+        end
+        result = yield socket if block_given?
+        result
+      rescue Error::ConnectionUnavailable, SystemCallError, IOError => e
+        warn "Safe socket encountered error (socket in failed state): #{e.class}: #{e}"
+        reconnect!
+        retry
+      rescue Celluloid::Error => e
+        warn "Internal error encountered. Allowing exception to bubble. #{e.class}: #{e}"
+        abort e
+      rescue Exception => e
+        warn "!!! Unexpected error encountered within safe socket: #{e.class}: #{e}"
+        raise
+      end
+    end
+
+    # Reconnect the underlying socket
+    #
+    # @return [nil]
+    def reconnect!
+      if(reconnector.try_lock)
+        begin
+          if(@socket_max_retries <= @socket_retries)
+            abort ConnectionFailure.new "Failed to re-establish connection after #{@socket_retries} tries."
+          end
+          pause_interval = @reconnect_pause * @socket_retries
+          @socket_retries += 1
+          warn "Pausing for #{pause_interval} seconds before reconnect"
+          sleep(pause_interval)
+          init!
+          @socket_retries = 0
+        rescue Celluloid::Error => e
+          warn "Internal error encountered. Allowing exception to bubble. #{e.class}: #{e}"
+          abort e
+        rescue SystemCallError, IOError => e
+          error "Reconnect error encountered: #{e.class} - #{e}"
+          retry
+        end
+        callback_for(:reconnect)
+        reconnect_notifier.broadcast(:connected)
+        reconnector.unlock
+      else
+        reconnect_notifier.wait(:connected)
+      end
+      nil
+    end
+
+    # Connect the underlying socket
+    #
+    # @return [nil]
+    def connect!
+      debug 'Initializing connection'
+      if(@socket)
+        @socket.close unless @socket.closed?
+        @socket = nil
+      end
+      @socket = Celluloid::IO::TCPSocket.new(host, port)
+      safe_socket{|socket| socket.write version.rjust(4).upcase}
+      identify_and_negotiate
+      async.process_to_queue!
+      info 'Connection initialized'
+      nil
     end
 
   end

data/lib/krakow/connection_features/deflate.rb

@@ -1,13 +1,20 @@
 require 'zlib'
+require 'krakow'
 
 module Krakow
   module ConnectionFeatures
+    # Deflate functionality
     module Deflate
+      # Deflatable IO
       class Io
 
         attr_reader :io, :buffer, :headers, :inflator, :deflator
 
-        def initialize(io)
+        # Create new deflatable IO
+        #
+        # @param io [IO] IO to wrap
+        # @return [Io]
+        def initialize(io, args={})
           @io = io
           @buffer = ''
           @inflator = Zlib::Inflate.new(-Zlib::MAX_WBITS)
@@ -15,10 +22,17 @@ module Krakow
         end
 
         # Proxy to underlying socket
+        #
+        # @param args [Object]
+        # @return [Object]
         def method_missing(*args)
           io.__send__(*args)
         end
 
+        # Receive bytes from the IO
+        #
+        # @param n [Integer] nuber of bytes
+        # @return [String]
         def recv(n)
           until(buffer.length >= n)
             read_stream
@@ -28,6 +42,9 @@ module Krakow
         end
         alias_method :read, :recv
 
+        # Read contents from stream
+        #
+        # @return [String]
         def read_stream
           str = io.read
           unless(str.empty?)
@@ -35,6 +52,10 @@ module Krakow
           end
         end
 
+        # Write string to IO
+        #
+        # @param string [String]
+        # @return [Integer] number of bytes written
         def write(string)
           unless(string.empty?)
             output = deflator.deflate(string)
@@ -45,10 +66,14 @@ module Krakow
           end
         end
 
+        # Close the IO
+        #
+        # @return [TrueClass]
         def close(*args)
           super
           deflator.deflate(nil, Zlib::FINISH)
           deflator.close
+          true
         end
 
       end

data/lib/krakow/connection_features/snappy_frames.rb

@@ -1,18 +1,24 @@
 require 'snappy'
 require 'digest/crc'
+require 'krakow'
 
-# TODO: Add support for max size + chunks
-# TODO: Include support for remaining types
 module Krakow
   module ConnectionFeatures
+    # Snappy functionality
+    # @todo Add support for max size + chunks
+    # @todo Include support for remaining types
     module SnappyFrames
+      # Snappy-able IO
       class Io
 
+        # Header identifier
         IDENTIFIER = "\x73\x4e\x61\x50\x70\x59".force_encoding('ASCII-8BIT')
         ident_size = [IDENTIFIER.size].pack('L<')
         ident_size.slice!(-1,1)
+        # Size of identifier
         IDENTIFIER_SIZE = ident_size
 
+        # Mapping of types
         CHUNK_TYPE = {
           "\xff".force_encoding('ASCII-8BIT') => :identifier,
           "\x00".force_encoding('ASCII-8BIT') => :compressed,
@@ -21,21 +27,36 @@ module Krakow
 
         attr_reader :io, :buffer
 
-        def initialize(io)
+        # Create new snappy-able IO
+        #
+        # @param io [IO] IO to wrap
+        # @return [Io]
+        def initialize(io, args={})
           @io = io
           @snappy_write_ident = false
           @buffer = ''
         end
 
         # Proxy to underlying socket
+        #
+        # @param args [Object]
+        # @return [Object]
         def method_missing(*args)
           io.__send__(*args)
         end
 
+        # Mask the checksum
+        #
+        # @param checksum [String]
+        # @return [String]
         def checksum_mask(checksum)
           (((checksum >> 15) | (checksum << 17)) + 0xa282ead8) & 0xffffffff
         end
 
+        # Receive bytes from the IO
+        #
+        # @param n [Integer] nuber of bytes
+        # @return [String]
         def recv(n)
           read_stream unless buffer.size >= n
           result = buffer.slice!(0,n)
@@ -43,6 +64,9 @@ module Krakow
         end
         alias_method :read, :recv
 
+        # Read contents from stream
+        #
+        # @return [String]
         def read_stream
           header = io.recv(4)
           ident = CHUNK_TYPE[header.slice!(0)]
@@ -68,6 +92,10 @@ module Krakow
           end
         end
 
+        # Write string to IO
+        #
+        # @param string [String]
+        # @return [Integer] number of bytes written
         def write(string)
           unless(@snappy_writer_ident)
             send_snappy_identifier
@@ -83,6 +111,9 @@ module Krakow
           io.write output
         end
 
+        # Send the identifier for snappy content
+        #
+        # @return [Integer] bytes written
         def send_snappy_identifier
           io.write [CHUNK_TYPE.key(:identifier), IDENTIFIER_SIZE, IDENTIFIER].pack('a*a*a*')
         end