HornsAndHooves-sidekiq-limit_fetch 4.5.0

data/lib/sidekiq/limit_fetch.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require 'forwardable'
+require 'sidekiq'
+require 'sidekiq/manager'
+require 'sidekiq/api'
+
+module Sidekiq
+  module 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'
+
+    TIMEOUT = Sidekiq::BasicFetch::TIMEOUT
+
+    extend self
+
+    RedisBaseConnectionError = RedisClient::ConnectionError
+    RedisCommandError = RedisClient::CommandError
+
+    def new(_)
+      self
+    end
+
+    def retrieve_work
+      queue, job = redis_brpop(Queues.acquire)
+      Queues.release_except(queue)
+      UnitOfWork.new(queue, job, capsule) if job
+    end
+
+    def capsule
+      Sidekiq.default_configuration.default_capsule
+    end
+
+    def bulk_requeue(*args)
+      Sidekiq::BasicFetch.new(capsule).bulk_requeue(*args)
+    end
+
+    def redis_retryable
+      yield
+    rescue RedisBaseConnectionError
+      sleep TIMEOUT
+      retry
+    rescue RedisCommandError => e
+      # If Redis was restarted and is still loading its snapshot,
+      # then we should treat this as a temporary connection error too.
+      raise unless e.message =~ /^LOADING/
+
+      sleep TIMEOUT
+      retry
+    end
+
+    private
+
+    # rubocop:disable Metrics/MethodLength
+    def redis_brpop(queues)
+      if queues.empty?
+        sleep TIMEOUT  # there are no queues to handle, so lets sleep
+        []             # and return nothing
+      else
+        redis_retryable do
+          Sidekiq.redis do |it|
+            it.blocking_call(false, 'brpop', *queues, TIMEOUT)
+          end
+        end
+      end
+    end
+    # rubocop:enable Metrics/MethodLength
+  end
+end

data/lib/sidekiq-limit_fetch.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require_relative 'sidekiq/limit_fetch'

data/sidekiq-limit_fetch.gemspec

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require 'English'
+Gem::Specification.new do |gem|
+  gem.name          = 'HornsAndHooves-sidekiq-limit_fetch'
+  gem.version       = '4.5.0'
+  gem.license       = 'MIT'
+  gem.authors       = ['HornsAndHooves', 'Peter Maneykowski']
+  gem.email         = ['maneyko@integracredit.com']
+  gem.summary       = 'Sidekiq strategy to support queue limits'
+  gem.homepage      = 'https://github.com/HornsAndHooves/sidekiq-limit_fetch'
+  gem.description   = 'Sidekiq strategy to restrict number of workers which are able to run specified ' \
+                      'queues simultaneously.'
+
+  gem.metadata['homepage_uri'] = gem.homepage
+  gem.metadata['source_code_uri'] = 'https://github.com/HornsAndHooves/sidekiq-limit_fetch'
+  gem.metadata['changelog_uri'] = 'https://github.com/HornsAndHooves/sidekiq-limit_fetch/blob/master/CHANGELOG.md'
+
+  gem.files         = `git ls-files`.split($INPUT_RECORD_SEPARATOR)
+  gem.require_paths = %w[lib]
+
+  gem.required_ruby_version = '>= 2.7.0'
+
+  gem.add_dependency 'sidekiq', '>= 8'
+  gem.add_development_dependency 'rake'
+  gem.add_development_dependency 'redis-namespace', '~> 1.5', '>= 1.5.2'
+  gem.add_development_dependency 'rspec'
+  gem.add_development_dependency 'rubocop'
+  gem.add_development_dependency 'simplecov'
+end

data/spec/sidekiq/extensions/manager_spec.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+RSpec.describe Sidekiq::Manager do
+  let(:capsule_or_options) do
+    Sidekiq.default_configuration.default_capsule
+  end
+
+  it 'can be instantiated' do
+    expect(described_class).to be < Sidekiq::Manager::InitLimitFetch
+    manager = described_class.new(capsule_or_options)
+    expect(manager).to respond_to(:start)
+  end
+end

