upperkut 1.0.3 → 1.0.4

data/lib/upperkut/strategies/buffered_queue.rb

@@ -0,0 +1,218 @@
+require 'upperkut/util'
+require 'upperkut/redis_pool'
+require 'upperkut/strategies/base'
+
+module Upperkut
+  module Strategies
+    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, {})
+        @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)
+        )
+
+        @batch_size = options.fetch(
+          :batch_size,
+          Integer(ENV['UPPERKUT_BATCH_SIZE'] || 1000)
+        )
+
+        @waiting_time = 0
+      end
+
+      def push_items(items = [])
+        items = normalize_items(items)
+        return false if items.empty?
+
+        redis do |conn|
+          conn.rpush(key, encode_json_items(items))
+        end
+
+        true
+      end
+
+      def fetch_items
+        batch_size = [@batch_size, size].min
+
+        items = redis do |conn|
+          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)
+      end
+
+      def clear
+        redis { |conn| conn.del(key) }
+      end
+
+      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: encode_json_items(items))
+        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: encode_json_items(items))
+        end
+      end
+
+      def process?
+        buff_size = size
+
+        if fulfill_condition?(buff_size)
+          @waiting_time = 0
+          return true
+        else
+          @waiting_time += @worker.setup.polling_interval
+          return false
+        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 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
+        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 - first_item.enqueued_at.to_f
+      end
+
+      def size
+        redis do |conn|
+          conn.llen(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(@redis_options).create
+                          end
+                        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: [encode_json_items(item)])
+          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

data/lib/upperkut/strategies/scheduled_queue.rb

@@ -0,0 +1,162 @@
+require 'time'
+require 'upperkut/util'
+require 'upperkut/redis_pool'
+require 'upperkut/strategies/base'
+
+module Upperkut
+  module Strategies
+    # Public: Encapsulates methods required to build a Scheculed Queue
+    # Items are queued, but are only fetched at a specific point in time.
+    class ScheduledQueue < Upperkut::Strategies::Base
+      include Upperkut::Util
+
+      ZPOPBYRANGE = %(
+        local score_from = ARGV[1]
+        local score_to = ARGV[2]
+        local limit = ARGV[3]
+
+        local values = redis.call('zrangebyscore', KEYS[1], score_from, score_to, 'LIMIT', '0', limit)
+
+        if table.getn(values) > 0 then
+          redis.call('zrem', KEYS[1], unpack(values))
+        end
+
+        return values
+      ).freeze
+
+      attr_reader :options
+
+      def initialize(worker, options = {})
+        @options = options
+        @redis_options = @options.fetch(:redis, {})
+        @worker = worker
+
+        @batch_size = @options.fetch(
+          :batch_size,
+          Integer(ENV['UPPERKUT_BATCH_SIZE'] || 1000)
+        )
+      end
+
+      def push_items(items = [])
+        items = normalize_items(items)
+        return false if items.empty?
+
+        redis do |conn|
+          items.each do |item|
+            schedule_item = ensure_timestamp_attr(item)
+            timestamp = schedule_item.body['timestamp']
+            conn.zadd(key, timestamp, encode_json_items(schedule_item))
+          end
+        end
+
+        true
+      end
+
+      def fetch_items
+        args = {
+          value_from: '-inf'.freeze,
+          value_to: Time.now.utc.to_f.to_s,
+          limit: @batch_size
+        }
+        items = []
+
+        redis do |conn|
+          items = pop_values(conn, args)
+        end
+
+        decode_json_items(items)
+      end
+
+      def clear
+        redis { |conn| conn.del(key) }
+      end
+
+      def ack(_items); end
+
+      def nack(items)
+        push_items(items)
+      end
+
+      def metrics
+        {
+          'latency' => latency,
+          'size' => size
+        }
+      end
+
+      def process?
+        buff_size = size('-inf', Time.now.utc.to_i)
+        return true if fulfill_condition?(buff_size)
+
+        false
+      end
+
+      private
+
+      def key
+        "upperkut:queued:#{to_underscore(@worker.name)}"
+      end
+
+      def ensure_timestamp_attr(item)
+        return item if item.body.key?('timestamp')
+
+        Item.new(
+          id: item.id,
+          body: item.body.merge('timestamp' => Time.now.utc.to_i),
+          enqueued_at: item.enqueued_at
+        )
+      end
+
+      def pop_values(redis_client, args)
+        value_from = args[:value_from]
+        value_to = args[:value_to]
+        limit = args[:limit]
+        redis_client.eval(ZPOPBYRANGE, keys: [key], argv: [value_from, value_to, limit])
+      end
+
+      def fulfill_condition?(buff_size)
+        !buff_size.zero?
+      end
+
+      def size(min = '-inf', max = '+inf')
+        redis do |conn|
+          conn.zcount(key, min, max)
+        end
+      end
+
+      def latency
+        now = Time.now.utc
+        timestamp = now.to_f
+
+        item = redis do |conn|
+          item = conn.zrangebyscore(key, '-inf', timestamp.to_s, limit: [0, 1]).first
+          decode_json_items([item]).first
+        end
+
+        return timestamp - item.body['timestamp'].to_f if item
+
+        0
+      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(@redis_options).create
+                          end
+                        end
+      end
+    end
+  end
+end

