pika_que 0.1.0

data/lib/pika_que/connection.rb

@@ -0,0 +1,47 @@
+module PikaQue
+  class Connection
+    extend Forwardable
+
+    def_delegators :@connection, :create_channel
+
+    include Logging
+
+    attr_reader :connection
+
+    def initialize(opts = {})
+      @opts = PikaQue.config.merge(opts)
+      @opts[:amqp]  = ENV.fetch('RABBITMQ_URL', 'amqp://guest:guest@localhost:5672')
+      @opts[:vhost] = AMQ::Settings.parse_amqp_url(@opts[:amqp]).fetch(:vhost, '/')
+    end
+
+    def self.create(opts = {})
+      new(opts).tap{ |conn| conn.connect! }
+    end
+
+    def connect!
+      @connection ||= Bunny.new(@opts[:amqp], :vhost => @opts[:vhost],
+                              :heartbeat => @opts[:heartbeat],
+                              :properties => @opts.fetch(:properties, {}),
+                              :logger => PikaQue::logger).tap do |conn|
+        conn.start
+      end
+    end
+
+    def connected?
+      @connection && @connection.connected?
+    end
+
+    def disconnect!
+      @connection.close if @connection
+      @connection = nil
+    end
+
+    def ensure_connection
+      unless connected?
+        @connection = nil
+        connect!
+      end
+    end
+    
+  end
+end

data/lib/pika_que/delay_worker.rb

@@ -0,0 +1,55 @@
+module PikaQue
+  class DelayWorker
+
+    attr_accessor :broker, :pool, :queue, :handler
+
+    def initialize(opts = {})
+      @opts = PikaQue.config.merge(opts)
+      @broker = @opts[:broker] || PikaQue::Broker.new(nil, @opts).tap{ |b| b.start }
+      @pool = @opts[:worker_pool] || Concurrent::FixedThreadPool.new(@opts[:concurrency] || 1)
+      @delay_name = "#{@opts[:exchange]}-delay"
+    end
+
+    def prepare
+      @queue = broker.queue(@delay_name, @opts[:queue_options])
+
+      @handler = broker.handler(@opts[:handler_class], @opts[:handler_options])
+      # TODO use routing keys?
+      @handler.bind_queue(@queue, @queue.name)
+    end
+
+    def run
+      @consumer = queue.subscribe(:block => false, :manual_ack => @opts[:ack]) do | delivery_info, metadata, msg |
+        pool.post do
+          work(delivery_info, metadata, msg)
+        end
+      end
+    end
+
+    def start
+      prepare
+      run
+    end
+
+    def stop
+      @consumer.cancel if @consumer
+      @consumer = nil
+
+      unless @opts[:worker_pool]
+        @pool.shutdown
+        @pool.wait_for_termination 12
+      end
+      broker.cleanup
+      broker.stop
+    end
+
+    def work(delivery_info, metadata, msg)
+      handler.handle(:ack, broker.channel, delivery_info, metadata, msg)
+    end
+
+    def logger
+      PikaQue.logger
+    end
+
+  end
+end

data/lib/pika_que/errors.rb

@@ -0,0 +1,5 @@
+module PikaQue
+
+  class SetupError < StandardError; end
+  
+end

data/lib/pika_que/handlers/default_handler.rb

@@ -0,0 +1,31 @@
+module PikaQue
+  module Handlers
+    class DefaultHandler
+     
+      def initialize(opts = {})
+        # nothing to do here
+      end
+
+      def bind_queue(queue, routing_key)
+      end
+
+      def handle(response_code, channel, delivery_info, metadata, msg, error = nil)
+        case response_code
+        when :ack
+          PikaQue.logger.debug "DefaultHandler acknowledge <#{msg}>"
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        when :requeue
+          PikaQue.logger.debug "DefaultHandler requeue <#{msg}>"
+          channel.reject(delivery_info.delivery_tag, true)
+        else
+          PikaQue.logger.debug "DefaultHandler reject <#{msg}>"
+          channel.reject(delivery_info.delivery_tag, false)
+        end
+      end
+
+      def close
+      end
+
+    end    
+  end
+end

data/lib/pika_que/handlers/delay_handler.rb