data/spec/sidekiq/extensions/queue_spec.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+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,114 @@
+# frozen_string_literal: true
+
+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' }
+  let(:timeout) { 0.5 }
+
+  after { monitor.kill }
+
+  context 'old locks' do
+    before { monitor }
+
+    it 'should remove invalidated old locks' do
+      2.times { queue.acquire }
+      sleep ttl * 2
+      expect(queue.probed).to eq 2
+
+      allow(described_class).to receive(:update_heartbeat)
+      sleep ttl * 2
+      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 ttl * 2
+      expect(queue.probed).to eq 0
+    end
+  end
+
+  context 'dynamic queue' do
+    let(:limits) do
+      {
+        'queue1' => 3,
+        'queue2' => 3
+      }
+    end
+    let(:queues) { %w[queue1 queue2] }
+    let(:queue) { Sidekiq::LimitFetch::Queues }
+
+    let(:config) { Sidekiq::Config.new(options) }
+    let(:capsule) do
+      config.capsule('default') do |cap|
+        cap.concurrency = 1
+        cap.queues = config[:queues]
+      end
+    end
+
+    let(:capsule_or_options) do
+      capsule
+    end
+
+    context 'without excluded queue' do
+      let(:options) do
+        {
+          limits: limits,
+          queues: queues,
+          dynamic: true
+        }
+      end
+
+      it 'should add dynamic queue' do
+        queue.start(capsule_or_options)
+        monitor
+
+        expect(queue.instance_variable_get(:@queues)).not_to include('queue3')
+
+        Sidekiq.redis do |it|
+          it.sadd 'queues', 'queue3'
+        end
+
+        sleep ttl * 2
+        expect(queue.instance_variable_get(:@queues)).to include('queue3')
+
+        Sidekiq.redis do |it|
+          it.srem 'queues', 'queue3'
+        end
+      end
+    end
+
+    context 'with excluded queue' do
+      let(:options) do
+        {
+          limits: limits,
+          queues: queues,
+          dynamic: { exclude: ['queue4'] }
+        }
+      end
+
+      it 'should exclude excluded dynamic queue' do
+        queue.start(capsule_or_options)
+        monitor
+
+        expect(queue.instance_variable_get(:@queues)).not_to include('queue4')
+
+        Sidekiq.redis do |it|
+          it.sadd 'queues', 'queue4'
+        end
+
+        sleep ttl * 2
+        expect(queue.instance_variable_get(:@queues)).not_to include('queue4')
+
+        Sidekiq.redis do |it|
+          it.srem 'queues', 'queue4'
+        end
+      end
+    end
+  end
+end

