upperkut 0.7.2 → 1.0.0.rc

data/lib/upperkut/redis_pool.rb

@@ -8,12 +8,12 @@ module Upperkut
       size: 2, # pool related option
       connect_timeout: 0.2,
       read_timeout: 5.0,
-      write_timeout: 0.5,
-      url: ENV['REDIS_URL']
+      write_timeout: 0.5
     }.freeze
 
     def initialize(options)
-      @options      = DEFAULT_OPTIONS.merge(options)
+      @options      = DEFAULT_OPTIONS.merge(url: ENV['REDIS_URL'])
+                                     .merge(options)
 
       # Extract pool related options
       @size         = @options.delete(:size)

data/lib/upperkut/strategies/base.rb

@@ -24,6 +24,20 @@ module Upperkut
         raise NotImplementedError
       end
 
+      # Public: Confirms that items have been processed successfully.
+      #
+      # items - The Array of items do be confirmed.
+      def ack(_items)
+        raise NotImplementedError
+      end
+
+      # Public: Informs that items have been not processed successfully and therefore must be re-processed.
+      #
+      # items - The Array of items do be unacknowledged.
+      def nack(_items)
+        raise NotImplementedError
+      end
+
       # Public: Tells when to execute the event processing,
       # when this condition is met so the events are dispatched to
       # the worker.

data/lib/upperkut/strategies/buffered_queue.rb

@@ -7,14 +7,65 @@ module Upperkut
     class BufferedQueue < Upperkut::Strategies::Base
       include Upperkut::Util
 
+      DEQUEUE_ITEMS = %(
+        local key = KEYS[1]
+        local waiting_ack_key = KEYS[2]
+        local batch_size = ARGV[1]
+        local current_timestamp = ARGV[2]
+        local expired_ack_timestamp = ARGV[3] + 1
+
+        -- move expired items back to the queue
+        local expired_ack_items = redis.call("ZRANGEBYSCORE", waiting_ack_key, 0, expired_ack_timestamp)
+        if table.getn(expired_ack_items) > 0 then
+          redis.call("ZREMRANGEBYSCORE", waiting_ack_key, 0, expired_ack_timestamp)
+          for i, item in ipairs(expired_ack_items) do
+            redis.call("RPUSH", key, item)
+          end
+        end
+
+        -- now fetch a batch
+        local items = redis.call("LRANGE", key, 0, batch_size - 1)
+        for i, item in ipairs(items) do
+          redis.call("ZADD", waiting_ack_key, current_timestamp + tonumber('0.' .. i), item)
+        end
+        redis.call("LTRIM", key, batch_size, -1)
+
+        return items
+      ).freeze
+
+      ACK_ITEMS = %(
+        local waiting_ack_key = KEYS[1]
+        local items = ARGV
+
+        for i, item in ipairs(items) do
+          redis.call("ZREM", waiting_ack_key, item)
+        end
+      ).freeze
+
+      NACK_ITEMS = %(
+        local key = KEYS[1]
+        local waiting_ack_key = KEYS[2]
+        local items = ARGV
+
+        for i, item in ipairs(items) do
+          redis.call("ZREM", waiting_ack_key, item)
+          redis.call("RPUSH", key, item)
+        end
+      ).freeze
+
       attr_reader :options
 
       def initialize(worker, options = {})
         @options = options
         @redis_options = options.fetch(:redis, {})