@@ -0,0 +1,124 @@
+module PikaQue
+  module Handlers
+    class DelayHandler
+
+      # Create following exchanges with delay_name = pika-que-delay
+      # pika-que-delay
+      # pika-que-delay-requeue
+      # and following queues
+      # pika-que-delay-60
+      # pika-que-delay-600
+      # pika-que-delay-3600
+      # pika-que-delay-86400
+      #
+
+      # default delays are 1min, 10min, 1hr, 24hr
+      DEFAULT_DELAY_OPTS = {
+        :delay_periods            => [60, 600, 3600, 86400],
+        :delay_backoff_multiplier => 1000,
+      }.freeze
+
+      def initialize(opts = {})
+        @opts = PikaQue.config.merge(DEFAULT_DELAY_OPTS).merge(opts)
+        @connection = opts[:connection] || PikaQue.connection
+        @channel = @connection.create_channel
+        @delay_monitor = Monitor.new
+        @root_monitor = Monitor.new
+
+        # make sure it is in descending order
+        @delay_periods = @opts[:delay_periods].sort!{ |x,y| y <=> x }
+        @backoff_multiplier = @opts[:delay_backoff_multiplier] # This is for example/dev/test
+
+        @delay_name = "#{@opts[:exchange]}-delay"
+        @requeue_name = "#{@opts[:exchange]}-delay-requeue"
+        @root_name = @opts[:exchange]
+
+        setup_exchanges
+        setup_queues
+      end
+
+      def bind_queue(queue, routing_key)
+        # bind the worker queue to requeue exchange
+        queue.bind(@requeue_exchange, :routing_key => routing_key)
+      end
+
+      def handle(response_code, channel, delivery_info, metadata, msg, error = nil)
+        delay_period = next_delay_period(metadata[:headers])
+        if delay_period > 0
+          # We will publish the message to the delay exchange              
+          PikaQue.logger.info "DelayHandler msg=delaying, delay=#{delay_period}, headers=#{metadata[:headers]}"
+
+          publish_delay(delivery_info, msg, metadata[:headers].merge({ 'delay' => delay_period }))
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        else
+          # Publish the original message with the routing_key to the root exchange
+          work_queue = metadata[:headers]['work_queue']
+          PikaQue.logger.info "DelayHandler msg=publishing, queue=#{work_queue}, headers=#{metadata[:headers]}"
+
+          publish_work(work_queue, msg)
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        end
+      end
+
+      def close
+        @channel.close unless @channel.closed?
+      end
+
+      private
+
+      def setup_exchanges
+        PikaQue.logger.debug "DelayHandler creating exchange=#{@delay_name}"
+        @delay_exchange = @channel.exchange(@delay_name, :type => 'headers', :durable => exchange_durable?)
+
+        PikaQue.logger.debug "DelayHandler creating exchange=#{@requeue_name}"
+        @requeue_exchange = @channel.exchange(@requeue_name, :type => 'topic', :durable => exchange_durable?)
+
+        PikaQue.logger.debug "DelayHandler getting exchange=#{@root_name}"
+        @root_exchange = @channel.exchange(@root_name, :type => 'direct', :durable => exchange_durable?)
+      end
+
+      def setup_queues
+        @delay_periods.each do |t|
+          # Create the queues and bindings
+          PikaQue.logger.debug "DelayHandler creating queue=#{@delay_name}-#{t} x-dead-letter-exchange=#{@requeue_name}"
+          
+          delay_queue = @channel.queue("#{@delay_name}-#{t}", 
+                                      :durable => queue_durable?,
+                                      :arguments => {
+                                        :'x-dead-letter-exchange' => @requeue_name,
+                                        :'x-message-ttl' => t * @backoff_multiplier
+                                      })
+          delay_queue.bind(@delay_exchange, :arguments => { :delay => t })
+        end
+      end
+
+      def queue_durable?
+        @opts.fetch(:queue_options, {}).fetch(:durable, false)
+      end
+
+      def exchange_durable?
+        @opts.fetch(:exchange_options, {}).fetch(:durable, false)
+      end
+
+      def publish_delay(delivery_info, msg, headers)
+        @delay_monitor.synchronize do
+          @delay_exchange.publish(msg, routing_key: delivery_info.routing_key, headers: headers)
+        end
+      end
+
+      def publish_work(routing_key, msg)
+        @root_monitor.synchronize do
+          @root_exchange.publish(msg, routing_key: routing_key)
+        end
+      end
+
+      def next_delay_period(headers)
+        work_at = headers['work_at']
+        t = (work_at - Time.now.to_f).round
+        # greater check is to ignore remainder of time (seconds) smaller than the last delay
+        @delay_periods.bsearch{ |e| t >= e && (t / e.to_f).round > 0 } || 0
+      end
+
+    end
+  end
+end

data/lib/pika_que/handlers/error_handler.rb