data/lib/upperkut/util.rb

@@ -0,0 +1,73 @@
+require 'json'
+require 'upperkut/item'
+
+module Upperkut
+  module Util
+    def to_underscore(object)
+      klass_name = object.dup
+      klass_name.gsub!(/::/, '_')
+      klass_name.gsub!(/([A-Z\d]+)([A-Z][a-z])/, '\1_\2')
+      klass_name.gsub!(/([a-z\d])([A-Z])/, '\1_\2')
+      klass_name.tr!('-', '_')
+      klass_name.downcase!
+      klass_name
+    end
+
+    # Public:
+    #  Normalize hash and hash arrays into a hash of Items.
+    #  An Item object contains metadata, for example the timestamp from the moment it was enqueued,
+    #   that we need to carry through multiple execution tries.
+    #
+    #  When the execution fails, we need to schedule the whole batch for retry, and scheduling
+    #   an Item will make Upperkut understand that we're not dealing with a new batch,
+    #   so metrics like latency will increase.
+    def normalize_items(items)
+      items = [items] unless items.is_a?(Array)
+
+      items.map do |item|
+        next item if item.is_a?(Item)
+
+        Item.new(id: SecureRandom.uuid, body: item)
+      end
+    end
+
+    def encode_json_items(items)
+      items = [items] unless items.is_a?(Array)
+
+      items.map do |item|
+        JSON.generate(
+          'id' => item.id,
+          'body' => item.body,
+          'enqueued_at' => item.enqueued_at
+        )
+      end
+    end
+
+    def decode_json_items(items)
+      items.each_with_object([]) do |item_json, memo|
+        next unless item_json
+
+        hash = JSON.parse(item_json)
+        id, body, enqueued_at = hash.values_at('id', 'body', 'enqueued_at')
+        memo << Item.new(id: id, body: body, enqueued_at: enqueued_at)
+      end
+    end
+
+    def retry_block(retries_limit = 3, base_sleep = 2)
+      retries = 0
+
+      begin
+        yield
+      rescue StandardError => err
+        if retries < retries_limit
+          retries += 1
+          sleep_time = base_sleep**retries
+          Kernel.sleep(sleep_time)
+          retry
+        end
+
+        raise err
+      end
+    end
+  end
+end

data/lib/upperkut/version.rb

@@ -0,0 +1,3 @@
+module Upperkut
+  VERSION = '1.0.4'.freeze
+end

data/lib/upperkut/worker.rb

@@ -0,0 +1,42 @@
+require 'forwardable'
+require 'upperkut/strategies/buffered_queue'
+require 'upperkut/middleware'
+require 'upperkut'
+
+module Upperkut
+  module Worker
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      extend Forwardable
+
+      def_delegators :setup, :strategy, :server_middlewares, :client_middlewares
+      def_delegators :strategy, :metrics, :clear
+
+      def push_items(items)
+        client_middlewares.invoke(self, items) do
+          strategy.push_items(items)
+        end
+      end
+
+      def fetch_items
+        strategy.fetch_items
+      end
+
+      def setup_upperkut
+        yield(setup) if block_given?
+      end
+
+      def setup
+        @config ||=
+          begin
+            config = Upperkut::Configuration.default.clone
+            config.strategy ||= Upperkut::Strategies::BufferedQueue.new(self)
+            config
+          end
+      end
+    end
+  end
+end