nsq-krakow 0.1.0

data/lib/krakow/consumer/queue.rb

@@ -0,0 +1,151 @@
+require 'krakow'
+
+module Krakow
+  # Consume messages from a server
+  class Consumer
+
+    class Queue
+
+      include Celluloid
+      include Utils::Lazy
+
+      # @return [Consumer]
+      attr_reader :consumer
+      # @return [Array] order of message removal
+      attr_reader :pop_order
+      # @return [Symbol] callback method name
+      attr_reader :removal_callback
+
+      # Create new consumer queue instance
+      #
+      # @param consumer [Consumer]
+      # @return [self]
+      def initialize(consumer, *args)
+        opts = args.detect{|x| x.is_a?(Hash)}
+        @consumer = consumer
+        @removal_callback = opts[:removal_callback]
+        @messages = {}
+        @pop_order = []
+        @cleaner = nil
+      end
+
+      # Message container
+      #
+      # @yieldparam [Hash] messages
+      # @return [Hash] messages or block result
+      def messages
+        if(block_given?)
+          yield @messages
+        else
+          @messages
+        end
+      end
+
+      # Register a new connection
+      #
+      # @param connection [Connection]
+      # @return [TrueClass]
+      def register_connection(connection)
+        messages do |collection|
+          collection[connection.identifier] = []
+        end
+        true
+      end
+
+      # Remove connection registration and remove all messages
+      #
+      # @param identifier [String] connection identifier
+      # @return [Array<FrameType::Message>] messages queued for deregistered connection
+      def deregister_connection(identifier)
+        messages do |collection|
+          removed = collection.delete(identifier)
+          pop_order.delete(identifier)
+          removed
+        end
+      end
+
+      # Push new message into queue
+      #
+      # @param message [FrameType::Message]
+      # @return [self]
+      def push(message)
+        unless(message.is_a?(FrameType::Message))
+          abort TypeError.new "Expecting `FrameType::Message` but received `#{message.class}`!"
+        end
+        messages do |collection|
+          begin
+            collection[message.connection.identifier] << message
+            pop_order << message.connection.identifier
+          rescue Celluloid::DeadActorError
+            abort Error::ConnectionUnavailable.new
+          end
+        end
+        signal(:new_message)
+        current_actor
+      end
+      alias_method :<<, :push
+      alias_method :enq, :push
+
+      # Pop first item off the queue
+      #
+      # @return [Object]
+      def pop
+        message = nil
+        until(message)
+          wait(:new_message) if pop_order.empty?
+          messages do |collection|
+            key = pop_order.shift
+            if(key)
+              message = collection[key].shift
+              message = validate_message(message)
+            end
+          end
+        end
+        message
+      end
+      alias_method :deq, :pop
+
+      # @return [Integer] number of queued messages
+      def size
+        messages do |collection|
+          collection.values.map(&:size).inject(&:+)
+        end
+      end
+
+      # Remove duplicate message from queue if possible
+      #
+      # @param message [FrameType::Message]
+      # @return [TrueClass, FalseClass]
+      def scrub_duplicate_message(message)
+        messages do |collection|
+          idx = collection[message.connection.identifier].index do |msg|
+            msg.message_id == message.message_id
+          end
+          if(idx)
+            msg = collection[message.connection.identifier].delete_at(idx)
+            if(removal_callback)
+              consumer.send(removal_callback, [message])
+            end
+            true
+          else
+            false
+          end
+        end
+      end
+
+      # Validate message
+      def validate_message(message)
+        if(message.instance_stamp > message.instance_stamp + (message.connection.endpoint_settings[:msg_timeout] / 1000.0))
+          warn "Message exceeded timeout! Discarding. (#{message})"
+          if(removal_callback)
+            consumer.send(removal_callback, [message])
+          end
+          nil
+        else
+          message
+        end
+      end
+
+    end
+  end
+end

data/lib/krakow/discovery.rb