@@ -0,0 +1,69 @@
+module PikaQue
+  module Handlers
+    class ErrorHandler
+
+      DEFAULT_ERROR_OPTS = {
+        :exchange     => 'pika-que-error',
+        :exchange_options => { :type => :topic },
+        :queue        => 'pika-que-error',
+        :routing_key  => '#'
+      }.freeze
+     
+      def initialize(opts = {})
+        @opts = PikaQue.config.merge(DEFAULT_ERROR_OPTS).merge(opts)
+        @connection = @opts[:connection] || PikaQue.connection
+        @channel = @connection.create_channel
+        @exchange = @channel.exchange(@opts[:exchange], type: exchange_type, durable: exchange_durable?)
+        @queue = @channel.queue(@opts[:queue], durable: queue_durable?)
+        @queue.bind(@exchange, routing_key: @opts[:routing_key])
+        @monitor = Monitor.new
+      end
+
+      def bind_queue(queue, routing_key)
+      end
+
+      def handle(response_code, channel, delivery_info, metadata, msg, error = nil)
+        case response_code
+        when :ack
+          PikaQue.logger.debug "ErrorHandler acknowledge <#{msg}>"
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        when :reject
+          PikaQue.logger.debug "ErrorHandler reject <#{msg}>"
+          channel.reject(delivery_info.delivery_tag, false)
+        when :requeue
+          PikaQue.logger.debug "ErrorHandler requeue <#{msg}>"
+          channel.reject(delivery_info.delivery_tag, true)
+        else
+          PikaQue.logger.debug "ErrorHandler publishing <#{msg}> to [#{@queue.name}]"
+          publish(delivery_info, msg)
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        end
+      end
+
+      def close
+        @channel.close unless @channel.closed?
+      end
+
+      private
+
+      def queue_durable?
+        @opts.fetch(:queue_options, {}).fetch(:durable, false)
+      end
+
+      def exchange_durable?
+        @opts.fetch(:exchange_options, {}).fetch(:durable, false)
+      end
+
+      def exchange_type
+        @opts.fetch(:exchange_options, {}).fetch(:type, :topic)
+      end
+
+      def publish(delivery_info, msg)
+        @monitor.synchronize do
+          @exchange.publish(msg, routing_key: delivery_info.routing_key)
+        end
+      end
+
+    end    
+  end
+end

data/lib/pika_que/handlers/retry_handler.rb

