HornsAndHooves-sidekiq-limit_fetch 4.5.0

data/lib/sidekiq/limit_fetch/global/monitor.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module LimitFetch
+    module Global
+      module Monitor
+        extend self
+
+        HEARTBEAT_PREFIX = 'limit:heartbeat:'
+        PROCESS_SET = 'limit:processes'
+        HEARTBEAT_TTL = 20
+        REFRESH_TIMEOUT = 5
+
+        def start!(ttl = HEARTBEAT_TTL, timeout = REFRESH_TIMEOUT)
+          Thread.new do
+            loop do
+              Sidekiq::LimitFetch.redis_retryable do
+                handle_dynamic_queues
+                update_heartbeat ttl
+                invalidate_old_processes
+              end
+
+              sleep timeout
+            end
+          end
+        end
+
+        def all_processes
+          Sidekiq.redis { |it| it.smembers PROCESS_SET }
+        end
+
+        def old_processes
+          all_processes.reject do |process|
+            Sidekiq.redis { |it| it.get heartbeat_key process } == '1'
+          end
+        end
+
+        def remove_old_processes!
+          Sidekiq.redis do |it|
+            old_processes.each { |process| it.srem PROCESS_SET, [process] }
+          end
+        end
+
+        def handle_dynamic_queues
+          queues = Sidekiq::LimitFetch::Queues
+          return unless queues.dynamic?
+
+          available_queues = Sidekiq::Queue.all.map(&:name).reject do |it|
+            queues.dynamic_exclude.include? it
+          end
+          queues.handle available_queues
+        end
+
+        private
+
+        def update_heartbeat(ttl)
+          Sidekiq.redis do |it|
+            it.multi do |pipeline|
+              pipeline.set heartbeat_key, '1'
+              pipeline.sadd PROCESS_SET, [Selector.uuid]
+              pipeline.expire heartbeat_key, ttl
+            end
+          end
+        end
+
+        def invalidate_old_processes
+          Sidekiq.redis do |_it|
+            remove_old_processes!
+            processes = all_processes
+
+            Sidekiq::Queue.instances.each do |queue|
+              queue.remove_locks_except! processes
+            end
+          end
+        end
+
+        def heartbeat_key(process = Selector.uuid)
+          HEARTBEAT_PREFIX + process
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/limit_fetch/global/selector.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module LimitFetch
+    module Global
+      module Selector
+        extend self
+
+        MUTEX_FOR_UUID = Mutex.new
+
+        def acquire(queues, namespace)
+          redis_eval :acquire, [namespace, uuid, queues]
+        end
+
+        def release(queues, namespace)
+          redis_eval :release, [namespace, uuid, queues]
+        end
+
+        def uuid
+          # - if we'll remove "@uuid ||=" from inside of mutex
+          # then @uuid can be overwritten
+          # - if we'll remove "@uuid ||=" from outside of mutex
+          # then each read will lead to mutex
+          @uuid ||= MUTEX_FOR_UUID.synchronize { @uuid || SecureRandom.uuid }
+        end
+
+        private
+
+        def redis_eval(script_name, args)
+          Sidekiq.redis do |it|
+            it.evalsha send("redis_#{script_name}_sha"), [], args
+          rescue Sidekiq::LimitFetch::RedisCommandError => e
+            raise unless e.message.include? 'NOSCRIPT'
+
+            it.eval send("redis_#{script_name}_script"), 0, *args
+          end
+        end
+
+        def redis_acquire_sha
+          @redis_acquire_sha ||= OpenSSL::Digest::SHA1.hexdigest redis_acquire_script
+        end
+
+        def redis_release_sha
+          @redis_release_sha ||= OpenSSL::Digest::SHA1.hexdigest redis_release_script
+        end
+
+        def redis_acquire_script
+          <<-LUA
+        local namespace   = table.remove(ARGV, 1)..'limit_fetch:'
+        local worker_name = table.remove(ARGV, 1)
+        local queues      = ARGV
+        local available   = {}
+        local unblocked   = {}
+        local locks
+        local process_locks
+        local blocking_mode
+
+        for _, queue in ipairs(queues) do
+          if not blocking_mode or unblocked[queue] then
+            local probed_key        = namespace..'probed:'..queue
+            local pause_key         = namespace..'pause:'..queue
+            local limit_key         = namespace..'limit:'..queue
+            local process_limit_key = namespace..'process_limit:'..queue
+            local block_key         = namespace..'block:'..queue
+
+            local paused, limit, process_limit, can_block =
+              unpack(redis.call('mget',
+                pause_key,
+                limit_key,
+                process_limit_key,
+                block_key
+              ))
+
+            if not paused then
+              limit = tonumber(limit)
+              process_limit = tonumber(process_limit)
+
+              if can_block or limit then
+                locks = redis.call('llen', probed_key)
+              end
+
+              if process_limit then
+                local all_locks = redis.call('lrange', probed_key, 0, -1)
+                process_locks = 0
+                for _, process in ipairs(all_locks) do
+                  if process == worker_name then
+                    process_locks = process_locks + 1
+                  end
+                end
+              end
+
+              if not blocking_mode then
+                blocking_mode = can_block and locks > 0
+              end
+
+              if blocking_mode and can_block ~= 'true' then
+                for unblocked_queue in string.gmatch(can_block, "[^,]+") do
+                  unblocked[unblocked_queue] = true
+                end
+              end
+
+              if (not limit or limit > locks) and
+                 (not process_limit or process_limit > process_locks) then
+                redis.call('rpush', probed_key, worker_name)
+                table.insert(available, queue)
+              end
+            end
+          end
+        end
+
+        return available
+          LUA
+        end
+
+        def redis_release_script
+          <<-LUA
+        local namespace   = table.remove(ARGV, 1)..'limit_fetch:'
+        local worker_name = table.remove(ARGV, 1)
+        local queues      = ARGV
+
+        for _, queue in ipairs(queues) do
+          local probed_key = namespace..'probed:'..queue
+          redis.call('lrem', probed_key, 1, worker_name)
+        end
+          LUA
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/limit_fetch/global/semaphore.rb