@@ -0,0 +1,57 @@
+require 'uri'
+require 'http'
+require 'multi_json'
+require 'krakow'
+
+module Krakow
+
+  # Provides queue topic discovery
+  class Discovery
+
+    include Utils::Lazy
+
+    # @!group Attributes
+
+    # @!macro [attach] attribute
+    #   @!method $1
+    #     @return [$2] the $1 $0
+    #   @!method $1?
+    #     @return [TrueClass, FalseClass] truthiness of the $1 $0
+    attribute :nsqlookupd, [Array, String], :required => true
+
+    # @!endgroup
+
+    # Get list of end points with given topic name available
+    #
+    # @param topic [String] topic name
+    # @return [Array<Hash>]
+    def lookup(topic)
+      result = [nsqlookupd].flatten.map do |location|
+        uri = URI.parse(location)
+        uri.path = '/lookup'
+        uri.query = "topic=#{topic}&ts=#{Time.now.to_i}"
+        begin
+          debug "Requesting lookup for topic #{topic} - #{uri}"
+          content = HTTP.with(:accept => 'application/octet-stream').get(uri.to_s)
+          unless(content.respond_to?(:to_hash))
+            data = MultiJson.load(content.to_s)
+          else
+            data = content.to_hash
+          end
+          debug "Lookup response (#{uri.to_s}): #{data.inspect}"
+          if(data['data'] && data['data']['producers'])
+            data['data']['producers'].map do |producer|
+              Hash[*producer.map{|k,v| [k.to_sym, v]}.flatten]
+            end
+          end
+        rescue => e
+          warn "Lookup exception encountered: #{e.class.name} - #{e}"
+          nil
+        end
+      end.compact.flatten(1).uniq
+      debug "Discovery lookup result: #{result.inspect}"
+      result
+    end
+
+  end
+end

data/lib/krakow/distribution.rb