data/spec/sidekiq/limit_fetch/queues_spec.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+RSpec.describe Sidekiq::LimitFetch::Queues do
+  let(:queues)         { %w[queue1 queue2] }
+  let(:limits)         { { 'queue1' => 3 } }
+  let(:strict)         { true }
+  let(:blocking)       { nil }
+  let(:process_limits) { { 'queue2' => 3 } }
+
+  let(:options) do
+    { queues: queues,
+      limits: limits,
+      strict: strict,
+      blocking: blocking,
+      process_limits: process_limits }
+  end
+
+  let(:config) { Sidekiq::Config.new(options) }
+  let(:capsule) do
+    config.capsule('default') do |cap|
+      cap.concurrency = 1
+      cap.queues = config[:queues]
+    end
+  end
+
+  let(:capsule_or_options) do
+    capsule
+  end
+
+  before do
+    subject.start(capsule_or_options)
+  end
+
+  def in_thread(&block)
+    thr = Thread.new(&block)
+    thr.join
+  end
+
+  it 'should acquire queues' do
+    in_thread { subject.acquire }
+    expect(Sidekiq::Queue['queue1'].probed).to eq 1
+    expect(Sidekiq::Queue['queue2'].probed).to eq 1
+  end
+
+  it 'should acquire dynamically blocking queues' do
+    in_thread { subject.acquire }
+    expect(Sidekiq::Queue['queue1'].probed).to eq 1
+    expect(Sidekiq::Queue['queue2'].probed).to eq 1
+
+    Sidekiq::Queue['queue1'].block
+
+    in_thread { subject.acquire }
+    expect(Sidekiq::Queue['queue1'].probed).to eq 2
+    expect(Sidekiq::Queue['queue2'].probed).to eq 1
+  end
+
+  it 'should block except given queues' do
+    Sidekiq::Queue['queue1'].block_except 'queue2'
+    in_thread { subject.acquire }
+    expect(Sidekiq::Queue['queue1'].probed).to eq 1
+    expect(Sidekiq::Queue['queue2'].probed).to eq 1
+
+    Sidekiq::Queue['queue1'].block_except 'queue404'
+    in_thread { subject.acquire }
+    expect(Sidekiq::Queue['queue1'].probed).to eq 2
+    expect(Sidekiq::Queue['queue2'].probed).to eq 1
+  end
+
+  it 'should release queues' do
+    in_thread do
+      subject.acquire
+      subject.release_except nil
+    end
+    expect(Sidekiq::Queue['queue1'].probed).to eq 0
+    expect(Sidekiq::Queue['queue2'].probed).to eq 0
+  end
+
+  it 'should release queues except selected' do
+    in_thread do
+      subject.acquire
+      subject.release_except 'queue:queue1'
+    end
+    expect(Sidekiq::Queue['queue1'].probed).to eq 1
+    expect(Sidekiq::Queue['queue2'].probed).to eq 0
+  end
+
+  it 'should release when no queues was acquired' do
+    queues.each { |name| Sidekiq::Queue[name].pause }
+    in_thread do
+      subject.acquire
+      expect { subject.release_except nil }.not_to raise_exception
+    end
+  end
+
+  context 'blocking' do
+    let(:blocking) { %w[queue1] }
+
+    it 'should acquire blocking queues' do
+      3.times { in_thread { subject.acquire } }
+      expect(Sidekiq::Queue['queue1'].probed).to eq 3
+      expect(Sidekiq::Queue['queue2'].probed).to eq 1
+    end
+  end
+
+  it 'should set limits' do
+    subject
+    expect(Sidekiq::Queue['queue1'].limit).to eq 3
+    expect(Sidekiq::Queue['queue2'].limit).not_to be
+  end
+
+  it 'should set process_limits' do
+    subject
+    expect(Sidekiq::Queue['queue2'].process_limit).to eq 3
+  end
+
+  context 'without strict flag' do
+    let(:strict) { false }
+
+    it 'should retrieve weighted queues' do
+      expect(subject.ordered_queues).to match_array(%w[queue1 queue2])
+    end
+  end
+
+  it 'with strict flag should retrieve strictly ordered queues' do
+    expect(subject.ordered_queues).to eq %w[queue1 queue2]
+  end
+end

data/spec/sidekiq/limit_fetch/semaphore_spec.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+RSpec.describe 'semaphore' do
+  let(:name) { 'default' }
+  subject { Sidekiq::LimitFetch::Global::Semaphore.new name }
+
+  it 'should have no limit by default' do
+    expect(subject.limit).not_to be
+  end
+
+  it 'should set limit' do
+    subject.limit = 4
+    expect(subject.limit).to eq 4
+  end
+
+  it 'should acquire and count active tasks' do
+    3.times { subject.acquire }
+    expect(subject.probed).to eq 3
+  end
+
+  it 'should acquire tasks with regard to limit' do
+    subject.limit = 4
+    6.times { subject.acquire }
+    expect(subject.probed).to eq 4
+  end
+
+  it 'should acquire tasks with regard to process limit' do
+    subject.process_limit = 4
+    6.times { subject.acquire }
+    expect(subject.probed).to eq 4
+  end
+
+  it 'should release active tasks' do
+    6.times { subject.acquire }
+    3.times { subject.release }
+    expect(subject.probed).to eq 3
+  end
+
+  it 'should pause tasks' do
+    3.times { subject.acquire }
+    subject.pause
+    2.times { subject.acquire }
+    expect(subject.probed).to eq 3
+    2.times { subject.release }
+    expect(subject.probed).to eq 1
+  end
+
+  it 'should unpause tasks' do
+    subject.pause
+    3.times { subject.acquire }
+    subject.unpause
+    2.times { subject.acquire }
+    expect(subject.probed).to eq 2
+  end
+
+  it 'should pause tasks for a limited time' do
+    3.times { subject.acquire }
+    subject.pause_for_ms 50
+    2.times { subject.acquire }
+    expect(subject.probed).to eq 3
+    sleep(100.0 / 1000)
+    2.times { subject.acquire }
+    expect(subject.probed).to eq 5
+  end
+end