@@ -0,0 +1,186 @@
+module PikaQue
+  module Handlers
+    class RetryHandler
+
+      # Create following exchanges with retry_prefix = pika-que
+      # pika-que-retry
+      # pika-que-retry-requeue
+      # pika-que-error
+      # and following queues
+      # pika-que-retry-60 (default for const mode)
+      # pika-que-retry-120
+      # pika-que-retry-240
+      # pika-que-retry-480
+      # pika-que-retry-960
+      #
+      # retry_mode can be either :exp or :const
+
+      DEFAULT_RETRY_OPTS = {
+        :retry_prefix             => 'pika-que',
+        :retry_mode               => :exp,
+        :retry_max_times          => 5,
+        :retry_backoff_base       => 0,
+        :retry_backoff_multiplier => 1000,
+        :retry_const_backoff      => 60
+      }.freeze
+
+      def initialize(opts = {})
+        @opts = PikaQue.config.merge(DEFAULT_RETRY_OPTS).merge(opts)
+        @connection = opts[:connection] || PikaQue.connection
+        @channel = @connection.create_channel
+        @retry_monitor = Monitor.new
+        @error_monitor = Monitor.new
+
+        @max_retries = @opts[:retry_max_times]
+        @backoff_base = @opts[:retry_backoff_base]
+        @backoff_multiplier = @opts[:retry_backoff_multiplier] # This is for example/dev/test
+
+        @retry_name = "#{@opts[:retry_prefix]}-retry"
+        @requeue_name = "#{@opts[:retry_prefix]}-retry-requeue"
+        @error_name = "#{@opts[:retry_prefix]}-error"
+
+        setup_exchanges
+        setup_queues
+      end
+
+      def bind_queue(queue, routing_key)
+        # bind the worker queue to requeue exchange
+        queue.bind(@requeue_exchange, :routing_key => routing_key)
+      end
+
+      def handle(response_code, channel, delivery_info, metadata, msg, error = nil)
+        case response_code
+        when :ack
+          PikaQue.logger.debug "RetryHandler acknowledge <#{msg}>"
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        when :reject
+          PikaQue.logger.debug "RetryHandler reject retry <#{msg}>"
+          handle_retry(channel, delivery_info, metadata, msg, :reject)
+        when :requeue
+          PikaQue.logger.debug "RetryHandler requeue <#{msg}>"
+          channel.reject(delivery_info.delivery_tag, true)
+        else
+          PikaQue.logger.debug "RetryHandler error retry <#{msg}>"
+          handle_retry(channel, delivery_info, metadata, msg, error)
+        end
+      end
+
+      def close
+        @channel.close unless @channel.closed?
+      end
+
+
+      #####################################################
+      # formula
+      # base X = 0, 30, 60, 120, 180, etc defaults to 0
+      # (X + 15) * 2 ** (count + 1)
+      def self.backoff_periods(max_retries, backoff_base)
+        (1..max_retries).map{ |c| next_ttl(c, backoff_base) }
+      end
+
+      def self.next_ttl(count, backoff_base)
+        (backoff_base + 15) * 2 ** (count + 1)
+      end
+
+      private
+
+      def setup_exchanges
+        PikaQue.logger.debug "RetryHandler creating exchange=#{@retry_name}"
+        @retry_exchange = @channel.exchange(@retry_name, :type => 'headers', :durable => exchange_durable?)
+        @error_exchange, @requeue_exchange = [@error_name, @requeue_name].map do |name|
+          PikaQue.logger.debug "RetryHandler creating exchange=#{name}"
+          @channel.exchange(name, :type => 'topic', :durable => exchange_durable?)
+        end
+      end
+
+      def setup_queues
+        if @opts[:retry_mode] == :const
+          bo = @opts[:retry_const_backoff]
+          PikaQue.logger.debug "RetryHandler creating queue=#{@retry_name}-#{bo} x-dead-letter-exchange=#{@requeue_name}"
+          backoff_queue = @channel.queue("#{@retry_name}-#{bo}",
+                                        :durable => queue_durable?,
+                                        :arguments => {
+                                          :'x-dead-letter-exchange' => @requeue_name,
+                                          :'x-message-ttl' => bo * @backoff_multiplier
+                                        })
+          backoff_queue.bind(@retry_exchange, :arguments => { :backoff => bo })
+        else
+          backoffs = Expbackoff.backoff_periods(@max_retries, @backoff_base)
+          backoffs.each do |bo|
+            PikaQue.logger.debug "RetryHandler creating queue=#{@retry_name}-#{bo} x-dead-letter-exchange=#{@requeue_name}"
+            backoff_queue = @channel.queue("#{@retry_name}-#{bo}",
+                                          :durable => queue_durable?,
+                                          :arguments => {
+                                            :'x-dead-letter-exchange' => @requeue_name,
+                                            :'x-message-ttl' => bo * @backoff_multiplier
+                                          })
+            backoff_queue.bind(@retry_exchange, :arguments => { :backoff => bo })
+          end
+        end
+
+        PikaQue.logger.debug "RetryHandler creating queue=#{@error_name}"
+        @error_queue = @channel.queue(@error_name, :durable => queue_durable?)
+        @error_queue.bind(@error_exchange, :routing_key => '#')
+      end
+
+      def queue_durable?
+        @opts.fetch(:queue_options, {}).fetch(:durable, false)
+      end
+
+      def exchange_durable?
+        @opts.fetch(:exchange_options, {}).fetch(:durable, false)
+      end
+
+      def handle_retry(channel, delivery_info, metadata, msg, reason)
+        # +1 for the current attempt
+        num_attempts = failure_count(metadata[:headers]) + 1
+        if num_attempts <= @max_retries
+          # Publish message to the x-dead-letter-exchange (ie. retry exchange)
+          PikaQue.logger.info "RetryHandler msg=retrying, count=#{num_attempts}, headers=#{metadata[:headers] || {}}"
+          
+          if @opts[:retry_mode] == :exp
+            backoff_ttl = Expbackoff.next_ttl(num_attempts, @backoff_base)
+          else
+            backoff_ttl = @opts[:retry_const_backoff]
+          end
+
+          publish_retry(delivery_info, msg, { backoff: backoff_ttl, count: num_attempts })
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        else
+          PikaQue.logger.info "RetryHandler msg=failing, retry_count=#{num_attempts}, headers=#{metadata[:headers]}, reason=#{reason}"
+
+          publish_error(delivery_info, msg)
+          channel.acknowledge(delivery_info.delivery_tag, false)
+        end
+      end
+
+      # Uses the header to determine the number of failures this job has
+      # seen in the past. This does not count the current failure. So for
+      # instance, the first time the job fails, this will return 0, the second
+      # time, 1, etc.
+      # @param headers [Hash] Hash of headers that Rabbit delivers as part of
+      #   the message
+      # @return [Integer] Count of number of failures.
+      def failure_count(headers)
+        if headers.nil? || headers['count'].nil?
+          0
+        else
+          headers['count']
+        end
+      end
+
+      def publish_retry(delivery_info, msg, headers)
+        @retry_monitor.synchronize do
+          @retry_exchange.publish(msg, routing_key: delivery_info.routing_key, headers: headers)
+        end
+      end
+
+      def publish_error(delivery_info, msg)
+        @error_monitor.synchronize do
+          @error_exchange.publish(msg, routing_key: delivery_info.routing_key)
+        end
+      end
+
+    end
+  end
+end