@@ -0,0 +1,229 @@
+require 'krakow'
+
+module Krakow
+  # Message distribution
+  # @abstract
+  class Distribution
+
+    autoload :Default, 'krakow/distribution/default'
+#    autoload :ProducerWeighted, 'krakow/distribution/producer_weighted'
+#    autoload :ConsumerWeighted, 'krakow/distribution/consumer_weighted'
+
+    include Celluloid
+    include Utils::Lazy
+    # @!parse include Krakow::Utils::Lazy::InstanceMethods
+    # @!parse extend Krakow::Utils::Lazy::ClassMethods
+
+    attr_accessor :ideal, :flight_record, :registry
+
+    # @!group Attributes
+
+    # @!macro [attach] attribute
+    #   @!method $1
+    #     @return [$2] the $1 $0
+    #   @!method $1?
+    #     @return [TrueClass, FalseClass] truthiness of the $1 $0
+    attribute :consumer, Krakow::Consumer, :required => true
+    attribute :watch_dog_interval, Numeric, :default => 1.0
+    attribute :backoff_interval, Numeric
+    attribute :max_in_flight, Integer, :default => 1
+
+    # @!endgroup
+
+    def initialize(args={})
+      super
+      @ideal = 0
+      @flight_record = {}
+      @registry = {}
+    end
+
+    # [Abstract] Reset flight distributions
+    def redistribute!
+      raise NotImplementedError.new 'Custom `#redistrubute!` method must be provided!'
+    end
+
+    # [Abstract] Determine RDY value for given connection
+    # @param connection_identifier [String]
+    # @return [Integer]
+    def calculate_ready!(connection_identifier)
+      raise NotImplementedError.new 'Custom `#calculate_ready!` method must be provided!'
+    end
+
+    # Remove message metadata from registry
+    #
+    # @param message [Krakow::FrameType::Message, String] message or ID
+    # @return [Krakow::Connection, NilClass]
+    def unregister_message(message)
+      msg_id = message.respond_to?(:message_id) ? message.message_id : message.to_s
+      connection = connection_lookup(flight_record[msg_id])
+      flight_record.delete(msg_id)
+      if(connection)
+        begin
+          ident = connection.identifier
+          registry_info = registry_lookup(ident)
+          registry_info[:in_flight] -= 1
+          calculate_ready!(ident)
+          connection
+        rescue Celluloid::DeadActorError
+          warn 'Connection is dead. No recalculation applied on ready.'
+        end
+      else
+        warn 'No connection associated to message via lookup. No recalculation applied on ready.'
+      end
+    end
+
+    # Return the currently configured RDY value for given connnection
+    #
+    # @param connection_identifier [String]
+    # @return [Integer]
+    def ready_for(connection_identifier)
+      registry_lookup(connection_identifier)[:ready]
+    end
+
+
+    # Send RDY for given connection
+    #
+    # @param connection [Krakow::Connection]
+    # @return [Krakow::FrameType::Error,nil]
+    def set_ready_for(connection, *_)
+      connection.transmit(
+        Command::Rdy.new(
+          :count => ready_for(connection.identifier)
+        )
+      )
+    end
+
+    # Initial ready value used for new connections
+    #
+    # @return [Integer]
+    def initial_ready
+      ideal > 0 ? 1 : 0
+    end
+
+    # Registers message into registry and configures for distribution
+    #
+    # @param message [FrameType::Message]
+    # @param connection_identifier [String]
+    # @return [Integer]
+    def register_message(message, connection_identifier)
+      if(flight_record[message.message_id])
+        abort KeyError.new "Message is already registered in flight record! (#{message.message_id})"
+      else
+        registry_info = registry_lookup(connection_identifier)
+        registry_info[:in_flight] += 1
+        flight_record[message.message_id] = connection_identifier
+        calculate_ready!(connection_identifier)
+      end
+    end
+
+    # Add connection to make available for RDY distribution
+    #
+    # @param connection [Krakow::Connection]
+    # @return [TrueClass]
+    def add_connection(connection)
+      unless(registry[connection.identifier])
+        registry[connection.identifier] = {
+          :ready => initial_ready,
+          :in_flight => 0,
+          :failures => 0,
+          :backoff_until => 0
+        }
+      end
+      true
+    end
+
+    # Remove connection from RDY distribution
+    #
+    # @param connection_identifier [String]
+    # @return [TrueClass]
+    def remove_connection(connection_identifier, *args)
+      # remove connection from registry
+      registry.delete(connection_identifier)
+      # remove any in flight messages
+      flight_record.delete_if do |k,v|
+        if(v == connection_identifier)
+          warn "Removing in flight reference due to failed connection: #{v}"
+          true
+        end
+      end
+      true
+    end
+
+    # Return connection associated with given registry key
+    #
+    # @param identifier [String] connection identifier
+    # @return [Krakow::Connection, nil]
+    def connection_lookup(identifier)
+      consumer.connection(identifier)
+    end
+
+    # Return source connection for given message ID
+    #
+    # @param msg_id [String]
+    # @yield execute with connection
+    # @yieldparam connection [Krakow::Connection]
+    # @return [Krakow::Connection, Object]
+    def in_flight_lookup(msg_id)
+      connection = connection_lookup(flight_record[msg_id])
+      unless(connection)
+        abort Krakow::Error::LookupFailed.new("Failed to locate in flight message (ID: #{msg_id})")
+      end
+      if(block_given?)
+        begin
+          yield connection
+        rescue => e
+          abort e
+        end
+      else
+        connection
+      end
+    end
+
+    # Return registry information for given connection
+    # @param connection_identifier [String]
+    # @return [Hash] registry information
+    # @raise [Krakow::Error::LookupFailed]
+    def registry_lookup(connection_identifier)
+      registry[connection_identifier] ||
+        abort(Krakow::Error::LookupFailed.new("Failed to locate connection information in registry (#{connection_identifier})"))
+    end
+
+    # @return [Array<Krakow::Connection>] connections in registry
+    def connections
+      registry.keys.map do |identifier|
+        connection_lookup(identifier)
+      end.compact
+    end
+
+    # Log failure of processed message
+    #
+    # @param connection_identifier [String]
+    # @return [TrueClass]
+    def failure(connection_identifier)
+      if(backoff_interval)
+        registry_info = registry_lookup(connection_identifier)
+        registry_info[:failures] += 1
+        registry_info[:backoff_until] = Time.now.to_i + (registry_info[:failures] * backoff_interval)
+      end
+      true
+    end
+
+    # Log success of processed message
+    #
+    # @param connection_identifier [String]
+    # @return [TrueClass]
+    def success(connection_identifier)
+      if(backoff_interval)
+        registry_info = registry_lookup(connection_identifier)
+        if(registry_info[:failures] > 1)
+          registry_info[:failures] -= 1
+          registry_info[:backoff_until] = Time.now.to_i + (registry_info[:failures] * backoff_interval)
+        else
+          registry_info[:failures] = 0
+        end
+      end
+      true
+    end
+
+  end
+end

data/lib/krakow/distribution/default.rb