-        @redis_pool = setup_redis_pool
-        @worker     = worker
-        @max_wait   = options.fetch(
+        @worker = worker
+
+        @ack_wait_limit = options.fetch(
+          :ack_wait_limit,
+          Integer(ENV['UPPERKUT_ACK_WAIT_LIMIT'] || 120)
+        )
+
+        @max_wait = options.fetch(
           :max_wait,
           Integer(ENV['UPPERKUT_MAX_WAIT'] || 20)
         )
@@ -28,23 +79,23 @@ module Upperkut
       end
 
       def push_items(items = [])
-        items = [items] if items.is_a?(Hash)
+        items = normalize_items(items)
         return false if items.empty?
 
         redis do |conn|
-          conn.rpush(key, encode_json_items(items))
+          conn.rpush(key, items.map(&:to_json))
         end
 
         true
       end
 
       def fetch_items
-        stop = [@batch_size, size].min
+        batch_size = [@batch_size, size].min
 
         items = redis do |conn|
-          conn.multi do
-            stop.times { conn.lpop(key) }
-          end
+          conn.eval(DEQUEUE_ITEMS,
+                    keys: [key, processing_key],
+                    argv: [batch_size, Time.now.utc.to_i, Time.now.utc.to_i - @ack_wait_limit])
         end
 
         decode_json_items(items)
@@ -54,11 +105,24 @@ module Upperkut
         redis { |conn| conn.del(key) }
       end
 
-      def metrics
-        {
-          'latency' => latency,
-          'size' => size
-        }
+      def ack(items)
+        raise ArgumentError, 'Invalid item' unless items.all? { |item| item.is_a?(Item) }
+
+        redis do |conn|
+          conn.eval(ACK_ITEMS,
+                    keys: [processing_key],
+                    argv: items.map(&:to_json))
+        end
+      end
+
+      def nack(items)
+        raise ArgumentError, 'Invalid item' unless items.all? { |item| item.is_a?(Item) }
+
+        redis do |conn|
+          conn.eval(NACK_ITEMS,
+                    keys: [key, processing_key],
+                    argv: items.map(&:to_json))
+        end
       end
 
       def process?
@@ -73,45 +137,81 @@ module Upperkut
         end
       end
 
+      def metrics
+        current_latency = latency
+
+        {
+          'latency' => current_latency,
+          'oldest_unacked_item_age' => oldest_item_age(current_latency),
+          'size' => size
+        }
+      end
+
       private
 
+      def key
+        "upperkut:buffers:#{to_underscore(@worker.name)}"
+      end
+
+      def processing_key
+        "#{key}:processing"
+      end
+
       def fulfill_condition?(buff_size)
         return false if buff_size.zero?
 
         buff_size >= @batch_size || @waiting_time >= @max_wait
       end
 
-      def size
-        redis do |conn|
-          conn.llen(key)
+      def oldest_item_age(current_latency)
+        oldest_processing_item = redis do |conn|
+          items = conn.zrange(processing_key, 0, 0)
+          decode_json_items(items).first
         end
+
+        oldest_processing_age = if oldest_processing_item
+                                  now = Time.now.to_f
+                                  now - oldest_processing_item.enqueued_at.to_f
+                                else
+                                  0
+                                end
+
+        [current_latency, oldest_processing_age].max
       end
 
       def latency
-        item = redis { |conn| conn.lrange(key, 0, 0) }
-        item = decode_json_items(item).first
-        return 0 unless item
+        items = redis { |conn| conn.lrange(key, 0, 0) }
+        first_item = decode_json_items(items).first
+        return 0 unless first_item
 
         now = Time.now.to_f
-        now - item.fetch('enqueued_at', Time.now).to_f
+        now - first_item.enqueued_at.to_f
       end
 
-      def setup_redis_pool
-        return @redis_options if @redis_options.is_a?(ConnectionPool)
-
-        RedisPool.new(options.fetch(:redis, {})).create
+      def size
+        redis do |conn|
+          conn.llen(key)
+        end
       end
 
       def redis
         raise ArgumentError, 'requires a block' unless block_given?
 
-        @redis_pool.with do |conn|
-          yield conn
+        retry_block do
+          redis_pool.with do |conn|
+            yield conn
+          end
         end
       end
 
-      def key
-        "upperkut:buffers:#{to_underscore(@worker.name)}"
+      def redis_pool
+        @redis_pool ||= begin
+                          if @redis_options.is_a?(ConnectionPool)
+                            @redis_options
+                          else
+                            RedisPool.new(@redis_options).create
+                          end
+                        end
       end
     end
   end

data/lib/upperkut/strategies/priority_queue.rb

@@ -0,0 +1,217 @@
+require 'upperkut/util'
+require 'upperkut/redis_pool'
+require 'upperkut/strategies/base'
+
+module Upperkut
+  module Strategies
+    # Public: Queue that prevent a single tenant from taking over.
+    class PriorityQueue < Upperkut::Strategies::Base
+      include Upperkut::Util
+
+      ONE_DAY_IN_SECONDS = 86400
+
+      # Logic as follows:
+      #
+      # We keep the last score used for each tenant key. One tenant_key is
+      #  an tenant unique id. To calculate the next_score we use
+      #  max(current_tenant_score, current_global_score) + increment we store
+      #  the queue in a sorted set using the next_score as ordering key if one
+      #  tenant sends lots of messages, this tenant ends up with lots of
+      #  messages in the queue spaced by increment if another tenant then
+      #  sends a message, since it previous_tenant_score is lower than the
+      #  first tenant, it will be inserted before it in the queue.
+      #
+      # In other words, the idea of this queue is to not allowing an tenant
+      #  that sends a lot of messages to dominate processing and give a chance
+      #  for tenants that sends few messages to have a fair share of
+      #  processing time.
+      ENQUEUE_ITEM = %(
+        local increment = 1
+        local checkpoint_key = KEYS[1]
+        local counter_key = KEYS[2]
+        local score_key = KEYS[3]
+        local queue_key = KEYS[4]
+        local current_checkpoint = tonumber(redis.call("GET", checkpoint_key)) or 0
+        local current_counter = tonumber(redis.call("INCR", counter_key))
+        local current_score = tonumber(redis.call("GET", score_key)) or 0
+        local next_score = nil
+
+        if current_score >= current_checkpoint then
+          next_score = current_score + increment
+        else
+          next_score = current_checkpoint + increment
+        end
+
+        redis.call("SETEX", score_key, #{ONE_DAY_IN_SECONDS}, next_score)
+        redis.call("ZADD", queue_key, next_score + tonumber('0.' .. current_counter), ARGV[1])
+
+        return next_score
+      ).freeze
+
+      # Uses ZPOP* functions available only on redis 5.0.0+
+      DEQUEUE_ITEM = %(
+        local checkpoint_key = KEYS[1]
+        local queue_key = KEYS[2]
+        local batch_size = ARGV[1]
+        local popped_items = redis.call("ZPOPMIN", queue_key, batch_size)
+        local items = {}
+        local last_score = 0
+
+        for i, v in ipairs(popped_items) do
+          if i % 2 == 1 then
+            table.insert(items, v)
+          else
+            last_score = v
+          end
+        end
+
+        redis.call("SETEX", checkpoint_key, 86400, last_score)
+        return items
+      ).freeze
+
+      def initialize(worker, options)
+        @worker = worker
+        @options = options
+        @priority_key = options.fetch(:priority_key)
+        @redis_options = options.fetch(:redis, {})
+
+        @max_wait = options.fetch(
+          :max_wait,
+          Integer(ENV['UPPERKUT_MAX_WAIT'] || 20)
+        )
+
+        @batch_size = options.fetch(
+          :batch_size,
+          Integer(ENV['UPPERKUT_BATCH_SIZE'] || 1000)
+        )
+
+        @waiting_time = 0
+
+        raise ArgumentError, 'Invalid priority_key. ' \
+          'Must be a lambda' unless @priority_key.respond_to?(:call)
+      end
+
+      # Public: Ingests the event into strategy.
+      #
+      # items - The Array of items do be inserted.
+      #
+      # Returns true when success, raise when error.
+      def push_items(items = [])
+        items = normalize_items(items)
+        return false if items.empty?
+
+        redis do |conn|
+          items.each do |item|
+            priority_key = @priority_key.call(item)
+            score_key = "#{queue_key}:#{priority_key}:score"
+
+            keys = [checkpoint_key,
+                    counter_key,
+                    score_key,
+                    queue_key]
+
+            conn.eval(ENQUEUE_ITEM,
+                      keys: keys,
+                      argv: [item.to_json])
+          end
+        end
+
+        true
+      end
+
+      # Public: Retrieve events from Strategy.
+      #
+      # Returns an Array containing events as hash.
+      def fetch_items
+        batch_size = [@batch_size, size].min
+
+        items = redis do |conn|
+          conn.eval(DEQUEUE_ITEM,
+                    keys: [checkpoint_key, queue_key],
+                    argv: [batch_size])
+        end
+
+        decode_json_items(items)
+      end
+
+      # Public: Clear all data related to the strategy.
+      def clear
+        redis { |conn| conn.del(queue_key) }
+      end
+
+      def ack(_items); end
+
+      def nack(items)
+        push_items(items)
+      end
+
+      # Public: Tells when to execute the event processing,
+      # when this condition is met so the events are dispatched to
+      # the worker.
+      def process?
+        if fulfill_condition?(size)
+          @waiting_time = 0
+          return true
+        end
+
+        @waiting_time += @worker.setup.polling_interval
+        false
+      end
+
+      # Public: Consolidated strategy metrics.
+      #
+      # Returns hash containing metric name and values.
+      def metrics
+        {
+          'size' => size
+        }
+      end
+
+      private
+
+      def checkpoint_key
+        "#{queue_key}:checkpoint"
+      end
+
+      def counter_key
+        "#{queue_key}:counter"
+      end
+
+      def queue_key
+        "upperkut:priority_queue:#{to_underscore(@worker.name)}"
+      end
+
+      def fulfill_condition?(buff_size)
+        return false if buff_size.zero?
+
+        buff_size >= @batch_size || @waiting_time >= @max_wait
+      end
+
+      def size
+        redis do |conn|
+          conn.zcard(queue_key)
+        end
+      end
+
+      def redis
+        raise ArgumentError, 'requires a block' unless block_given?
+
+        retry_block do
+          redis_pool.with do |conn|
+            yield conn
+          end
+        end
+      end
+
+      def redis_pool
+        @redis_pool ||= begin
+                          if @redis_options.is_a?(ConnectionPool)
+                            @redis_options
+                          else
+                            RedisPool.new(@options.fetch(:redis, {})).create
+                          end
+                        end
+      end
+    end
+  end
+end