rcgt-sidekiq-limit_fetch 3.4.1

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

@@ -0,0 +1,125 @@
+module Sidekiq::LimitFetch::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|
+        begin
+          it.evalsha send("redis_#{script_name}_sha"), argv: args
+        rescue Redis::CommandError => error
+          raise unless error.message.include? 'NOSCRIPT'
+          it.eval send("redis_#{script_name}_script"), argv: args
+        end
+      end
+    end
+
+    def redis_acquire_sha
+      @acquire_sha ||= Digest::SHA1.hexdigest redis_acquire_script
+    end
+
+    def redis_release_sha
+      @release_sha ||= 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

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

@@ -0,0 +1,175 @@
+module Sidekiq::LimitFetch::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 if value
+    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 if value
+    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 > 0
+    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", true }
+    end
+
+    def pause_for_ms ms
+      redis {|it| it.psetex "#{PREFIX}:pause:#@name", ms, true }
+    end
+
+    def unpause
+      redis {|it| it.del "#{PREFIX}:pause:#@name" }
+    end
+
+    def paused?
+      redis {|it| it.get "#{PREFIX}:pause:#@name" }
+    end
+
+    def block
+      redis {|it| it.set "#{PREFIX}:block:#@name", true }
+    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" }
+    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 > 0
+    end
+
+    def explain
+      <<-END.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?}
+      END
+    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

data/lib/sidekiq/limit_fetch/instances.rb

@@ -0,0 +1,19 @@
+module Sidekiq::LimitFetch::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
+end

data/lib/sidekiq/limit_fetch/queues.rb

@@ -0,0 +1,124 @@
+module Sidekiq::LimitFetch::Queues
+  extend self
+
+  THREAD_KEY = :acquired_queues
+
+  def start(options)
+    @queues         = options[:queues]
+    @dynamic        = options[:dynamic]
+
+    @limits         = options[:limits] || {}
+    @process_limits = options[:process_limits] || {}
+    @blocks         = options[:blocking] || []
+
+    options[:strict] ? strict_order! : weighted_order!
+
+    apply_process_limit_to_queues
+    apply_limit_to_queues
+    apply_blocks_to_queues
+  end
+
+  def acquire
+    selector.acquire(ordered_queues, namespace)
+      .tap {|it| save it }
+      .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 add(queues)
+    queues.each do |queue|
+      unless @queues.include? queue
+        apply_process_limit_to_queue(queue)
+        apply_limit_to_queue(queue)
+
+        @queues.push queue
+      end
+    end
+  end
+
+  def strict_order!
+    @queues.uniq!
+    def ordered_queues; @queues end
+  end
+
+  def weighted_order!
+    def ordered_queues; @queues.shuffle.uniq end
+  end
+
+  def namespace
+    @namespace ||= Sidekiq.redis do |it|
+      if it.respond_to?(:namespace) and 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]
+    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]
+
+    unless queue.limit_changed?
+      queue.limit = @limits[queue_name.to_s] || @limits[queue_name.to_sym]
+    end
+  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 selector
+    Sidekiq::LimitFetch::Global::Selector
+  end
+
+  def save(queues)
+    Thread.current[THREAD_KEY] = queues
+  end
+
+  def restore
+    Thread.current[THREAD_KEY] || []
+  ensure
+    Thread.current[THREAD_KEY] = nil
+  end
+end

data/lib/sidekiq/limit_fetch/unit_of_work.rb

@@ -0,0 +1,24 @@
+module Sidekiq
+  class LimitFetch::UnitOfWork < BasicFetch::UnitOfWork
+    def initialize(queue, job)
+      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

data/lib/sidekiq/limit_fetch.rb

@@ -0,0 +1,58 @@
+require 'forwardable'
+require 'sidekiq'
+require 'sidekiq/manager'
+require 'sidekiq/api'
+
+module Sidekiq::LimitFetch
+  autoload :UnitOfWork, 'sidekiq/limit_fetch/unit_of_work'
+
+  require_relative 'limit_fetch/instances'
+  require_relative 'limit_fetch/queues'
+  require_relative 'limit_fetch/global/semaphore'
+  require_relative 'limit_fetch/global/selector'
+  require_relative 'limit_fetch/global/monitor'
+  require_relative 'extensions/queue'
+  require_relative 'extensions/manager'
+
+  extend self
+
+  def new(_)
+    self
+  end
+
+  def retrieve_work
+    queue, job = redis_brpop(Queues.acquire)
+    Queues.release_except(queue)
+    UnitOfWork.new(queue, job) if job
+  end
+
+  # Backwards compatibility for sidekiq v6.1.0
+  # @see https://github.com/mperham/sidekiq/pull/4602
+  def bulk_requeue(*args)
+    if Sidekiq::BasicFetch.respond_to?(:bulk_requeue) # < 6.1.0
+      Sidekiq::BasicFetch.bulk_requeue(*args)
+    else # 6.1.0+
+      Sidekiq::BasicFetch.new(Sidekiq.options).bulk_requeue(*args)
+    end
+  end
+
+  def redis_retryable
+    yield
+  rescue Redis::BaseConnectionError
+    sleep 1
+    retry
+  end
+
+  private
+
+  TIMEOUT = Sidekiq::BasicFetch::TIMEOUT
+
+  def redis_brpop(queues)
+    if queues.empty?
+      sleep TIMEOUT  # there are no queues to handle, so lets sleep
+      []             # and return nothing
+    else
+      redis_retryable { Sidekiq.redis { |it| it.brpop *queues, TIMEOUT } }
+    end
+  end
+end