@@ -0,0 +1,190 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module LimitFetch
+    module Global
+      class Semaphore
+        PREFIX = 'limit_fetch'
+
+        attr_reader :local_busy
+
+        def initialize(name)
+          @name = name
+          @lock = Mutex.new
+          @local_busy = 0
+        end
+
+        def limit
+          value = redis { |it| it.get "#{PREFIX}:limit:#{@name}" }
+          value&.to_i
+        end
+
+        def limit=(value)
+          @limit_changed = true
+
+          if value
+            redis { |it| it.set "#{PREFIX}:limit:#{@name}", value }
+          else
+            redis { |it| it.del "#{PREFIX}:limit:#{@name}" }
+          end
+        end
+
+        def limit_changed?
+          @limit_changed
+        end
+
+        def process_limit
+          value = redis { |it| it.get "#{PREFIX}:process_limit:#{@name}" }
+          value&.to_i
+        end
+
+        def process_limit=(value)
+          if value
+            redis { |it| it.set "#{PREFIX}:process_limit:#{@name}", value }
+          else
+            redis { |it| it.del "#{PREFIX}:process_limit:#{@name}" }
+          end
+        end
+
+        def acquire
+          Selector.acquire([@name], namespace).size.positive?
+        end
+
+        def release
+          redis { |it| it.lrem "#{PREFIX}:probed:#{@name}", 1, Selector.uuid }
+        end
+
+        def busy
+          redis { |it| it.llen "#{PREFIX}:busy:#{@name}" }
+        end
+
+        def busy_processes
+          redis { |it| it.lrange "#{PREFIX}:busy:#{@name}", 0, -1 }
+        end
+
+        def increase_busy
+          increase_local_busy
+          redis { |it| it.rpush "#{PREFIX}:busy:#{@name}", Selector.uuid }
+        end
+
+        def decrease_busy
+          decrease_local_busy
+          redis { |it| it.lrem "#{PREFIX}:busy:#{@name}", 1, Selector.uuid }
+        end
+
+        def probed
+          redis { |it| it.llen "#{PREFIX}:probed:#{@name}" }
+        end
+
+        def probed_processes
+          redis { |it| it.lrange "#{PREFIX}:probed:#{@name}", 0, -1 }
+        end
+
+        def pause
+          redis { |it| it.set "#{PREFIX}:pause:#{@name}", '1' }
+        end
+
+        def pause_for_ms(milliseconds)
+          redis { |it| it.psetex "#{PREFIX}:pause:#{@name}", milliseconds, 1 }
+        end
+
+        def unpause
+          redis { |it| it.del "#{PREFIX}:pause:#{@name}" }
+        end
+
+        def paused?
+          redis { |it| it.get "#{PREFIX}:pause:#{@name}" } == '1'
+        end
+
+        def block
+          redis { |it| it.set "#{PREFIX}:block:#{@name}", '1' }
+        end
+
+        def block_except(*queues)
+          raise ArgumentError if queues.empty?
+
+          redis { |it| it.set "#{PREFIX}:block:#{@name}", queues.join(',') }
+        end
+
+        def unblock
+          redis { |it| it.del "#{PREFIX}:block:#{@name}" }
+        end
+
+        def blocking?
+          redis { |it| it.get "#{PREFIX}:block:#{@name}" } == '1'
+        end
+
+        def clear_limits
+          redis do |it|
+            %w[block busy limit pause probed process_limit].each do |key|
+              it.del "#{PREFIX}:#{key}:#{@name}"
+            end
+          end
+        end
+
+        def increase_local_busy
+          @lock.synchronize { @local_busy += 1 }
+        end
+
+        def decrease_local_busy
+          @lock.synchronize { @local_busy -= 1 }
+        end
+
+        def local_busy?
+          @local_busy.positive?
+        end
+
+        def explain
+          <<-INFO.gsub(/^ {8}/, '')
+        Current sidekiq process: #{Selector.uuid}
+
+          All processes:
+        #{Monitor.all_processes.join "\n"}
+
+          Stale processes:
+        #{Monitor.old_processes.join "\n"}
+
+          Locked queue processes:
+        #{probed_processes.sort.join "\n"}
+
+          Busy queue processes:
+        #{busy_processes.sort.join "\n"}
+
+          Limit:
+        #{limit.inspect}
+
+          Process limit:
+        #{process_limit.inspect}
+
+          Blocking:
+        #{blocking?}
+          INFO
+        end
+
+        def remove_locks_except!(processes)
+          locked_processes = probed_processes.uniq
+          (locked_processes - processes).each do |dead_process|
+            remove_lock! dead_process
+          end
+        end
+
+        def remove_lock!(process)
+          redis do |it|
+            it.lrem "#{PREFIX}:probed:#{@name}", 0, process
+            it.lrem "#{PREFIX}:busy:#{@name}", 0, process
+          end
+        end
+
+        private
+
+        def redis(&block)
+          Sidekiq.redis(&block)
+        end
+
+        def namespace
+          Sidekiq::LimitFetch::Queues.namespace
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/limit_fetch/instances.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module LimitFetch
+    module Instances
+      def self.extended(klass)
+        klass.instance_variable_set :@instances, {}
+      end
+
+      def new(*args)
+        @instances[args] ||= super
+      end
+
+      alias [] new
+
+      def instances
+        @instances.values
+      end
+
+      def reset_instances!
+        @instances = {}
+      end
+
+      def delete_instance(name)
+        @instances.delete [name]
+      end
+    end
+  end
+end

