nsq-krakow 0.1.0

data/lib/krakow/connection_features.rb

@@ -0,0 +1,10 @@
+require 'krakow'
+
+module Krakow
+  # Features that wrap the connection
+  module ConnectionFeatures
+    autoload :SnappyFrames, 'krakow/connection_features/snappy_frames'
+    autoload :Deflate, 'krakow/connection_features/deflate'
+    autoload :Ssl, 'krakow/connection_features/ssl'
+  end
+end

data/lib/krakow/connection_features/deflate.rb

@@ -0,0 +1,82 @@
+require 'zlib'
+require 'krakow'
+
+module Krakow
+  module ConnectionFeatures
+    # Deflate functionality
+    module Deflate
+      # Deflatable IO
+      class Io
+
+        attr_reader :io, :buffer, :headers, :inflator, :deflator
+
+        # 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)
+          @deflator = Zlib::Deflate.new(nil, -Zlib::MAX_WBITS)
+        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
+            sleep(0.1) unless buffer.length >= n
+          end
+          buffer.slice!(0, n)
+        end
+        alias_method :read, :recv
+
+        # Read contents from stream
+        #
+        # @return [String]
+        def read_stream
+          str = io.read
+          unless(str.empty?)
+            buffer << inflator.inflate(str)
+          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)
+            output << deflator.flush
+            io.write(output)
+          else
+            0
+          end
+        end
+
+        # Close the IO
+        #
+        # @return [TrueClass]
+        def close(*args)
+          super
+          deflator.deflate(nil, Zlib::FINISH)
+          deflator.close
+          true
+        end
+
+      end
+    end
+  end
+end

data/lib/krakow/connection_features/snappy_frames.rb

@@ -0,0 +1,129 @@
+begin
+  require 'snappy'
+rescue LoadError
+  $stderr.puts 'ERROR: Failed to locate `snappy` gem. Install `snappy` gem into system or bundle.'
+  raise
+end
+require 'digest/crc'
+require 'krakow'
+
+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,
+          "\x01".force_encoding('ASCII-8BIT') => :uncompressed
+        }
+
+        attr_reader :io, :buffer
+
+        # 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)
+          result.empty? ? nil : result
+        end
+        alias_method :read, :recv
+
+        # Read contents from stream
+        #
+        # @return [String]
+        def read_stream
+          header = io.recv(4)
+          ident = CHUNK_TYPE[header.slice!(0)]
+          size = (header << CHUNK_TYPE.key(:compressed)).unpack('L<').first
+          content = io.recv(size)
+          case ident
+          when :identifier
+            unless(content == IDENTIFIER)
+              raise "Invalid stream identification encountered (content: #{content.inspect})"
+            end
+            read_stream
+          when :compressed
+            checksum = content.slice!(0, 4).unpack('L<').first
+            deflated = Snappy.inflate(content)
+            digest = Digest::CRC32c.new
+            digest << deflated
+            unless(checksum == checksum_mask(digest.checksum))
+              raise 'Checksum mismatch!'
+            end
+            buffer << deflated
+          when :uncompressed
+            buffer << content
+          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
+          end
+          digest = Digest::CRC32c.new
+          digest << string
+          content = Snappy.deflate(string)
+          size = content.length + 4
+          size = [size].pack('L<')
+          size.slice!(-1,1)
+          checksum = [checksum_mask(digest.checksum)].pack('L<')
+          output = [CHUNK_TYPE.key(:compressed), size, checksum, content].pack('a*a*a*a*')
+          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
+
+      end
+    end
+  end
+end

data/lib/krakow/connection_features/ssl.rb