@@ -0,0 +1,159 @@
+require 'krakow'
+
+module Krakow
+  class Distribution
+    # Default distribution implementation. This uses a round-robin
+    # approach for less than ideal states.
+    class Default < Distribution
+
+      attr_reader :less_than_ideal_stack, :watch_dog
+
+      # recalculate `ideal` and update RDY on connections
+      def redistribute!
+        @ideal = registry.size < 1 ? 0 : max_in_flight / registry.size
+        debug "Distribution calculated ideal: #{ideal}"
+        if(less_than_ideal?)
+          registry.each do |connection_id, reg_info|
+            reg_info[:ready] = 0
+          end
+          max_in_flight.times do
+            less_than_ideal_ready!
+          end
+          connections.each do |connection|
+            set_ready_for(connection, :force)
+          end
+          watch_dog.cancel if watch_dog
+          @watch_dog = every(watch_dog_interval) do
+            force_unready
+          end
+        else
+          if(watch_dog)
+            watch_dog.cancel
+            @watch_dog = nil
+          end
+          connections.each do |connection|
+            current_ready = ready_for(connection.identifier)
+            calculate_ready!(connection.identifier)
+            unless(current_ready == ready_for(connection.identifier))
+              debug "Redistribution ready setting update for connection #{connection}"
+              set_ready_for(connection)
+            end
+          end
+        end
+      end
+
+      # Is ideal less than 1
+      #
+      # @return [TrueClass, FalseClass]
+      def less_than_ideal?
+        ideal < 1
+      end
+
+      # Find next connection to receive RDY count
+      #
+      # @return [Krakow::Connection, nil]
+      def less_than_ideal_ready!
+        admit_defeat = false
+        connection = nil
+        until(connection || (admit_defeat && less_than_ideal_stack.empty?))
+          if(less_than_ideal_stack.nil? || less_than_ideal_stack.empty?)
+            @less_than_ideal_stack = waiting_connections
+            admit_defeat = true
+          end
+          con = less_than_ideal_stack.pop
+          if(con)
+            unless(registry_lookup(con.identifier)[:backoff_until] > Time.now.to_i)
+              connection = con
+            end
+          end
+        end
+        if(connection)
+          registry_lookup(connection.identifier)[:ready] = 1
+          connection
+        end
+      end
+
+      # Adds extra functionality to provide round robin RDY setting
+      # when in less than ideal state
+      #
+      # @param connection [Krakow::Connection]
+      # @param args [Symbol]
+      # @return [Krakow::FrameType::Error, nil]
+      def set_ready_for(connection, *args)
+        super connection
+        if(less_than_ideal? && !args.include?(:force))
+          debug "RDY set ignored due to less than ideal state (con: #{connection})"
+          con = less_than_ideal_ready!
+          if(con)
+            watch_dog.reset if watch_dog
+            super con
+          else
+            warn 'Failed to set RDY state while less than ideal. Connection stack is empty!'
+          end
+        end
+      end
+
+      # Update connection ready count
+      # @param connection_identifier [String]
+      # @return [Integer, nil]
+      def calculate_ready!(connection_identifier)
+        begin
+          registry_info = registry_lookup(connection_identifier)
+          unless(less_than_ideal?)
+            registry_info[:ready] = ideal - registry_info[:in_flight]
+            if(registry_info[:ready] < 0 || registry_info[:backoff_until] > Time.now.to_i)
+              registry_info[:ready] = 0
+              registry_info[:backoff_timer].cancel if registry[:backoff_timer]
+              registry_info[:backoff_timer] = after(registry_info[:backoff_until] - Time.now.to_i) do
+                calculate_ready!(connection_identifier)
+                set_ready_for(connection_lookup(connection_identifier)) unless less_than_ideal?
+              end
+            end
+            registry_info[:ready]
+          else
+            registry_info[:ready] = 0
+          end
+        rescue Error::ConnectionFailure
+          warn 'Failed connection encountered!'
+        rescue Error::ConnectionUnavailable
+          warn 'Unavailable connection encountered!'
+        end
+      end
+
+      # All connections without RDY state
+      #
+      # @return [Array<Krakow::Connection>]
+      def waiting_connections
+        registry.find_all do |conn_id, info|
+          info[:ready] < 1 && info[:in_flight] < 1 && info[:backoff_until] < Time.now.to_i
+        end.map{|conn_id, info| connection_lookup(conn_id) }.compact
+      end
+
+      # All connections with RDY state
+      #
+      # @return [Array<Krakow::Connection>]
+      def rdy_connections
+        registry.find_all do |conn_id, info|
+          info[:ready] > 0
+        end.map{|conn_id, info| connection_lookup(conn_id) }.compact
+      end
+
+      # Force a connection to give up RDY state so next in stack can receive
+      #
+      # @return [nil]
+      def force_unready
+        debug 'Forcing a connection into an unready state due to less than ideal state'
+        connection = rdy_connections.shuffle.first
+        if(connection)
+          debug "Stripping RDY state from connection: #{connection}"
+          calculate_ready!(connection.identifier)
+          set_ready_for(connection)
+        else
+          warn "Failed to locate available connection for RDY aquisition!"
+        end
+        nil
+      end
+
+    end
+  end
+end