data/lib/sidekiq/limit_fetch/queues.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module LimitFetch
+    module Queues
+      extend self
+
+      THREAD_KEY = :acquired_queues
+
+      # rubocop:disable Metrics/AbcSize
+      # rubocop:disable Metrics/CyclomaticComplexity
+      # rubocop:disable Metrics/MethodLength
+      # rubocop:disable Metrics/PerceivedComplexity
+      def start(capsule)
+        config  = capsule.config
+        @queues = capsule.queues.map { |queue| queue.is_a?(Array) ? queue.first : queue }.uniq
+
+        @startup_queues = @queues.dup
+
+        if config[:dynamic].is_a? Hash
+          @dynamic         = true
+          @dynamic_exclude = config[:dynamic][:exclude] || []
+        else
+          @dynamic = config[:dynamic]
+          @dynamic_exclude = []
+        end
+
+        @limits         = config[:limits] || {}
+        @process_limits = config[:process_limits] || {}
+        @blocks         = config[:blocking] || []
+
+        config[:strict] ? strict_order! : weighted_order!
+
+        apply_process_limit_to_queues
+        apply_limit_to_queues
+        apply_blocks_to_queues
+      end
+      # rubocop:enable Metrics/AbcSize
+      # rubocop:enable Metrics/CyclomaticComplexity
+      # rubocop:enable Metrics/MethodLength
+      # rubocop:enable Metrics/PerceivedComplexity
+
+      def acquire
+        queues = saved
+        queues ||= Sidekiq::LimitFetch.redis_retryable do
+          selector.acquire(ordered_queues, namespace)
+        end
+        save queues
+        queues.map { |it| "queue:#{it}" }
+      end
+
+      def release_except(full_name)
+        queues = restore
+        queues.delete full_name[/queue:(.*)/, 1] if full_name
+        Sidekiq::LimitFetch.redis_retryable do
+          selector.release queues, namespace
+        end
+      end
+
+      def dynamic?
+        @dynamic
+      end
+
+      def startup_queue?(queue)
+        @startup_queues.include?(queue)
+      end
+
+      def dynamic_exclude
+        @dynamic_exclude
+      end
+
+      def add(queues)
+        return unless queues
+
+        queues.each do |queue|
+          next if @queues.include? queue
+
+          if startup_queue?(queue)
+            apply_process_limit_to_queue(queue)
+            apply_limit_to_queue(queue)
+          end
+
+          @queues.push queue
+        end
+      end
+
+      def remove(queues)
+        return unless queues
+
+        queues.each do |queue|
+          next unless @queues.include? queue
+
+          clear_limits_for_queue(queue)
+          @queues.delete queue
+          Sidekiq::Queue.delete_instance(queue)
+        end
+      end
+
+      def handle(queues)
+        add(queues - @queues)
+        remove(@queues - queues)
+      end
+
+      # rubocop:disable Lint/NestedMethodDefinition
+      def strict_order!
+        @queues.uniq!
+        def ordered_queues
+          @queues
+        end
+      end
+
+      def weighted_order!
+        def ordered_queues
+          @queues.shuffle.uniq
+        end
+      end
+      # rubocop:enable Lint/NestedMethodDefinition
+
+      def namespace
+        @namespace ||= Sidekiq.redis do |it|
+          if it.respond_to?(:namespace) && it.namespace
+            "#{it.namespace}:"
+          else
+            ''
+          end
+        end
+      end
+
+      private
+
+      def apply_process_limit_to_queues
+        @queues.uniq.each do |queue_name|
+          apply_process_limit_to_queue(queue_name)
+        end
+      end
+
+      def apply_process_limit_to_queue(queue_name)
+        queue = Sidekiq::Queue[queue_name]
+        return unless queue.process_limit.nil? # honor existing deployed limits
+
+        queue.process_limit = @process_limits[queue_name.to_s] || @process_limits[queue_name.to_sym]
+      end
+
+      def apply_limit_to_queues
+        @queues.uniq.each do |queue_name|
+          apply_limit_to_queue(queue_name)
+        end
+      end
+
+      def apply_limit_to_queue(queue_name)
+        queue = Sidekiq::Queue[queue_name]
+
+        return if queue.limit_changed?
+        return unless queue.limit.nil? # honor existing deployed limits
+
+        queue.limit = @limits[queue_name.to_s] || @limits[queue_name.to_sym]
+      end
+
+      def apply_blocks_to_queues
+        @queues.uniq.each do |queue_name|
+          Sidekiq::Queue[queue_name].unblock
+        end
+
+        @blocks.to_a.each do |it|
+          if it.is_a? Array
+            it.each { |name| Sidekiq::Queue[name].block_except it }
+          else
+            Sidekiq::Queue[it].block
+          end
+        end
+      end
+
+      def clear_limits_for_queue(queue_name)
+        queue = Sidekiq::Queue[queue_name]
+        queue.clear_limits
+      end
+
+      def selector
+        Sidekiq::LimitFetch::Global::Selector
+      end
+
+      def saved
+        Thread.current[THREAD_KEY]
+      end
+
+      def save(queues)
+        Thread.current[THREAD_KEY] = queues
+      end
+
+      def restore
+        saved || []
+      ensure
+        save nil
+      end
+    end
+  end
+end

data/lib/sidekiq/limit_fetch/unit_of_work.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module LimitFetch
+    class UnitOfWork < BasicFetch::UnitOfWork
+      def initialize(...)
+        super
+        redis_retryable { Queue[queue_name].increase_busy }
+      end
+
+      def acknowledge
+        redis_retryable { Queue[queue_name].decrease_busy }
+        redis_retryable { Queue[queue_name].release }
+      end
+
+      def requeue
+        super
+        acknowledge
+      end
+
+      private
+
+      def redis_retryable(&block)
+        Sidekiq::LimitFetch.redis_retryable(&block)
+      end
+    end
+  end
+end