@@ -0,0 +1,75 @@
+require 'openssl'
+require 'krakow'
+
+module Krakow
+  module ConnectionFeatures
+    # SSL functionality
+    module Ssl
+      # SSL-able IO
+      class Io
+
+        attr_reader :_socket
+
+        # Create new SSL-able IO
+        #
+        # @param io [IO] IO to wrap
+        # @param args [Hash]
+        # @option args [Hash] :ssl_context
+        # @return [Io]
+        def initialize(io, args={})
+          ssl_socket_arguments = [io]
+          if(args[:ssl_context])
+            validate_ssl_args!(args[:ssl_context])
+            context = OpenSSL::SSL::SSLContext.new
+            context.cert = OpenSSL::X509::Certificate.new(File.open(args[:ssl_context][:certificate]))
+            context.key = OpenSSL::PKey::RSA.new(File.open(args[:ssl_context][:key]))
+            ssl_socket_arguments << context
+          end
+          @_socket = Celluloid::IO::SSLSocket.new(*ssl_socket_arguments)
+          _socket.sync = true
+          _socket.connect
+        end
+
+        # Proxy to underlying socket
+        #
+        # @param args [Object]
+        # @return [Object]
+        def method_missing(*args)
+          _socket.send(*args)
+        end
+
+        # Receive bytes from the IO
+        #
+        # @param len [Integer] nuber of bytes
+        # @return [String]
+        def recv(len)
+          str = readpartial(len)
+          if(len > str.length)
+            str << sysread(len - str.length)
+          end
+          str
+        end
+
+        private
+
+        # Validate the SSL configuration provided
+        #
+        # @param args [Hash]
+        # @option args [String] :certificate path to certificate
+        # @option args [String] :key path to key
+        # @raise [ArgumentError, LoadError]
+        def validate_ssl_args!(args)
+          [:key, :certificate].each do |arg_key|
+            unless(args.has_key?(arg_key))
+              raise ArgumentError.new "The `:ssl_context` option requires `#{arg_key.inspect}` to be set"
+            end
+            unless(File.readable?(args[arg_key]))
+              raise LoadError.new "Unable to read the `#{arg_key.inspect}` file from the `:ssl_context` arguments"
+            end
+          end
+        end
+
+      end
+    end
+  end
+end

data/lib/krakow/consumer.rb