data/lib/sidekiq-limit_fetch.rb

@@ -0,0 +1 @@
+require_relative 'sidekiq/limit_fetch'

data/rcgt-sidekiq-limit_fetch.gemspec

@@ -0,0 +1,24 @@
+Gem::Specification.new do |gem|
+  gem.name          = 'rcgt-sidekiq-limit_fetch'
+  gem.version       = '3.4.1'
+  gem.license       = 'MIT'
+  gem.authors       = 'RCGT Consulting Inc.'
+  gem.email         = 'kumanan.yogaratnam@rcgtconsulting.com'
+  gem.summary       = 'Sidekiq strategy to support queue limits'
+  gem.homepage      = 'https://github.com/RCGTConsulting/sidekiq-limit_fetch'
+  gem.description   = <<-DESCRIPTION
+    A fork of https://github.com/brainopia/sidekiq-limit_fetch
+    
+    Sidekiq strategy to restrict number of workers
+    which are able to run specified queues simultaneously.
+  DESCRIPTION
+
+  gem.files         = `git ls-files`.split($/)
+  gem.test_files    = gem.files.grep %r{^spec/}
+  gem.require_paths = %w(lib)
+
+  gem.add_dependency 'sidekiq', '>= 4'
+  gem.add_development_dependency 'redis-namespace', '~> 1.5', '>= 1.5.2'
+  gem.add_development_dependency 'rspec'
+  gem.add_development_dependency 'rake'
+end

data/spec/sidekiq/extensions/queue_spec.rb

@@ -0,0 +1,94 @@
+RSpec.describe Sidekiq::Queue do
+  context 'singleton' do
+    shared_examples :constructor do
+      it 'with default name' do
+        new_object = -> { described_class.send constructor }
+        expect(new_object.call).to eq new_object.call
+      end
+
+      it 'with given name' do
+        new_object = ->(name) { described_class.send constructor, name }
+        expect(new_object.call('name')).to eq new_object.call('name')
+      end
+    end
+
+    context '.new' do
+      let(:constructor) { :new }
+      it_behaves_like :constructor
+    end
+
+    context '.[]' do
+      let(:constructor) { :[] }
+      it_behaves_like :constructor
+    end
+
+    context '#lock' do
+      let(:name) { 'example' }
+      let(:queue) { Sidekiq::Queue[name] }
+
+      it 'should be available' do
+        expect(queue.acquire).to be
+      end
+
+      it 'should be pausable' do
+        queue.pause
+        expect(queue.acquire).not_to be
+      end
+
+      it 'should be continuable' do
+        queue.pause
+        queue.unpause
+        expect(queue.acquire).to be
+      end
+
+      it 'should be limitable' do
+        queue.limit = 1
+        expect(queue.acquire).to be
+        expect(queue.acquire).not_to be
+      end
+
+      it 'should be resizable' do
+        queue.limit = 0
+        expect(queue.acquire).not_to be
+        queue.limit = nil
+        expect(queue.acquire).to be
+      end
+
+      it 'should be countable' do
+        queue.limit = 3
+        5.times { queue.acquire }
+        expect(queue.probed).to eq 3
+      end
+
+      it 'should be releasable' do
+        queue.acquire
+        expect(queue.probed).to eq 1
+        queue.release
+        expect(queue.probed).to eq 0
+      end
+
+      it 'should tell if paused' do
+        expect(queue).not_to be_paused
+        queue.pause
+        expect(queue).to be_paused
+        queue.unpause
+        expect(queue).not_to be_paused
+      end
+
+      it 'should tell if blocking' do
+        expect(queue).not_to be_blocking
+        queue.block
+        expect(queue).to be_blocking
+        queue.unblock
+        expect(queue).not_to be_blocking
+      end
+
+      it 'should be marked as changed' do
+        queue = Sidekiq::Queue["uniq_#{name}"]
+        expect(queue).not_to be_limit_changed
+        queue.limit = 3
+        expect(queue).to be_limit_changed
+      end
+    end
+  end
+end

data/spec/sidekiq/limit_fetch/global/monitor_spec.rb

@@ -0,0 +1,33 @@
+RSpec.describe Sidekiq::LimitFetch::Global::Monitor do
+  let(:monitor) { described_class.start! ttl, timeout }
+  let(:ttl) { 1 }
+  let(:queue) { Sidekiq::Queue[name] }
+  let(:name) { 'default' }
+
+  before { monitor }
+  after { monitor.kill }
+
+  context 'old locks' do
+    let(:timeout) { 0.5 }
+
+    it 'should remove invalidated old locks' do
+      2.times { queue.acquire }
+      sleep 2*ttl
+      expect(queue.probed).to eq 2
+
+      allow(described_class).to receive(:update_heartbeat)
+      sleep 2*ttl
+      expect(queue.probed).to eq 0
+    end
+
+    it 'should remove invalid locks' do
+      2.times { queue.acquire }
+      allow(described_class).to receive(:update_heartbeat)
+      Sidekiq.redis do |it|
+        it.del Sidekiq::LimitFetch::Global::Monitor::PROCESS_SET
+      end
+      sleep 2*ttl
+      expect(queue.probed).to eq 0
+    end
+  end
+end