├── .gitignore ├── docker-compose.yml ├── Gemfile ├── Rakefile ├── LICENSE ├── lib └── sidekiq │ ├── priority_queue.rb │ └── priority_queue │ ├── scripts.rb │ ├── web │ └── views │ │ ├── _paging.erb │ │ ├── priority_queues.erb │ │ └── priority_queue.erb │ ├── testing.rb │ ├── combined_fetch.rb │ ├── web.rb │ ├── client.rb │ ├── api.rb │ ├── fetch.rb │ └── reliable_fetch.rb ├── .github └── workflows │ └── tests.yml ├── test ├── test_api.rb ├── test_web.rb ├── test_fetch.rb ├── test_client.rb ├── test_combined_fetch.rb ├── helper.rb └── test_reliable_fetch.rb ├── sidekiq-priority_queue.gemspec ├── Gemfile.lock ├── README.md └── bin └── sidekiqload /.gitignore: -------------------------------------------------------------------------------- 1 | .vagrant 2 | /.idea/ 3 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | --- 2 | version: '3.3' 3 | services: 4 | redis: 5 | image: redis:alpine 6 | ports: 7 | - 6379:6379 -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | source 'https://rubygems.org' 2 | gemspec 3 | 4 | 5 | group :test do 6 | gem 'pry-byebug' 7 | gem 'minitest' 8 | gem 'rake' 9 | gem 'simplecov' 10 | gem 'rack-test' 11 | end 12 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require 'bundler/gem_tasks' 2 | require 'rake/testtask' 3 | Rake::TestTask.new(:test) do |test| 4 | test.warning = true 5 | test.pattern = 'test/**/test_*.rb' 6 | end 7 | 8 | task :default => :test 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) ChartMogul Ltd 2 | 3 | Sidekiq-priority_queue is an Open Source project licensed under the terms of 4 | the LGPLv3 license. Please see 5 | for license text. 6 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require 'sidekiq' 4 | require 'sidekiq/priority_queue/api' 5 | require 'sidekiq/priority_queue/client' 6 | require 'sidekiq/priority_queue/combined_fetch' 7 | require 'sidekiq/priority_queue/fetch' 8 | require 'sidekiq/priority_queue/reliable_fetch' 9 | require 'sidekiq/priority_queue/scripts' 10 | require 'sidekiq/priority_queue/web' 11 | 12 | module Sidekiq 13 | module PriorityQueue 14 | end 15 | end 16 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: Tests 2 | on: 3 | pull_request: 4 | push: 5 | branches: 6 | - main 7 | jobs: 8 | tests: 9 | name: Run tests 10 | outputs: 11 | job-status: ${{ job.status }} 12 | runs-on: ubuntu-18.04 13 | timeout-minutes: 10 14 | services: 15 | redis: 16 | image: redis:alpine 17 | ports: 18 | - 6379:6379 19 | steps: 20 | - name: Checkout 21 | uses: actions/checkout@v2.3.1 22 | - name: Set up correct version of Ruby 23 | uses: actions/setup-ruby@v1 24 | with: 25 | ruby-version: 2.7 26 | - name: Install dependencies via Bundler 27 | run: bundle install --jobs 4 --retry 3 28 | - name: Run tests 29 | run: bundle exec rake 30 | -------------------------------------------------------------------------------- /test/test_api.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | class TestApi < Sidekiq::Test 5 | describe 'Queue' do 6 | 7 | before do 8 | Sidekiq.redis = { :url => REDIS_URL } 9 | Sidekiq.redis do |conn| 10 | conn.flushdb 11 | conn.sadd('priority-queues', 'priority-queue:foo') 12 | conn.zadd('priority-queue:foo', 0, 'blah') 13 | conn.zadd("priority-queue-counts:foo", 1, 'blah') 14 | end 15 | end 16 | 17 | it 'works' do 18 | assert_equal 1, Sidekiq::PriorityQueue::Queue.all.size 19 | assert_equal 1, Sidekiq::PriorityQueue::Queue.all.first.size 20 | end 21 | 22 | it 'can enumerate jobs' do 23 | assert_equal ["blah"], Sidekiq::PriorityQueue::Queue.new('foo').first.args 24 | end 25 | 26 | end 27 | end 28 | -------------------------------------------------------------------------------- /sidekiq-priority_queue.gemspec: -------------------------------------------------------------------------------- 1 | Gem::Specification.new do |s| 2 | s.name = 'sidekiq-priority_queue' 3 | s.version = '1.0.6' 4 | s.date = '2018-07-31' 5 | s.summary = "Priority Queuing for Sidekiq" 6 | s.description = "An extension for Sidekiq allowing jobs in a single queue to be executed by a priority score rather than FIFO" 7 | s.authors = ["Jacob Matthews", "Petr Kopac"] 8 | s.email = 'petr@chartmogul.com' 9 | s.files = `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features|pkg)/}) } 10 | s.homepage = 'https://github.com/chartmogul/sidekiq-priority_queue' 11 | s.license = 'MIT' 12 | s.required_ruby_version = '>= 2.5.0' 13 | 14 | s.add_dependency 'sidekiq', '>= 6.2.2' 15 | s.add_development_dependency 'minitest', '~> 5.10', '>= 5.10.1' 16 | end 17 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/scripts.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | 4 | module Sidekiq 5 | module PriorityQueue 6 | module Scripts 7 | 8 | ZPOPMIN = %q( 9 | local resp = redis.call('zrange', KEYS[1], '0', '0') 10 | if (resp[1] ~= nil) then 11 | local val = resp[# resp] 12 | redis.call('zrem', KEYS[1], val) 13 | return val 14 | else 15 | return false 16 | end 17 | ) 18 | 19 | ZPOPMIN_SADD = %q( 20 | local resp = redis.call('zrange', KEYS[1], '0', '0') 21 | if (resp[1] ~= nil) then 22 | local val = resp[# resp] 23 | redis.call('zrem', KEYS[1], val) 24 | redis.call('sadd', KEYS[2], val) 25 | return val 26 | else 27 | return false 28 | end 29 | ) 30 | 31 | end 32 | end 33 | end 34 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/web/views/_paging.erb: -------------------------------------------------------------------------------- 1 | 2 | <% if @total_size > @count %> 3 | 24 | <% end %> -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/testing.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | # Don't require in production code. 4 | # This disables the middleware and falls back to normal push, meaning in tests it will use inline/fake mode. 5 | # Prioritization doesn't make any sense in inline or fake tests anyway. 6 | module Sidekiq 7 | module PriorityQueue 8 | module TestingClient 9 | def call(worker_class, item, queue, redis_pool) 10 | testing_verify_subqueue(item) if item['subqueue'] && !item['priority'] 11 | yield # continue pushing the normal Sidekiq way 12 | end 13 | 14 | # Help testing the lambda; raise in case it's invalid. 15 | def testing_verify_subqueue(item) 16 | subqueue = resolve_subqueue(item['subqueue'], item['args']) 17 | serialized = "#{subqueue}" 18 | 19 | raise "subqueue shouldn't be nil" if subqueue.nil? 20 | raise "subqueue shouldn't be empty" if serialized == "" 21 | end 22 | end 23 | 24 | Sidekiq::PriorityQueue::Client.prepend TestingClient 25 | end 26 | end -------------------------------------------------------------------------------- /Gemfile.lock: -------------------------------------------------------------------------------- 1 | PATH 2 | remote: . 3 | specs: 4 | sidekiq-priority_queue (1.0.5) 5 | sidekiq (>= 6.2.2) 6 | 7 | GEM 8 | remote: https://rubygems.org/ 9 | specs: 10 | byebug (11.1.3) 11 | coderay (1.1.3) 12 | connection_pool (2.2.5) 13 | docile (1.3.4) 14 | method_source (1.0.0) 15 | minitest (5.14.3) 16 | pry (0.13.1) 17 | coderay (~> 1.1) 18 | method_source (~> 1.0) 19 | pry-byebug (3.9.0) 20 | byebug (~> 11.0) 21 | pry (~> 0.13.0) 22 | rack (2.2.3) 23 | rack-test (1.1.0) 24 | rack (>= 1.0, < 3) 25 | rake (13.0.3) 26 | redis (4.5.1) 27 | sidekiq (6.2.2) 28 | connection_pool (>= 2.2.2) 29 | rack (~> 2.0) 30 | redis (>= 4.2.0) 31 | simplecov (0.21.2) 32 | docile (~> 1.1) 33 | simplecov-html (~> 0.11) 34 | simplecov_json_formatter (~> 0.1) 35 | simplecov-html (0.12.3) 36 | simplecov_json_formatter (0.1.2) 37 | 38 | PLATFORMS 39 | ruby 40 | 41 | DEPENDENCIES 42 | minitest 43 | pry-byebug 44 | rack-test 45 | rake 46 | sidekiq-priority_queue! 47 | simplecov 48 | 49 | BUNDLED WITH 50 | 2.2.3 51 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/web/views/priority_queues.erb: -------------------------------------------------------------------------------- 1 |

<%= t('Queues') %>

2 | 3 |
4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | <% @queues.each do |queue| %> 12 | 13 | 16 | 17 | 18 | 24 | 25 | <% end %> 26 |
<%= t('Queue') %><%= t('Size') %><%= t('Latency') %><%= t('Actions') %>
14 | <%= h queue.name %> 15 | <%= queue.size %> <%# number_with_delimiter(queue.latency.round(2)) %> 19 |
20 | <%= csrf_tag %> 21 | 22 |
23 |
27 |
28 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/combined_fetch.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module Sidekiq 4 | module PriorityQueue 5 | class CombinedFetch 6 | attr_reader :fetches 7 | 8 | def initialize(fetches = []) 9 | @fetches = fetches 10 | end 11 | 12 | def retrieve_work 13 | fetches.each do |fetch| 14 | work = fetch.retrieve_work 15 | return work if work 16 | end 17 | end 18 | 19 | def self.configure(&block) 20 | combined_fetch = self.new 21 | yield combined_fetch 22 | 23 | combined_fetch 24 | end 25 | 26 | def add(fetch) 27 | fetches << fetch 28 | end 29 | 30 | def bulk_requeue(inprogress, options) 31 | # ReliableFetch#bulk_equeue ignores inprogress, so it's safe to call both 32 | fetches.each do |f| 33 | if [Fetch, ReliableFetch].any? { |klass| f.is_a?(klass) } 34 | jobs_to_requeue = inprogress.select{|uow| uow.queue.start_with?('priority-queue:') } 35 | f.bulk_requeue(jobs_to_requeue, options) 36 | else 37 | jobs_to_requeue = inprogress.reject{|uow| uow.queue.start_with?('priority-queue:') } 38 | f.bulk_requeue(jobs_to_requeue, options) 39 | end 40 | end 41 | end 42 | end 43 | end 44 | end 45 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/web.rb: -------------------------------------------------------------------------------- 1 | require 'sidekiq/web' 2 | 3 | 4 | module Sidekiq::PriorityQueue 5 | module Web 6 | 7 | ROOT = File.expand_path('../web', __FILE__) 8 | 9 | def self.registered(app) 10 | app.tabs['Priority Queues'] = 'priority_queues' 11 | 12 | app.get '/priority_queues' do 13 | @queues = Queue.all 14 | render(:erb, File.read("#{ROOT}/views/priority_queues.erb")) 15 | end 16 | 17 | app.get '/priority_queues/:name' do 18 | @name = route_params[:name] 19 | halt(404) unless @name 20 | 21 | @count = (params['count'] || 25).to_i 22 | @queue = Sidekiq::Queue.new(@name) 23 | (@current_page, @total_size, @messages) = page("priority-queue:#{@name}", params['page'], @count) 24 | @subqueue_counts = Sidekiq.redis do |con| 25 | con.zrevrange("priority-queue-counts:#{@name}", 0, params['subqueue_count'] || 10, withscores: true) 26 | end.map { |name, count| SubqueueCount.new(name, count) } 27 | 28 | @messages = @messages.map{ |msg| Job.new(msg.first, @name, msg.last) } 29 | render(:erb, File.read("#{ROOT}/views/priority_queue.erb")) 30 | end 31 | 32 | app.post "/priority_queues/:name/delete" do 33 | name = route_params[:name] 34 | Job.new(params['key_val'], name).delete 35 | redirect_with_query("#{root_path}priority_queues/#{CGI.escape(name)}") 36 | end 37 | 38 | end 39 | end 40 | end 41 | 42 | ::Sidekiq::Web.register Sidekiq::PriorityQueue::Web 43 | -------------------------------------------------------------------------------- /test/test_web.rb: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | # frozen_string_literal: true 3 | require_relative 'helper' 4 | require 'sidekiq/web' 5 | require 'rack/test' 6 | 7 | class TestWeb < Sidekiq::Test 8 | describe 'sidekiq web' do 9 | include Rack::Test::Methods 10 | 11 | def app 12 | Sidekiq::Web 13 | end 14 | 15 | class PrioritizedWebWorker 16 | include Sidekiq::Worker 17 | sidekiq_options subqueue: ->(args){ args[0] } 18 | 19 | def perform(a,b,c) 20 | end 21 | end 22 | 23 | job = { 'jid' => 'blah', 'class' => 'FakeWorker', 'args' => [1,2,3], 'subqueue' => 1 } 24 | job_2 = { 'jid' => 'blahah', 'class' => 'FakeWorkerDef', 'args' => [1,2,3,4], 'subqueue' => 1 } 25 | 26 | before do 27 | Sidekiq.redis = { :url => REDIS_URL } 28 | Sidekiq.redis do |conn| 29 | conn.flushdb 30 | conn.sadd('priority-queues', 'priority-queue:default') 31 | conn.zadd('priority-queue:default', 0, job.to_json) 32 | conn.zadd("priority-queue-counts:default", 2, job_2['subqueue']) 33 | end 34 | end 35 | 36 | it 'can display queues' do 37 | get '/priority_queues' 38 | assert_equal 200, last_response.status 39 | assert_match(/default/, last_response.body) 40 | end 41 | 42 | it 'can display queue' do 43 | get '/priority_queues/default' 44 | assert_equal 200, last_response.status 45 | assert_match(/FakeWorker/, last_response.body) 46 | end 47 | 48 | it 'can delete job' do 49 | post '/priority_queues/default/delete', key_val: job.to_json 50 | assert_equal 302, last_response.status 51 | assert_equal 0, Sidekiq::PriorityQueue::Queue.new('default').size 52 | end 53 | 54 | end 55 | end 56 | -------------------------------------------------------------------------------- /test/test_fetch.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | class TestFetcher < Sidekiq::Test 5 | describe 'fetcher' do 6 | job = {'jid' => 'blah', 'args' => [1,2,3], 'subqueue' => 1 } 7 | 8 | before do 9 | Sidekiq.redis = { :url => REDIS_URL } 10 | Sidekiq.redis do |conn| 11 | conn.flushdb 12 | conn.zadd('priority-queue:foo', 0, job.to_json) 13 | conn.zadd("priority-queue-counts:foo", 1, job['subqueue']) 14 | end 15 | end 16 | 17 | after do 18 | Sidekiq.redis = REDIS 19 | end 20 | 21 | it 'retrieves' do 22 | fetch = Sidekiq::PriorityQueue::Fetch.new(:queues => ['foo']) 23 | uow = fetch.retrieve_work 24 | refute_nil uow 25 | assert_equal 'foo', uow.queue_name 26 | assert_equal job.to_json, uow.job 27 | q = Sidekiq::PriorityQueue::Queue.new('foo') 28 | assert_equal 0, q.size 29 | uow.requeue 30 | assert_equal 1, q.size 31 | assert uow.acknowledge 32 | Sidekiq.redis do |conn| 33 | assert_nil conn.zscore("priority-queue-counts:foo", job['subqueue']) 34 | end 35 | end 36 | 37 | it 'retrieves with strict setting' do 38 | fetch = Sidekiq::PriorityQueue::Fetch.new(:queues => ['basic', 'bar', 'bar'], :strict => true) 39 | cmd = fetch.queues_cmd 40 | assert_equal cmd, ['priority-queue:basic', 'priority-queue:bar'] 41 | end 42 | 43 | it 'bulk requeues only priority-queue jobs' do 44 | q1 = Sidekiq::PriorityQueue::Queue.new('foo') 45 | q2 = Sidekiq::PriorityQueue::Queue.new('bar') 46 | assert_equal 1, q1.size 47 | assert_equal 0, q2.size 48 | uow = Sidekiq::PriorityQueue::Fetch::UnitOfWork 49 | Sidekiq::PriorityQueue::Fetch.new(queues: []).bulk_requeue( 50 | [ uow.new('priority-queue:foo', 'bob'), uow.new('fuzzy:queue:foo', 'bar') ], 51 | queues: [] 52 | ) 53 | assert_equal 2, q1.size 54 | assert_equal 0, q2.size 55 | end 56 | 57 | end 58 | end 59 | -------------------------------------------------------------------------------- /test/test_client.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | class TestClient < Sidekiq::Test 5 | 6 | describe 'client' do 7 | class Worker 8 | include Sidekiq::Worker 9 | end 10 | 11 | it 'allows standard sidekiq functionality' do 12 | Sidekiq.redis {|c| c.flushdb } 13 | assert Worker.perform_async(1) 14 | assert_equal 1, Sidekiq::Queue.new().size 15 | end 16 | 17 | it 'pushes to a priority queue and not a normal queue' do 18 | Sidekiq.redis {|c| c.flushdb } 19 | assert Worker.set(priority: 0).perform_async(1) 20 | assert_equal 0, Sidekiq::Queue.new().size 21 | q = Sidekiq::PriorityQueue::Queue.new() 22 | assert_equal 1, q.size 23 | assert_equal 0, q.first.priority 24 | end 25 | 26 | class PrioritizedWorker 27 | include Sidekiq::Worker 28 | sidekiq_options subqueue: ->(args){ args[0] } 29 | end 30 | 31 | it 'prioritises based on already enqueued jobs for the same key' do 32 | Sidekiq.redis {|c| c.flushdb } 33 | # NOTE: The ordering of keys with the same score is lexicographical: https://redis.io/commands/zrange 34 | jobs_with_expected_priority = [ ['a',1], ['b',1], ['a',2] ] 35 | jobs_with_expected_priority.each{|arg,_| PrioritizedWorker.perform_async(arg) } 36 | 37 | queue = Sidekiq::PriorityQueue::Queue.new 38 | assert_equal 3, queue.size 39 | 40 | assert_equal jobs_with_expected_priority, queue.map{ |q| [q.subqueue, q.priority] } 41 | end 42 | 43 | class SubqueuePersistedWorker 44 | include Sidekiq::Worker 45 | sidekiq_options subqueue: 123 # emualte the persisted subqueue argument 46 | end 47 | 48 | it 'works when subqueue argument is already resolved / persisted' do 49 | Sidekiq.redis {|c| c.flushdb } 50 | 51 | SubqueuePersistedWorker.perform_async(123) 52 | 53 | queue = Sidekiq::PriorityQueue::Queue.new 54 | assert_equal 1, queue.size 55 | assert_equal [[123, 1]], queue.map{ |q| [q.subqueue, q.priority] } 56 | end 57 | end 58 | end 59 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/client.rb: -------------------------------------------------------------------------------- 1 | module Sidekiq 2 | module PriorityQueue 3 | class Client 4 | 5 | # inserted into Sidekiq's Client as middleware 6 | def call(worker_class, item, queue, redis_pool) 7 | if item['priority'] 8 | sadd('priority-queues', queue) 9 | zadd(queue, item['priority'], item) 10 | return item['jid'] 11 | elsif item['subqueue'] 12 | # replace the proc with what it returns 13 | sadd('priority-queues', queue) 14 | item['subqueue'] = resolve_subqueue(item['subqueue'], item['args']) 15 | priority = fetch_and_add(queue, item['subqueue'], item) 16 | zadd(queue, priority, item) 17 | return item['jid'] 18 | else 19 | # continue pushing the normal Sidekiq way 20 | yield 21 | end 22 | end 23 | 24 | private 25 | 26 | def resolve_subqueue(subqueue, job_args) 27 | return subqueue unless subqueue.respond_to?(:call) 28 | 29 | subqueue.call(job_args) 30 | end 31 | 32 | def zadd(queue, score, item) 33 | Sidekiq.redis do |conn| 34 | queue = "priority-queue:#{queue}" 35 | conn.zadd(queue, score, item.to_json) 36 | return item 37 | end 38 | end 39 | 40 | def sadd(set, member) 41 | Sidekiq.redis do |conn| 42 | conn.sadd(set,member) 43 | end 44 | end 45 | 46 | def fetch_and_add(queue, subqueue, item) 47 | Sidekiq.redis do |conn| 48 | priority = conn.zincrby("priority-queue-counts:#{queue}", 1, subqueue) 49 | end 50 | end 51 | end 52 | end 53 | end 54 | 55 | Sidekiq::Client.class_eval do 56 | def push(item) 57 | normed = normalize_item(item) 58 | payload = process_single(item['class'], normed) 59 | 60 | # if payload is a JID because the middleware already pushed then just return the JID 61 | return payload if payload.is_a?(String) 62 | 63 | if payload 64 | raw_push([payload]) 65 | payload['jid'] 66 | end 67 | end 68 | end 69 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/api.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'sidekiq/api' 3 | 4 | 5 | module Sidekiq 6 | module PriorityQueue 7 | class Queue 8 | include Enumerable 9 | 10 | attr_reader :name 11 | 12 | def initialize(name='default') 13 | @name = name 14 | @rname = "priority-queue:#{name}" 15 | end 16 | 17 | def size 18 | Sidekiq.redis { |con| con.zcard(@rname) } 19 | end 20 | 21 | def each 22 | initial_size = size 23 | deleted_size = 0 24 | page = 0 25 | page_size = 50 26 | 27 | while true do 28 | range_start = page * page_size - deleted_size 29 | range_end = range_start + page_size - 1 30 | entries = Sidekiq.redis do |conn| 31 | conn.zrange @rname, range_start, range_end, withscores: true 32 | end 33 | break if entries.empty? 34 | page += 1 35 | entries.each do |entry, priority| 36 | yield Job.new(entry, @name, priority) 37 | end 38 | deleted_size = initial_size - size 39 | end 40 | end 41 | 42 | def self.all 43 | Sidekiq.redis { |con| con.smembers('priority-queues') } 44 | .map{ |key| key.gsub('priority-queue:', '') } 45 | .sort 46 | .map { |q| Queue.new(q) } 47 | end 48 | 49 | end 50 | 51 | SubqueueCount = Struct.new(:name, :size) 52 | 53 | class Job < Sidekiq::JobRecord 54 | 55 | attr_reader :priority 56 | attr_reader :subqueue 57 | 58 | def initialize(item, queue_name = nil, priority = nil) 59 | @args = nil 60 | @value = item 61 | @item = item.is_a?(Hash) ? item : parse(item) 62 | @queue = queue_name || @item['queue'] 63 | @subqueue = @item['subqueue'] 64 | @priority = priority 65 | end 66 | 67 | def delete 68 | count = Sidekiq.redis do |conn| 69 | conn.zrem("priority-queue:#{@queue}", @value) 70 | end 71 | count != 0 72 | end 73 | end 74 | end 75 | end 76 | -------------------------------------------------------------------------------- /test/test_combined_fetch.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | require 'sidekiq/fetch' 4 | 5 | class TestFetcher < Sidekiq::Test 6 | describe 'fetcher' do 7 | normal_job = {'jid' => 'normal_job', 'args' => [1,2,3] } 8 | priority_job = {'jid' => 'priority_job', 'args' => [1,2,3], 'subqueue' => 1 } 9 | 10 | before do 11 | Sidekiq.redis = { :url => REDIS_URL } 12 | Sidekiq.redis do |conn| 13 | conn.flushdb 14 | conn.lpush('queue:foo', normal_job.to_json) 15 | conn.zadd('priority-queue:foo', 0, priority_job.to_json) 16 | conn.zadd("priority-queue-counts:foo", 1, priority_job['subqueue']) 17 | end 18 | end 19 | 20 | after do 21 | Sidekiq.redis = REDIS 22 | end 23 | 24 | it 'retrieves from both normal and priority queues' do 25 | fetch = Sidekiq::PriorityQueue::CombinedFetch.configure do |fetches| 26 | fetches.add Sidekiq::BasicFetch.new(queues: ['foo']) 27 | fetches.add Sidekiq::PriorityQueue::Fetch.new(queues: ['foo']) 28 | end 29 | 30 | uow = fetch.retrieve_work 31 | refute_nil uow 32 | assert_equal 'foo', uow.queue_name 33 | assert_equal normal_job.to_json, uow.job 34 | 35 | uow = fetch.retrieve_work 36 | refute_nil uow 37 | assert_equal 'foo', uow.queue_name 38 | assert_equal priority_job.to_json, uow.job 39 | end 40 | 41 | it 'bulk requeues all jobs only once' do 42 | fetch = Sidekiq::PriorityQueue::CombinedFetch.configure do |fetches| 43 | fetches.add Sidekiq::BasicFetch.new(queues: ['foo'], index: 0) 44 | fetches.add Sidekiq::PriorityQueue::Fetch.new(queues: ['foo'], index: 0) 45 | end 46 | 47 | q1 = Sidekiq::PriorityQueue::Queue.new('foo') 48 | q2 = Sidekiq::Queue.new('foo') 49 | assert_equal 1, q1.size 50 | assert_equal 1, q2.size 51 | uow = Sidekiq::PriorityQueue::Fetch::UnitOfWork 52 | 53 | Sidekiq.redis do |conn| 54 | conn.sadd("priority-queue:foo_#{Socket.gethostname}_0", 'bob') 55 | end 56 | 57 | fetch.bulk_requeue( 58 | [ uow.new('priority-queue:foo', 'bob'), uow.new('queue:foo', 'bar') ], 59 | { queues: ['foo'], index: 0 } 60 | ) 61 | 62 | assert_equal 2, q1.size 63 | assert_equal 2, q2.size 64 | end 65 | end 66 | end 67 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Sidekiq Priority Queue 2 | ============== 3 | Extends Sidekiq with support for queuing jobs with a fine grained priority and emulating multiple queues using a single Redis sorted set, ideal for multi-tenant applications. 4 | 5 | The standard Sidekiq setup performs really well using Redis lists, but lists can only be strict FIFO queues, which can be hugely problematic when they process slowly and one user may need to wait hours behind a backlog of jobs. 6 | 7 | Sidekiq Priority Queue offers a plug-in solution retaining the simplicity and performance of Sidekiq. The priority queue is a building block for emulating sub-queues (per tenant or user) by de-prioritising jobs according to how many jobs are already in this sub-queue. 8 | 9 | Sources of inspiration are naturally Sidekiq itself, the fantastic Redis documentation, and https://github.com/gocraft/work 10 | 11 | Installation 12 | ----------------- 13 | 14 | gem install sidekiq-priority_queue 15 | 16 | Configuration 17 | ----------------- 18 | ``` 19 | Sidekiq.configure_server do |config| 20 | config.options[:fetch] = Sidekiq::PriorityQueue::Fetch 21 | end 22 | 23 | Sidekiq.configure_client do |config| 24 | config.client_middleware do |chain| 25 | chain.add Sidekiq::PriorityQueue::Client 26 | end 27 | end 28 | ``` 29 | Usage 30 | ----------------- 31 | ``` 32 | class Worker 33 | include Sidekiq::Worker 34 | sidekiq_options priority: 1000 35 | end 36 | ``` 37 | Alternatively, you can split jobs into subqueues (via a proc) which are deprioritised based on the subqueue size: 38 | ``` 39 | class Worker 40 | include Sidekiq::Worker 41 | 42 | # args[0] will take the `user_id` argument below, and assign priority dynamically. 43 | sidekiq_options subqueue: ->(args){ args[0] } 44 | 45 | def perform(user_id, other_args) 46 | # do jobs 47 | end 48 | end 49 | ``` 50 | 51 | Testing 52 | ----------------- 53 | For example in your `spec/rails_helper.rb` you can include this line: 54 | ```ruby 55 | require 'sidekiq/priority_queue/testing' 56 | ``` 57 | next to the call to `Sidekiq::Testing.inline!`. It disables the feature for testing and falls back to `inline`/`fake` modes. 58 | If you accidentally require this in production code, it will likewise fall back to normal Sidekiq scheduling. 59 | 60 | Development 61 | ----------------- 62 | - Run `docker-compose up -d` to start up a temporary redis instance. 63 | - Run `bundle install` to install dependencies. 64 | - Run the tests with `bundle exec rake` 65 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/fetch.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'sidekiq' 3 | 4 | module Sidekiq 5 | module PriorityQueue 6 | class Fetch 7 | 8 | UnitOfWork = Struct.new(:queue, :job) do 9 | def acknowledge 10 | Sidekiq.redis do |conn| 11 | unless subqueue.nil? 12 | count = conn.zincrby(subqueue_counts, -1, subqueue) 13 | conn.zrem(subqueue_counts, subqueue) if count < 1 14 | end 15 | end 16 | end 17 | 18 | def queue_name 19 | queue.sub(/.*queue:/, '') 20 | end 21 | 22 | def subqueue 23 | @parsed_job ||= JSON.parse(job) 24 | @parsed_job['subqueue'] 25 | end 26 | 27 | def subqueue_counts 28 | "priority-queue-counts:#{queue_name}" 29 | end 30 | 31 | def requeue 32 | Sidekiq.redis do |conn| 33 | conn.zadd(queue, 0, job) 34 | end 35 | end 36 | end 37 | 38 | def initialize(options) 39 | @strictly_ordered_queues = !!options[:strict] 40 | @queues = options[:queues].map { |q| "priority-queue:#{q}" } 41 | @queues = @queues.uniq if @strictly_ordered_queues 42 | end 43 | 44 | def retrieve_work 45 | work = @queues.detect{ |q| job = zpopmin(q); break [q,job] if job } 46 | UnitOfWork.new(*work) if work 47 | end 48 | 49 | def zpopmin(queue) 50 | Sidekiq.redis do |con| 51 | @script_sha ||= con.script(:load, Sidekiq::PriorityQueue::Scripts::ZPOPMIN) 52 | con.evalsha(@script_sha, [queue]) 53 | end 54 | end 55 | 56 | def queues_cmd 57 | if @strictly_ordered_queues 58 | @queues 59 | else 60 | @queues.shuffle.uniq 61 | end 62 | end 63 | 64 | def bulk_requeue(inprogress, options) 65 | return if inprogress.empty? 66 | 67 | Sidekiq.logger.debug { "Re-queueing terminated jobs" } 68 | jobs_to_requeue = {} 69 | inprogress.each do |unit_of_work| 70 | jobs_to_requeue[unit_of_work.queue] ||= [] 71 | jobs_to_requeue[unit_of_work.queue] << unit_of_work.job 72 | end 73 | 74 | Sidekiq.redis do |conn| 75 | conn.pipelined do 76 | jobs_to_requeue.each do |queue, jobs| 77 | conn.zadd(queue, jobs.map{|j| [0,j] }) 78 | end 79 | end 80 | end 81 | Sidekiq.logger.info("Pushed #{inprogress.size} jobs back to Redis") 82 | rescue => ex 83 | Sidekiq.logger.warn("Failed to requeue #{inprogress.size} jobs: #{ex.message}") 84 | end 85 | end 86 | end 87 | end 88 | -------------------------------------------------------------------------------- /test/helper.rb: -------------------------------------------------------------------------------- 1 | # Copyright (c) Contributed Systems LLC, , ChartMogul Ltd 2 | 3 | # Sidekiq-priority_queue is an Open Source project licensed under the terms of 4 | # the LGPLv3 license. Please see 5 | # for license text. 6 | 7 | # frozen_string_literal: true 8 | $TESTING = true 9 | # disable minitest/parallel threads 10 | ENV["MT_CPU"] = "0" 11 | 12 | if ENV["COVERAGE"] 13 | require 'simplecov' 14 | SimpleCov.start do 15 | add_filter "/test/" 16 | add_filter "/myapp/" 17 | end 18 | end 19 | ENV['RACK_ENV'] = ENV['RAILS_ENV'] = 'test' 20 | 21 | trap 'TSTP' do 22 | threads = Thread.list 23 | 24 | puts 25 | puts "=" * 80 26 | puts "Received TSTP signal; printing all #{threads.count} thread backtraces." 27 | 28 | threads.each do |thr| 29 | description = thr == Thread.main ? "Main thread" : thr.inspect 30 | puts 31 | puts "#{description} backtrace: " 32 | puts thr.backtrace.join("\n") 33 | end 34 | 35 | puts "=" * 80 36 | end 37 | 38 | require 'pry-byebug' 39 | require 'minitest/autorun' 40 | 41 | require 'sidekiq' 42 | require 'sidekiq/api' 43 | require 'sidekiq/priority_queue' 44 | Sidekiq.logger.level = Logger::ERROR 45 | 46 | Sidekiq::Test = Minitest::Test 47 | 48 | require 'sidekiq/redis_connection' 49 | REDIS_URL = ENV['REDIS_URL'] || 'redis://localhost/15' 50 | REDIS = Sidekiq::RedisConnection.create(:url => REDIS_URL) 51 | 52 | Sidekiq.configure_client do |config| 53 | config.redis = { :url => REDIS_URL } 54 | end 55 | 56 | 57 | Sidekiq.configure_client do |config| 58 | config.client_middleware do |chain| 59 | chain.add Sidekiq::PriorityQueue::Client 60 | end 61 | end 62 | 63 | def capture_logging(lvl=Logger::INFO) 64 | old = Sidekiq.logger 65 | begin 66 | out = StringIO.new 67 | logger = Logger.new(out) 68 | logger.level = lvl 69 | Sidekiq.logger = logger 70 | yield 71 | out.string 72 | ensure 73 | Sidekiq.logger = old 74 | end 75 | end 76 | 77 | def with_logging(lvl=Logger::DEBUG) 78 | old = Sidekiq.logger.level 79 | begin 80 | Sidekiq.logger.level = lvl 81 | yield 82 | ensure 83 | Sidekiq.logger.level = old 84 | end 85 | end 86 | 87 | def reset_sidekiq_lifecycle_events 88 | Sidekiq.options[:lifecycle_events] = { 89 | startup: [], 90 | quiet: [], 91 | shutdown: [], 92 | heartbeat: [] 93 | } 94 | end 95 | 96 | def setup_reliable_fetcher(queues = ['foo']) 97 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: queues, index: 0) 98 | fetch.setup 99 | SidekiqUtilInstance.new.fire_event(:startup) 100 | fetch 101 | end 102 | 103 | class SidekiqUtilInstance 104 | include Sidekiq::Util 105 | end 106 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/web/views/priority_queue.erb: -------------------------------------------------------------------------------- 1 |
2 |
3 |