data/spec/sidekiq/limit_fetch_spec.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+Thread.abort_on_exception = true
+
+RSpec.describe Sidekiq::LimitFetch do
+  let(:options) { { queues: queues, limits: limits } }
+  let(:queues) { %w[queue1 queue1 queue2 queue2] }
+  let(:limits) { { 'queue1' => 1, 'queue2' => 2 } }
+  let(:config) { Sidekiq::Config.new(options) }
+  let(:capsule) do
+    config.capsule('default') do |cap|
+      cap.concurrency = 1
+      cap.queues = config[:queues]
+    end
+  end
+  let(:capsule_or_config) { capsule }
+
+  before do
+    subject::Queues.start(capsule_or_config)
+
+    Sidekiq.redis do |it|
+      it.del 'queue:queue1'
+      it.lpush 'queue:queue1', 'task1'
+      it.lpush 'queue:queue1', 'task2'
+      it.expire 'queue:queue1', 30
+    end
+  end
+
+  it 'should acquire lock on queue for execution' do
+    work = subject.retrieve_work
+    expect(work.queue_name).to eq 'queue1'
+    expect(work.job).to eq 'task1'
+
+    expect(Sidekiq::Queue['queue1'].busy).to eq 1
+    expect(Sidekiq::Queue['queue2'].busy).to eq 0
+
+    expect(subject.retrieve_work).not_to be
+    work.requeue
+
+    expect(Sidekiq::Queue['queue1'].busy).to eq 0
+    expect(Sidekiq::Queue['queue2'].busy).to eq 0
+
+    work = subject.retrieve_work
+    expect(work.job).to eq 'task1'
+
+    expect(Sidekiq::Queue['queue1'].busy).to eq 1
+    expect(Sidekiq::Queue['queue2'].busy).to eq 0
+
+    expect(subject.retrieve_work).not_to be
+    work.acknowledge
+
+    expect(Sidekiq::Queue['queue1'].busy).to eq 0
+    expect(Sidekiq::Queue['queue2'].busy).to eq 0
+
+    work = subject.retrieve_work
+    expect(work.job).to eq 'task2'
+  end
+end

data/spec/spec_helper.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require 'simplecov'
+SimpleCov.start
+
+require 'sidekiq/limit_fetch'
+
+Sidekiq.configure_embed do |config|
+  config.logger = nil
+end
+
+RSpec.configure do |config|
+  config.order = :random
+  config.disable_monkey_patching!
+  config.raise_errors_for_deprecations!
+  config.before do
+    Sidekiq::Queue.reset_instances!
+    Sidekiq.redis do |it|
+      clean_redis = lambda do |queue|
+        it.pipelined do |pipeline|
+          pipeline.del "limit_fetch:limit:#{queue}"
+          pipeline.del "limit_fetch:process_limit:#{queue}"
+          pipeline.del "limit_fetch:busy:#{queue}"
+          pipeline.del "limit_fetch:probed:#{queue}"
+          pipeline.del "limit_fetch:pause:#{queue}"
+          pipeline.del "limit_fetch:block:#{queue}"
+        end
+      end
+
+      clean_redis.call(name) if defined?(name)
+      queues.each(&clean_redis) if defined?(queues) && queues.is_a?(Array)
+    end
+  end
+end