@@ -0,0 +1,355 @@
+require 'krakow'
+
+module Krakow
+  # Consume messages from a server
+  class Consumer
+
+    autoload :Queue, 'krakow/consumer/queue'
+
+    include Utils::Lazy
+    # @!parse include Krakow::Utils::Lazy::InstanceMethods
+    # @!parse extend Krakow::Utils::Lazy::ClassMethods
+
+    include Celluloid
+
+    trap_exit :connection_failure
+    finalizer :consumer_cleanup
+
+    attr_reader :connections, :discovery, :distribution, :queue
+
+    # @!group Attributes
+
+    # @!macro [attach] attribute
+    #   @!method $1
+    #     @return [$2] the $1 $0
+    #   @!method $1?
+    #     @return [TrueClass, FalseClass] truthiness of the $1 $0
+    attribute :topic, String, :required => true
+    attribute :channel, String, :required => true
+    attribute :host, String
+    attribute :port, [String, Integer]
+    attribute :nsqlookupd, [Array, String]
+    attribute :max_in_flight, Integer, :default => 1
+    attribute :backoff_interval, Numeric
+    attribute :discovery_interval, Numeric, :default => 30
+    attribute :discovery_jitter, Numeric, :default => 10.0
+    attribute :notifier, [Celluloid::Signals, Celluloid::Condition, Celluloid::Actor]
+    attribute :connection_options, Hash, :default => ->{ Hash.new }
+
+    # @!endgroup
+
+    def initialize(args={})
+      super
+      arguments[:connection_options] = {:features => {}, :config => {}}.merge(
+        arguments[:connection_options] || {}
+      )
+      @connections = {}
+      @queue = Queue.new(
+        current_actor,
+        :removal_callback => :remove_message
+      )
+      @distribution = Distribution::Default.new(
+        :max_in_flight => max_in_flight,
+        :backoff_interval => backoff_interval,
+        :consumer => current_actor
+      )
+      if(nsqlookupd)
+        debug "Connections will be established via lookup #{nsqlookupd.inspect}"
+        @discovery = Discovery.new(:nsqlookupd => nsqlookupd)
+        discover
+      elsif(host && port)
+        direct_connect
+      else
+        abort Error::ConfigurationError.new('No connection information provided!')
+      end
+    end
+
+    # @return [TrueClass, FalseClass] currently connected to at least
+    #   one nsqd
+    def connected?
+      !!connections.values.any? do |con|
+        begin
+          con.connected?
+        rescue Celluloid::DeadActorError
+          false
+        end
+      end
+    end
+
+    # Connect to nsqd instance directly
+    #
+    # @return [Connection]
+    def direct_connect
+      debug "Connection will be established via direct connection #{host}:#{port}"
+      connection = build_connection(host, port, queue)
+      if(register(connection))
+        info "Registered new connection #{connection}"
+        distribution.redistribute!
+      else
+        abort Error::ConnectionFailure.new("Failed to establish subscription at provided end point (#{host}:#{port}")
+      end
+      connection
+    end
+
+    # Returns [Krakow::Connection] associated to key
+    #
+    # @param key [Object] identifier
+    # @return [Krakow::Connection] associated connection
+    def connection(key)
+      @connections[key]
+    end
+
+    # @return [String] stringify object
+    def to_s
+      "<#{self.class.name}:#{object_id} T:#{topic} C:#{channel}>"
+    end
+
+    # Instance destructor
+    #
+    # @return [nil]
+    def consumer_cleanup
+      debug 'Tearing down consumer'
+      if(distribution && distribution.alive?)
+        distribution.terminate
+      end
+      if(queue && queue.alive?)
+        queue.terminate
+      end
+      connections.values.each do |con|
+        con.terminate if con.alive?
+      end
+      info 'Consumer torn down'
+      nil
+    end
+
+    # Build a new [Krakow::Connection]
+    #
+    # @param host [String] remote host
+    # @param port [String, Integer] remote port
+    # @param queue [Queue] queue for messages
+    # @return [Krakow::Connection, nil] new connection or nil
+    def build_connection(host, port, queue)
+      begin
+        connection = Connection.new(
+          :host => host,
+          :port => port,
+          :queue => queue,
+          :topic => topic,
+          :channel => channel,
+          :notifier => notifier,
+          :features => connection_options[:features],
+          :features_args => connection_options[:config],
+          :callbacks => {
+            :handle => {
+              :actor => current_actor,
+              :method => :process_message
+            }
+          }
+        )
+        queue.register_connection(connection)
+        connection
+      rescue => e
+        error "Failed to build connection (host: #{host} port: #{port} queue: #{queue}) - #{e.class}: #{e}"
+        debug "#{e.class}: #{e}\n#{e.backtrace.join("\n")}"
+        nil
+      end
+    end
+
+    # Process a given message if required
+    #
+    # @param message [Krakow::FrameType]
+    # @param connection [Krakow::Connection]
+    # @return [Krakow::FrameType]
+    # @note If we receive a message that is already in flight, attempt
+    #   to scrub message from wait queue. If message is found, retry
+    #   distribution registration. If message is not found, assume it
+    #   is currently being processed and do not allow new message to
+    #   be queued
+    def process_message(message, connection)
+      discard = false
+      if(message.is_a?(FrameType::Message))
+        message.origin = current_actor
+        message.connection = connection
+        retried = false
+        begin
+          distribution.register_message(message, connection.identifier)
+        rescue KeyError => e
+          if(!retried && queue.scrub_duplicate_message(message))
+            retried = true
+            retry
+          else
+            error "Received message is currently in flight and not in wait queue. Discarding! (#{message})"
+            discard = true
+          end
+        end
+      end
+      discard ? nil : message
+    end
+
+    # Send RDY for connection based on distribution rules
+    #
+    # @param connection [Krakow::Connection]
+    # @return [nil]
+    def update_ready!(connection)
+      distribution.set_ready_for(connection)
+      nil
+    end
+
+    # Initialize the consumer by starting lookup and adding connections
+    #
+    # @return [nil]
+    def init!
+      debug 'Running consumer `init!` connection builds'
+      found = discovery.lookup(topic)
+      debug "Discovery results: #{found.inspect}"
+      connection = nil
+      found.each do |node|
+        debug "Processing discovery result: #{node.inspect}"
+        key = Connection.identifier(node[:broadcast_address], node[:tcp_port], topic, channel)
+        unless(connections[key])
+          connection = build_connection(node[:broadcast_address], node[:tcp_port], queue)
+          info "Registered new connection #{connection}" if register(connection)
+        else
+          debug "Discovery result already registered: #{node.inspect}"
+        end
+      end
+      distribution.redistribute! if connection
+      nil
+    end
+
+    # Start the discovery interval lookup
+    #
+    # @return [nil]
+    def discover
+      init!
+      after(discovery_interval + (discovery_jitter * rand)){ discover }
+    end
+
+    # Register connection with distribution
+    #
+    # @param connection [Krakow::Connection]
+    # @return [TrueClass, FalseClass] true if subscription was successful
+    def register(connection)
+      begin
+        connection.init!
+        connection.transmit(Command::Sub.new(:topic_name => topic, :channel_name => channel))
+        self.link connection
+        connections[connection.identifier] = connection
+        distribution.add_connection(connection)
+        true
+      rescue Error::BadResponse => e
+        debug "Failed to establish connection: #{e.result ? e.result.error : '<No Response!>'}"
+        connection.terminate
+        false
+      end
+    end
+
+    # Remove connection references when connection is terminated
+    #
+    # @param actor [Object] terminated actor
+    # @param reason [Exception] reason for termination
+    # @return [nil]
+    def connection_failure(actor, reason)
+      if(reason && key = connections.key(actor))
+        warn "Connection failure detected. Removing connection: #{key} - #{reason}"
+        connections.delete(key)
+        begin
+          distribution.remove_connection(key)
+        rescue Error::ConnectionUnavailable, Error::ConnectionFailure
+          warn 'Caught connection unavailability'
+        end
+        queue.deregister_connection(key)
+        distribution.redistribute!
+        direct_connect unless discovery
+      end
+      nil
+    end
+
+    # Remove message
+    #
+    # @param messages [Array<FrameType::Message>]
+    # @return [NilClass]
+    # @note used mainly for queue callback
+    def remove_message(messages)
+      [messages].flatten.compact.each do |msg|
+        distribution.unregister_message(msg.message_id)
+        update_ready!(msg.connection)
+      end
+      nil
+    end
+
+    # Confirm message has been processed
+    #
+    # @param message_id [String, Krakow::FrameType::Message]
+    # @return [TrueClass]
+    # @raise [KeyError] connection not found
+    def confirm(message_id)
+      message_id = message_id.message_id if message_id.respond_to?(:message_id)
+      begin
+        begin
+          connection = distribution.in_flight_lookup(message_id)
+          connection.transmit(Command::Fin.new(:message_id => message_id))
+          distribution.success(connection.identifier)
+        rescue => e
+          abort e
+        end
+        true
+      rescue KeyError => e
+        error "Message confirmation failed: #{e}"
+        abort e
+      rescue Error::LookupFailed => e
+        error "Lookup of message for confirmation failed! <Message ID: #{message_id} - Error: #{e}>"
+        abort e
+      rescue Error::ConnectionUnavailable => e
+        abort e
+      rescue Celluloid::DeadActorError
+        abort Error::ConnectionUnavailable.new
+      ensure
+        con = distribution.unregister_message(message_id)
+        update_ready!(con) if con
+      end
+    end
+    alias_method :finish, :confirm
+
+    # Requeue message (generally due to processing failure)
+    #
+    # @param message_id [String, Krakow::FrameType::Message]
+    # @param timeout [Numeric]
+    # @return [TrueClass]
+    def requeue(message_id, timeout=0)
+      message_id = message_id.message_id if message_id.respond_to?(:message_id)
+      distribution.in_flight_lookup(message_id) do |connection|
+        distribution.unregister_message(message_id)
+        connection.transmit(
+          Command::Req.new(
+            :message_id => message_id,
+            :timeout => timeout
+          )
+        )
+        distribution.failure(connection.identifier)
+        update_ready!(connection)
+      end
+      true
+    end
+
+    # Touch message (to extend timeout)
+    #
+    # @param message_id [String, Krakow::FrameType::Message]
+    # @return [TrueClass]
+    def touch(message_id)
+      message_id = message_id.message_id if message_id.respond_to?(:message_id)
+      begin
+        distribution.in_flight_lookup(message_id) do |connection|
+          connection.transmit(
+            Command::Touch.new(:message_id => message_id)
+          )
+        end
+        true
+      rescue Error::LookupFailed => e
+        error "Lookup of message for touch failed! <Message ID: #{message_id} - Error: #{e}>"
+        abort e
+      end
+    end
+
+  end
+end