4 | <%= t('CurrentMessagesInQueue', :queue => h(@name)) %> 5 | 6 | <%= number_with_delimiter(@total_size) %> 7 |

8 |
9 |
10 | <%= erb :_paging, locals: { url: "#{root_path}priority_queues/#{CGI.escape(@name)}" } %> 11 |
12 |
13 |
14 |
15 |

Biggest subqueues

16 | 17 | 18 |
19 |
20 | 21 | 22 | 23 | 24 | 25 | <% @subqueue_counts.each do |subqueue_count| %> 26 | 27 | 28 | 29 | 30 | <% end %> 31 |
Subqueue nameSubqueue size
<%= subqueue_count.name %><%= subqueue_count.size %>
32 |
33 |
34 | 35 | 36 |
37 |
38 |
39 |
40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | <% @messages.each_with_index do |msg, index| %> 49 | 50 | 51 | 61 | 62 | 63 | 70 | 71 | <% end %> 72 |
<%= t('Job') %><%= t('Arguments') %><%= t('Priority') %><%= t('Subqueue') %>
<%= h(msg.display_class) %> 52 | <% a = msg.display_args %> 53 | <% if a.inspect.size > 100 %> 54 | <%= h(a.inspect[0..100]) + "... " %> 55 | 56 |
<%= display_args(a) %>
57 | <% else %> 58 | <%= display_args(msg.display_args) %> 59 | <% end %> 60 | 		<%= msg.priority %><%= msg.subqueue %> 64 |
65 | <%= csrf_tag %> 66 | 67 | 68 |
69 |
73 |
74 | <%= erb :_paging, locals: { url: "#{root_path}priority_queues/#{@name}" } %> -------------------------------------------------------------------------------- /bin/sidekiqload: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env ruby 2 | 3 | # Copyright (c) Contributed Systems LLC, ChartMogul Ltd 4 | 5 | # Sidekiq-priority_queue is an Open Source project licensed under the terms of 6 | # the LGPLv3 license. Please see 7 | # for license text. 8 | 9 | # Quiet some warnings we see when running in warning mode: 10 | # RUBYOPT=-w bundle exec sidekiq 11 | $TESTING = false 12 | 13 | #require 'ruby-prof' 14 | Bundler.require(:default) 15 | 16 | require 'sidekiq/cli' 17 | require 'sidekiq/launcher' 18 | require 'sidekiq/priority_queue' 19 | 20 | include Sidekiq::Util 21 | 22 | Sidekiq.configure_server do |config| 23 | #config.options[:concurrency] = 1 24 | config.redis = { db: 13 } 25 | config.options[:queues] << 'default' 26 | config.logger.level = Logger::ERROR 27 | config.average_scheduled_poll_interval = 2 28 | config.options[:fetch] = Sidekiq::PriorityQueue::Fetch 29 | end 30 | 31 | Sidekiq.configure_client do |config| 32 | config.client_middleware do |chain| 33 | chain.add Sidekiq::PriorityQueue::Client 34 | end 35 | end 36 | 37 | class LoadWorker 38 | include Sidekiq::Worker 39 | sidekiq_options retry: 1, subqueue: ->(args){ args[0] } 40 | sidekiq_retry_in do |x| 41 | 1 42 | end 43 | 44 | def perform(idx) 45 | #raise idx.to_s if idx % 100 == 1 46 | end 47 | end 48 | 49 | # brew tap shopify/shopify 50 | # brew install toxiproxy 51 | # gem install toxiproxy 52 | #require 'toxiproxy' 53 | # simulate a non-localhost network for realer-world conditions. 54 | # adding 1ms of network latency has an ENORMOUS impact on benchmarks 55 | #Toxiproxy.populate([{ 56 | #"name": "redis", 57 | #"listen": "127.0.0.1:6380", 58 | #"upstream": "127.0.0.1:6379" 59 | #}]) 60 | 61 | self_read, self_write = IO.pipe 62 | %w(INT TERM TSTP TTIN).each do |sig| 63 | begin 64 | trap sig do 65 | self_write.puts(sig) 66 | end 67 | rescue ArgumentError 68 | puts "Signal #{sig} not supported" 69 | end 70 | end 71 | 72 | Sidekiq.redis {|c| c.flushdb} 73 | def handle_signal(launcher, sig) 74 | Sidekiq.logger.debug "Got #{sig} signal" 75 | case sig 76 | when 'INT' 77 | # Handle Ctrl-C in JRuby like MRI 78 | # http://jira.codehaus.org/browse/JRUBY-4637 79 | raise Interrupt 80 | when 'TERM' 81 | # Heroku sends TERM and then waits 10 seconds for process to exit. 82 | raise Interrupt 83 | when 'TSTP' 84 | Sidekiq.logger.info "Received TSTP, no longer accepting new work" 85 | launcher.quiet 86 | when 'TTIN' 87 | Thread.list.each do |thread| 88 | Sidekiq.logger.warn "Thread TID-#{(thread.object_id ^ ::Process.pid).to_s(36)} #{thread['label']}" 89 | if thread.backtrace 90 | Sidekiq.logger.warn thread.backtrace.join("\n") 91 | else 92 | Sidekiq.logger.warn "" 93 | end 94 | end 95 | end 96 | end 97 | 98 | def Process.rss 99 | `ps -o rss= -p #{Process.pid}`.chomp.to_i 100 | end 101 | 102 | 103 | Sidekiq.redis do |con| 104 | count = 100_000 105 | 106 | count.times do |idx| 107 | #TODO why does Sidekiq::Client.push not work as expected? 108 | con.zadd('priority-queue:default', idx, { 'class' => LoadWorker, 'args' => [idx] }.to_json) 109 | end 110 | end 111 | Sidekiq.logger.error "Created #{ Sidekiq::PriorityQueue::Queue.new().size } jobs" 112 | 113 | Monitoring = Thread.new do 114 | watchdog("monitor thread") do 115 | while true 116 | sleep 1 117 | qsize, retries = Sidekiq.redis do |conn| 118 | conn.pipelined do 119 | conn.zcard "priority-queue:default" 120 | conn.zcard "retry" 121 | end 122 | end.map(&:to_i) 123 | total = qsize + retries 124 | #GC.start 125 | Sidekiq.logger.error("RSS: #{Process.rss} Pending: #{total}") 126 | if total == 0 127 | Sidekiq.logger.error("Done") 128 | exit(0) 129 | end 130 | end 131 | end 132 | end 133 | 134 | begin 135 | #RubyProf::exclude_threads = [ Monitoring ] 136 | #RubyProf.start 137 | fire_event(:startup) 138 | #Sidekiq.logger.error "Simulating 1ms of latency between Sidekiq and redis" 139 | #Toxiproxy[:redis].downstream(:latency, latency: 1).apply do 140 | launcher = Sidekiq::Launcher.new(Sidekiq.options) 141 | launcher.run 142 | 143 | while readable_io = IO.select([self_read]) 144 | signal = readable_io.first[0].gets.strip 145 | handle_signal(launcher, signal) 146 | end 147 | #end 148 | rescue SystemExit => e 149 | #Sidekiq.logger.error("Profiling...") 150 | #result = RubyProf.stop 151 | #printer = RubyProf::GraphHtmlPrinter.new(result) 152 | #printer.print(File.new("output.html", "w"), :min_percent => 1) 153 | # normal 154 | rescue => e 155 | raise e if $DEBUG 156 | STDERR.puts e.message 157 | STDERR.puts e.backtrace.join("\n") 158 | exit 1 159 | end 160 | -------------------------------------------------------------------------------- /test/test_reliable_fetch.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | 5 | class TestFetcher < Sidekiq::Test 6 | describe 'reliable fetcher' do 7 | job = { 'jid' => 'blah', 'args' => [1, 2, 3], 'subqueue' => 1 } 8 | 9 | before do 10 | Sidekiq.redis = { url: REDIS_URL } 11 | Sidekiq.redis do |conn| 12 | conn.flushdb 13 | conn.zadd('priority-queue:foo', 0, job.to_json) 14 | conn.zadd('priority-queue-counts:foo', 1, job['subqueue']) 15 | end 16 | reset_sidekiq_lifecycle_events 17 | end 18 | 19 | after do 20 | Sidekiq.redis = REDIS 21 | end 22 | 23 | it 'adds three new event triggers when #setup is called' do 24 | assert_equal 0, Sidekiq.options[:lifecycle_events].values.flatten.size 25 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['foo'], index: 0) 26 | fetch.setup 27 | assert_equal 3, Sidekiq.options[:lifecycle_events].values.flatten.size 28 | end 29 | 30 | it 'stops picking up work once shutdown event is triggered' do 31 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['foo'], index: 0) 32 | fetch.setup 33 | 34 | SidekiqUtilInstance.new.fire_event(:shutdown) 35 | 36 | assert_equal 1, Sidekiq::PriorityQueue::Queue.new('foo').size 37 | assert_nil fetch.retrieve_work 38 | end 39 | 40 | it 'cleans up dead jobs on startup' do 41 | # First we have to mimic old jobs lying around 42 | previous_process_identity = "sidekiq-pipeline-32152-xfwvw:42251:#{SecureRandom.hex(6)}" 43 | priority_queue = 'priority-queue:bar' 44 | previous_wip_queue = "queue:spriorityq|#{previous_process_identity}|#{priority_queue}" 45 | 46 | Sidekiq.redis do |conn| 47 | conn.sadd('super_processes_priority', previous_process_identity) 48 | conn.sadd("#{previous_process_identity}:super_priority_queues", previous_wip_queue) 49 | conn.sadd(previous_wip_queue, job.to_json) 50 | end 51 | 52 | # Then we setup the fetcher 53 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['bar'], index: 0) 54 | fetch.setup 55 | 56 | # And here we test is really on startup it brings back old jobs 57 | assert_equal 0, Sidekiq::PriorityQueue::Queue.new('bar').size 58 | SidekiqUtilInstance.new.fire_event(:startup) 59 | assert_equal 1, Sidekiq::PriorityQueue::Queue.new('bar').size 60 | refute Sidekiq.redis { |c| c.exists?(previous_wip_queue) } 61 | end 62 | 63 | it 'registers process and private queues on startup' do 64 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['foo'], index: 0) 65 | fetch.setup 66 | 67 | SidekiqUtilInstance.new.fire_event(:startup) 68 | 69 | Sidekiq.redis do |conn| 70 | registered_processes = conn.smembers('super_processes_priority') 71 | assert_equal 1, registered_processes.size 72 | assert_equal fetch.identity, registered_processes.first 73 | private_queues = conn.smembers("#{registered_processes.first}:super_priority_queues") 74 | assert_equal 1, private_queues.size 75 | identity = registered_processes.first 76 | assert_equal "queue:spriorityq|#{identity}|priority-queue:foo", private_queues.first 77 | end 78 | end 79 | 80 | it 'registers process and private queues on heartbeat' do 81 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['foo'], index: 0) 82 | fetch.setup 83 | 84 | SidekiqUtilInstance.new.fire_event(:heartbeat) 85 | 86 | Sidekiq.redis do |conn| 87 | registered_processes = conn.smembers('super_processes_priority') 88 | assert_equal 1, registered_processes.size 89 | assert_equal fetch.identity, registered_processes.first 90 | private_queues = conn.smembers("#{registered_processes.first}:super_priority_queues") 91 | assert_equal 1, private_queues.size 92 | identity = registered_processes.first 93 | assert_equal "queue:spriorityq|#{identity}|priority-queue:foo", private_queues.first 94 | end 95 | end 96 | 97 | it 'retrieves and puts into private set' do 98 | fetch = setup_reliable_fetcher 99 | uow = fetch.retrieve_work 100 | refute_nil uow 101 | assert_equal 'foo', uow.queue_name 102 | assert_equal job.to_json, uow.job 103 | identity = fetch.identity 104 | wip_queue = "queue:spriorityq|#{identity}|priority-queue:foo" 105 | Sidekiq.redis { |conn| assert conn.sismember(wip_queue, job.to_json) } 106 | q = Sidekiq::PriorityQueue::Queue.new('foo') 107 | assert_equal 0, q.size 108 | assert uow.acknowledge 109 | Sidekiq.redis do |conn| 110 | assert_nil conn.zscore('priority-queue-counts:foo', job['subqueue']) 111 | assert !conn.sismember(wip_queue, job.to_json) 112 | end 113 | end 114 | 115 | it 'pushes WIP jobs back to the head of the sorted set' do 116 | assert_equal 1, Sidekiq::PriorityQueue::Queue.new('foo').size 117 | fetch = setup_reliable_fetcher 118 | identity = fetch.identity 119 | wip_queue = "queue:spriorityq|#{identity}|priority-queue:foo" 120 | 121 | killed_job = { 'jid' => 'blah_blah', 'args' => [1, 2, 3], 'subqueue' => 1 } 122 | Sidekiq.redis { |conn| conn.sadd(wip_queue, killed_job.to_json) } 123 | 124 | fetch.bulk_requeue(nil, nil) 125 | assert_equal 2, Sidekiq::PriorityQueue::Queue.new('foo').size 126 | end 127 | 128 | it "doesn't retrieve orphan jobs if checked recently" do 129 | # Setup lost job 130 | previous_process_identity = "sidekiq-pipeline-32152-xfwvw:42251:#{SecureRandom.hex(6)}" 131 | priority_queue = 'priority-queue:bar' 132 | previous_wip_queue = "queue:spriorityq|#{previous_process_identity}|#{priority_queue}" 133 | 134 | Sidekiq.redis { |conn| conn.sadd(previous_wip_queue, job.to_json) } 135 | # Setup recent orphan check 136 | Sidekiq.redis { |conn| conn.set('priority_reliable_fetch_orphan_check', Time.now.to_f, ex: 3600, nx: true) } 137 | 138 | # Then we setup the fetcher 139 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['bar'], index: 0) 140 | fetch.setup 141 | 142 | # And here we test that orphan check didn't run 143 | SidekiqUtilInstance.new.fire_event(:startup) 144 | assert_equal 0, Sidekiq::PriorityQueue::Queue.new('bar').size 145 | assert_equal 1, Sidekiq.redis { |conn| conn.scard(previous_wip_queue) } 146 | end 147 | 148 | it 'retrieves orphan jobs and pushes them back to priority queues' do 149 | # First we have to mimic old jobs lying around 150 | previous_process_identity = "sidekiq-pipeline-32152-xfwvw:42251:#{SecureRandom.hex(6)}" 151 | priority_queue = 'priority-queue:bar' 152 | previous_wip_queue = "queue:spriorityq|#{previous_process_identity}|#{priority_queue}" 153 | 154 | Sidekiq.redis { |conn| conn.sadd(previous_wip_queue, job.to_json) } 155 | 156 | # Then we setup the fetcher 157 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['bar'], index: 0) 158 | fetch.setup 159 | 160 | # And here we test is really on startup it brings back old jobs 161 | assert_equal 0, Sidekiq::PriorityQueue::Queue.new('bar').size 162 | SidekiqUtilInstance.new.fire_event(:startup) 163 | assert_equal 1, Sidekiq::PriorityQueue::Queue.new('bar').size 164 | refute Sidekiq.redis { |c| c.exists?(previous_wip_queue) } 165 | end 166 | 167 | it 'works even if orphan check raises an error' do 168 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: ['bar'], index: 0) 169 | fetch.setup 170 | 171 | def fetch.orphan_check?; raise StandardError; end 172 | SidekiqUtilInstance.new.fire_event(:startup, reraise: true) 173 | # If it passes, it didn't rise which is correct 174 | end 175 | 176 | it 'retrieves with strict setting' do 177 | fetch = Sidekiq::PriorityQueue::ReliableFetch.new(queues: %w[basic bar bar], strict: true) 178 | cmd = fetch.queues_cmd 179 | assert_equal cmd, ['priority-queue:basic', 'priority-queue:bar'] 180 | end 181 | end 182 | end 183 | -------------------------------------------------------------------------------- /lib/sidekiq/priority_queue/reliable_fetch.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require 'sidekiq' 4 | require 'sidekiq/util' 5 | 6 | module Sidekiq 7 | module PriorityQueue 8 | class ReliableFetch 9 | include Sidekiq::Util 10 | 11 | SUPER_PROCESSES_REGISTRY_KEY = 'super_processes_priority' 12 | 13 | UnitOfWork = Struct.new(:queue, :job, :wip_queue) do 14 | def acknowledge 15 | Sidekiq.redis do |conn| 16 | conn.srem(wip_queue, job) 17 | unless subqueue.nil? 18 | count = conn.zincrby(subqueue_counts, -1, subqueue) 19 | conn.zrem(subqueue_counts, subqueue) if count < 1 20 | end 21 | end 22 | end 23 | 24 | def queue_name 25 | queue.sub(/.*queue:/, '') 26 | end 27 | 28 | def subqueue 29 | @parsed_job ||= JSON.parse(job) 30 | @parsed_job['subqueue'] 31 | end 32 | 33 | def subqueue_counts 34 | "priority-queue-counts:#{queue_name}" 35 | end 36 | 37 | def requeue 38 | # Nothing needed. Jobs are in WIP queue. 39 | end 40 | end 41 | 42 | def initialize(options) 43 | @options = options 44 | @strictly_ordered_queues = !!options[:strict] 45 | @queues = options[:queues].map { |q| "priority-queue:#{q}" } 46 | @queues = @queues.uniq if @strictly_ordered_queues 47 | @done = false 48 | @process_index = options[:index] || ENV['PROCESS_INDEX'] 49 | end 50 | 51 | def setup 52 | Sidekiq.on(:startup) do 53 | cleanup_the_dead 54 | register_myself 55 | check 56 | end 57 | Sidekiq.on(:shutdown) do 58 | @done = true 59 | end 60 | Sidekiq.on(:heartbeat) do 61 | register_myself 62 | end 63 | end 64 | 65 | def retrieve_work 66 | return nil if @done 67 | 68 | work = @queues.detect do |q| 69 | job = zpopmin_sadd(q, wip_queue(q)) 70 | break [q, job] if job 71 | end 72 | UnitOfWork.new(*work, wip_queue(work.first)) if work 73 | end 74 | 75 | def wip_queue(q) 76 | "queue:spriorityq|#{identity}|#{q}" 77 | end 78 | 79 | def zpopmin_sadd(queue, wip_queue) 80 | Sidekiq.redis do |conn| 81 | @script_sha ||= conn.script(:load, Sidekiq::PriorityQueue::Scripts::ZPOPMIN_SADD) 82 | conn.evalsha(@script_sha, [queue, wip_queue]) 83 | end 84 | end 85 | 86 | def spop(wip_queue) 87 | Sidekiq.redis { |con| con.spop(wip_queue) } 88 | end 89 | 90 | def queues_cmd 91 | if @strictly_ordered_queues 92 | @queues 93 | else 94 | @queues.shuffle.uniq 95 | end 96 | end 97 | 98 | # Below method is called when we close sidekiq process gracefully 99 | def bulk_requeue(_inprogress, _options) 100 | Sidekiq.logger.debug { 'Priority ReliableFetch: Re-queueing terminated jobs' } 101 | requeue_wip_jobs 102 | unregister_super_process 103 | end 104 | 105 | private 106 | 107 | def check 108 | check_for_orphans if orphan_check? 109 | rescue StandardError => e 110 | # orphan check is best effort, we don't want Redis downtime to 111 | # break Sidekiq 112 | Sidekiq.logger.warn { "Priority ReliableFetch: Failed to do orphan check: #{e.message}" } 113 | end 114 | 115 | def orphan_check? 116 | delay = @options.fetch(:reliable_fetch_orphan_check, 3600).to_i 117 | return false if delay.zero? 118 | 119 | Sidekiq.redis do |conn| 120 | conn.set('priority_reliable_fetch_orphan_check', Time.now.to_f, ex: delay, nx: true) 121 | end 122 | end 123 | 124 | # This method is extra paranoid verification to check Redis for any possible 125 | # orphaned queues with jobs. If we change queue names and lose jobs in the meantime, 126 | # this will find old queues with jobs and rescue them. 127 | def check_for_orphans 128 | orphans_count = 0 129 | queues_count = 0 130 | orphan_queues = Set.new 131 | Sidekiq.redis do |conn| 132 | ids = conn.smembers(SUPER_PROCESSES_REGISTRY_KEY) 133 | Sidekiq.logger.debug("Priority ReliableFetch found #{ids.size} super processes") 134 | 135 | conn.scan_each(match: 'queue:spriorityq|*', count: 100) do |wip_queue| 136 | queues_count += 1 137 | _, id, original_priority_queue_name = wip_queue.split('|') 138 | next if ids.include?(id) 139 | 140 | # Race condition in pulling super_processes and checking queue liveness. 141 | # Need to verify in Redis. 142 | unless conn.sismember(SUPER_PROCESSES_REGISTRY_KEY, id) 143 | orphan_queues << original_priority_queue_name 144 | queue_jobs_count = 0 145 | loop do 146 | break if conn.scard(wip_queue).zero? 147 | 148 | # Here we should wrap below two operations in Lua script 149 | item = conn.spop(wip_queue) 150 | conn.zadd(original_priority_queue_name, 0, item) 151 | orphans_count += 1 152 | queue_jobs_count += 1 153 | end 154 | if queue_jobs_count.positive? 155 | Sidekiq::Pro.metrics.increment('jobs.recovered.fetch', by: queue_jobs_count, tags: ["queue:#{original_priority_queue_name}"]) 156 | end 157 | end 158 | end 159 | end 160 | 161 | if orphans_count.positive? 162 | Sidekiq.logger.warn { "Priority ReliableFetch recovered #{orphans_count} orphaned jobs in queues: #{orphan_queues.to_a.inspect}" } 163 | elsif queues_count.positive? 164 | Sidekiq.logger.info { "Priority ReliableFetch found #{queues_count} working queues with no orphaned jobs" } 165 | end 166 | orphans_count 167 | end 168 | 169 | # Below method is only to make sure we get jobs from incorrectly closed process (for example force killed using kill -9 SIDEKIQ_PID) 170 | def cleanup_the_dead 171 | overall_moved_count = 0 172 | Sidekiq.redis do |conn| 173 | conn.sscan_each(SUPER_PROCESSES_REGISTRY_KEY) do |super_process| 174 | next if conn.exists?(super_process) # Don't clean up currently running processes 175 | 176 | Sidekiq.logger.debug { "Priority ReliableFetch: Moving job from #{super_process} back to original queues" } 177 | 178 | # We need to pushback any leftover jobs still in WIP 179 | previously_handled_queues = conn.smembers("#{super_process}:super_priority_queues") 180 | 181 | # Below previously_handled_queues are simply WIP queues of previous, dead processes 182 | previously_handled_queues.each do |previously_handled_queue| 183 | queue_moved_size = 0 184 | original_priority_queue_name = previously_handled_queue.split('|').last 185 | 186 | Sidekiq.logger.debug { "Priority ReliableFetch: Moving job from #{previously_handled_queue} back to original queue: #{original_priority_queue_name}" } 187 | loop do 188 | break if conn.scard(previously_handled_queue).zero? 189 | 190 | # Here we should wrap below two operations in Lua script 191 | item = conn.spop(previously_handled_queue) 192 | conn.zadd(original_priority_queue_name, 0, item) 193 | queue_moved_size += 1 194 | overall_moved_count += 1 195 | end 196 | # Below we simply remove old WIP queue 197 | conn.del(previously_handled_queue) if conn.scard(previously_handled_queue).zero? 198 | Sidekiq.logger.debug { "Priority ReliableFetch: Moved #{queue_moved_size} jobs from ##{previously_handled_queue} back to original_queue: #{original_priority_queue_name} " } 199 | end 200 | 201 | Sidekiq.logger.debug { "Priority ReliableFetch: Unregistering super process #{super_process}" } 202 | conn.del("#{super_process}:super_priority_queues") 203 | conn.srem(SUPER_PROCESSES_REGISTRY_KEY, super_process) 204 | end 205 | end 206 | Sidekiq.logger.debug { "Priority ReliableFetch: Moved overall #{overall_moved_count} jobs from WIP queues" } 207 | rescue StandardError => e 208 | # best effort, ignore Redis network errors 209 | Sidekiq.logger.warn { "Priority ReliableFetch: Failed to requeue: #{e.message}" } 210 | end 211 | 212 | def requeue_wip_jobs 213 | jobs_to_requeue = {} 214 | Sidekiq.redis do |conn| 215 | @queues.each do |q| 216 | wip_queue_name = wip_queue(q) 217 | jobs_to_requeue[q] = [] 218 | 219 | while job = conn.spop(wip_queue_name) 220 | jobs_to_requeue[q] << job 221 | end 222 | end 223 | 224 | conn.pipelined do 225 | jobs_to_requeue.each do |queue, jobs| 226 | next if jobs.empty? # ZADD doesn't work with empty arrays 227 | 228 | conn.zadd(queue, jobs.map { |j| [0, j] }) 229 | end 230 | end 231 | end 232 | Sidekiq.logger.info("Priority ReliableFetch: Pushed #{jobs_to_requeue.values.flatten.size} jobs back to Redis") 233 | rescue StandardError => e 234 | Sidekiq.logger.warn("Priority ReliableFetch: Failed to requeue #{jobs_to_requeue.values.flatten.size} jobs: #{e.message}") 235 | end 236 | 237 | def register_myself 238 | super_process_wip_queues = @queues.map { |q| wip_queue(q) } 239 | id = identity # This is from standard sidekiq, updated with every heartbeat 240 | 241 | # This method will run multiple times so seeing this message twice is no problem. 242 | Sidekiq.logger.debug { "Priority ReliableFetch: Registering super process #{id} with #{super_process_wip_queues}" } 243 | 244 | Sidekiq.redis do |conn| 245 | conn.multi do 246 | conn.sadd(SUPER_PROCESSES_REGISTRY_KEY, id) 247 | conn.sadd("#{id}:super_priority_queues", super_process_wip_queues) 248 | end 249 | end 250 | end 251 | 252 | def unregister_super_process 253 | id = identity 254 | Sidekiq.logger.debug { "Priority ReliableFetch: Unregistering super process #{id}" } 255 | Sidekiq.redis do |conn| 256 | conn.multi do 257 | conn.srem(SUPER_PROCESSES_REGISTRY_KEY, id) 258 | conn.del("#{id}:super_priority_queues") 259 | end 260 | end 261 | end 262 | end 263 | end 264 | end 265 | --------------------------------------------------------------------------------