Merge pull request #1 from jaceksamol/kb-update-fork

Update fork

Author: jaceksamol

.github/FUNDING.yml

@@ -0,0 +1,3 @@
+# These are supported funding model platforms
+github: breamware
+open_collective: sidekiq-batch

.github/workflows/ci.yml

@@ -1,16 +1,30 @@
 name: CI
 
-on: [push, pull_request]
+on: [pull_request, workflow_dispatch]
 
 jobs:
   test:
+    env:
+      REDIS_HOST: 'redis'
 
     runs-on: ubuntu-latest
+    services:
+      redis:
+        image: redis
+        # Set health checks to wait until redis has started
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          # Maps port 6379 on service container to the host
+          - 6379:6379
 
     strategy:
       fail-fast: false
       matrix:
-        ruby: ["2.5", "2.6", "2.7", "3.0", "3.1", ruby-head]
+        ruby: ["3.0", "3.1", "3.2", ruby-head]
 
     steps:
     - uses: actions/checkout@v2

README.md

@@ -32,7 +32,41 @@ Or install it yourself as:
 
 ## Usage
 
-Sidekiq Batch is drop-in replacement for the API from Sidekiq PRO. See https://github.com/mperham/sidekiq/wiki/Batches for usage.
+Sidekiq Batch is MOSTLY a drop-in replacement for the API from Sidekiq PRO. See https://github.com/mperham/sidekiq/wiki/Batches for usage.
+
+## Caveats/Gotchas
+
+Consider the following workflow:
+
+  * Batch Z created
+  * Worker A queued in batch Z
+  * Worker A starts Worker B in batch Z
+  * Worker B completes *before* worker A does
+  * Worker A completes
+
+In the standard configuration, the `on(:success)` and `on(:complete)` callbacks will be triggered when Worker B completes.
+This configuration is the default, simply for legacy reasons. This gem adds the following option to the sidekiq.yml options:
+
+```yaml
+:batch_push_interval: 0
+```
+
+When this value is *absent* (aka legacy), Worker A will only push the increment of batch jobs (aka Worker B) *when it completes*
+
+When this value is set to `0`, Worker A will increment the count as soon as `WorkerB.perform_async` is called
+
+When this value is a positive number, Worker A will wait a maximum of value-seconds before pushing the increment to redis, or until it's done, whichever comes first.
+
+This comes into play if Worker A is queueing thousands of WorkerB jobs, or has some other reason for WorkerB to complete beforehand.
+
+If you are queueing many WorkerB jobs, it is recommended to set this value to something like `3` to avoid thousands of calls to redis, and call WorkerB like so:
+```ruby
+WorkerB.perform_in(4.seconds, some, args)
+```
+this will ensure that the batch callback does not get triggered until WorkerA *and* the last WorkerB job complete.
+
+If WorkerA is just slow for whatever reason, setting to `0` will update the batch status immediately so that the callbacks don't fire.
+
 
 ## Contributing
 

lib/sidekiq/batch.rb

@@ -20,7 +20,11 @@ def initialize(existing_bid = nil)
       @initialized = false
       @created_at = Time.now.utc.to_f
       @bidkey = "BID-" + @bid.to_s
-      @ready_to_queue = []
+      @queued_jids = []
+      @pending_jids = []
+
+      @incremental_push = !Sidekiq.default_configuration[:batch_push_interval].nil?
+      @batch_push_interval = Sidekiq.default_configuration[:batch_push_interval]
     end
 
     def description=(description)
@@ -43,10 +47,10 @@ def on(event, callback, options = {})
       callback_key = "#{@bidkey}-callbacks-#{event}"
       Sidekiq.redis do |r|
         r.multi do |pipeline|
-          pipeline.sadd(callback_key, JSON.unparse({
+          pipeline.sadd(callback_key, [JSON.unparse({
             callback: callback,
             opts: options
-          }))
+          })])
           pipeline.expire(callback_key, BID_EXPIRE_TTL)
         end
       end
@@ -64,51 +68,86 @@ def jobs
           Sidekiq.redis do |r|
             r.multi do |pipeline|
               pipeline.hset(@bidkey, "created_at", @created_at)
-              pipeline.hset(@bidkey, "parent_bid", parent_bid.to_s) if parent_bid
               pipeline.expire(@bidkey, BID_EXPIRE_TTL)
+              if parent_bid
+                pipeline.hset(@bidkey, "parent_bid", parent_bid.to_s)
+                pipeline.hincrby("BID-#{parent_bid}", "children", 1)
+              end
             end
           end
 
           @initialized = true
         end
 
-        @ready_to_queue = []
+        @queued_jids = []
+        @pending_jids = []
 
         begin
           parent = Thread.current[:batch]
           Thread.current[:batch] = self
+          Thread.current[:parent_bid] = parent_bid
           yield
         ensure
           Thread.current[:batch] = parent
+          Thread.current[:parent_bid] = nil
         end
 
-        return [] if @ready_to_queue.size == 0
+        return [] if @queued_jids.size == 0
+        conditional_redis_increment!(true)
 
         Sidekiq.redis do |r|
           r.multi do |pipeline|
             if parent_bid
-              pipeline.hincrby("BID-#{parent_bid}", "children", 1)
-              pipeline.hincrby("BID-#{parent_bid}", "total", @ready_to_queue.size)
               pipeline.expire("BID-#{parent_bid}", BID_EXPIRE_TTL)
             end
 
-            pipeline.hincrby(@bidkey, "pending", @ready_to_queue.size)
-            pipeline.hincrby(@bidkey, "total", @ready_to_queue.size)
             pipeline.expire(@bidkey, BID_EXPIRE_TTL)
 
-            pipeline.sadd(@bidkey + "-jids", @ready_to_queue)
+            pipeline.sadd(@bidkey + "-jids", @queued_jids)
             pipeline.expire(@bidkey + "-jids", BID_EXPIRE_TTL)
           end
         end
 
-        @ready_to_queue
+        @queued_jids
       ensure
         Thread.current[:bid_data] = bid_data
       end
     end
 
     def increment_job_queue(jid)
-      @ready_to_queue << jid
+      @queued_jids << jid
+      @pending_jids << jid
+      conditional_redis_increment!
+    end
+
+    def conditional_redis_increment!(force=false)
+      if should_increment? || force
+        parent_bid = Thread.current[:parent_bid]
+        Sidekiq.redis do |r|
+          r.multi do |pipeline|
+            if parent_bid
+              pipeline.hincrby("BID-#{parent_bid}", "total", @pending_jids.length)
+              pipeline.expire("BID-#{parent_bid}", BID_EXPIRE_TTL)
+            end
+
+            pipeline.hincrby(@bidkey, "pending", @pending_jids.length)
+            pipeline.hincrby(@bidkey, "total", @pending_jids.length)
+            pipeline.expire(@bidkey, BID_EXPIRE_TTL)
+          end
+        end
+        @pending_jids = []
+      end
+    end
+
+    def should_increment?
+      return false unless @incremental_push
+      return true if @batch_push_interval == 0 || @queued_jids.length == 1
+      now = Time.now.to_f
+      @last_increment ||= now
+      if @last_increment + @batch_push_interval > now
+        @last_increment = now
+        return true
+      end
     end
 
     def invalidate_all
@@ -130,7 +169,7 @@ def parent
     end
 
     def valid?(batch = self)
-      valid = !Sidekiq.redis { |r| r.exists("invalidated-bid-#{batch.bid}") }
+      valid = Sidekiq.redis { |r| r.exists("invalidated-bid-#{batch.bid}") }.zero?
       batch.parent ? valid && valid?(batch.parent) : valid
     end
 
@@ -149,7 +188,7 @@ class << self
       def process_failed_job(bid, jid)
         _, pending, failed, children, complete, parent_bid = Sidekiq.redis do |r|
           r.multi do |pipeline|
-            pipeline.sadd("BID-#{bid}-failed", jid)
+            pipeline.sadd("BID-#{bid}-failed", [jid])
 
             pipeline.hincrby("BID-#{bid}", "pending", 0)
             pipeline.scard("BID-#{bid}-failed")
@@ -166,7 +205,7 @@ def process_failed_job(bid, jid)
           Sidekiq.redis do |r|
             r.multi do |pipeline|
               pipeline.hincrby("BID-#{parent_bid}", "pending", 1)
-              pipeline.sadd("BID-#{parent_bid}-failed", jid)
+              pipeline.sadd("BID-#{parent_bid}-failed", [jid])
               pipeline.expire("BID-#{parent_bid}-failed", BID_EXPIRE_TTL)
             end
           end
@@ -188,8 +227,8 @@ def process_successful_job(bid, jid)
             pipeline.hget("BID-#{bid}", "total")
             pipeline.hget("BID-#{bid}", "parent_bid")
 
-            pipeline.srem("BID-#{bid}-failed", jid)
-            pipeline.srem("BID-#{bid}-jids", jid)
+            pipeline.srem("BID-#{bid}-failed", [jid])
+            pipeline.srem("BID-#{bid}-jids", [jid])
             pipeline.expire("BID-#{bid}", BID_EXPIRE_TTL)
           end
         end
@@ -209,7 +248,7 @@ def enqueue_callbacks(event, bid)
         already_processed, _, callbacks, queue, parent_bid, callback_batch = Sidekiq.redis do |r|
           r.multi do |pipeline|
             pipeline.hget(batch_key, event_name)
-            pipeline.hset(batch_key, event_name, true)
+            pipeline.hset(batch_key, event_name, 'true')
             pipeline.smembers(callback_key)
             pipeline.hget(batch_key, "callback_queue")
             pipeline.hget(batch_key, "parent_bid")
@@ -253,7 +292,7 @@ def enqueue_callbacks(event, bid)
         else
           # Otherwise finalize in sub batch complete callback
           cb_batch = self.new
-          cb_batch.callback_batch = true
+          cb_batch.callback_batch = 'true'
           Sidekiq.logger.debug {"Adding callback batch: #{cb_batch.bid} for batch: #{bid}"}
           cb_batch.on(:complete, "Sidekiq::Batch::Callback::Finalize#dispatch", opts)
           cb_batch.jobs do

lib/sidekiq/batch/callback.rb

@@ -40,10 +40,10 @@ def success(bid, status, parent_bid)
 
           _, _, success, _, complete, pending, children, failure = Sidekiq.redis do |r|
             r.multi do |pipeline|
-              pipeline.sadd("BID-#{parent_bid}-success", bid)
+              pipeline.sadd("BID-#{parent_bid}-success", [bid])
               pipeline.expire("BID-#{parent_bid}-success", Sidekiq::Batch::BID_EXPIRE_TTL)
               pipeline.scard("BID-#{parent_bid}-success")
-              pipeline.sadd("BID-#{parent_bid}-complete", bid)
+              pipeline.sadd("BID-#{parent_bid}-complete", [bid])
               pipeline.scard("BID-#{parent_bid}-complete")
               pipeline.hincrby("BID-#{parent_bid}", "pending", 0)
               pipeline.hincrby("BID-#{parent_bid}", "children", 0)
@@ -81,7 +81,7 @@ def complete(bid, status, parent_bid)
             Sidekiq.logger.debug {"Finalize parent complete bid: #{parent_bid}"}
             _, complete, pending, children, failure = Sidekiq.redis do |r|
               r.multi do |pipeline|
-                pipeline.sadd("BID-#{parent_bid}-complete", bid)
+                pipeline.sadd("BID-#{parent_bid}-complete", [bid])
                 pipeline.scard("BID-#{parent_bid}-complete")
                 pipeline.hincrby("BID-#{parent_bid}", "pending", 0)
                 pipeline.hincrby("BID-#{parent_bid}", "children", 0)

lib/sidekiq/batch/version.rb

@@ -1,5 +1,5 @@
 module Sidekiq
   class Batch
-    VERSION = '0.1.8'.freeze
+    VERSION = '0.1.9'.freeze
   end
 end

sidekiq-batch.gemspec

@@ -19,10 +19,9 @@ Gem::Specification.new do |spec|
   spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]
 
-  spec.add_dependency "sidekiq", ">= 3"
+  spec.add_dependency "sidekiq", ">= 7", "<8"
 
   spec.add_development_dependency "bundler", "~> 2.1"
   spec.add_development_dependency "rake", "~> 13.0"
   spec.add_development_dependency "rspec", "~> 3.0"
-  spec.add_development_dependency "fakeredis", "~> 0.8.0"
 end

spec/sidekiq/batch/middleware_spec.rb

@@ -79,7 +79,7 @@
 end
 
 describe Sidekiq::Batch::Middleware do
-  let(:config) { class_double(Sidekiq) }
+  let(:config) { instance_double(Sidekiq::Config) }
   let(:client_middleware) { double(Sidekiq::Middleware::Chain) }
 
   context 'client' do

spec/sidekiq/batch_spec.rb

@@ -170,7 +170,7 @@ def was_performed; end
         Sidekiq::Batch.process_failed_job(bid, 'failed-job-id')
         Sidekiq::Batch.process_failed_job(bid, failed_jid)
         failed = Sidekiq.redis { |r| r.smembers("BID-#{bid}-failed") }
-        expect(failed).to eq(['xxx', 'failed-job-id'])
+        expect(failed).to eq(['failed-job-id', 'xxx'])
       end
     end
   end
@@ -245,7 +245,7 @@ def was_performed; end
       it 'returns and does not enqueue callbacks' do
         batch = Sidekiq::Batch.new
         batch.on(event, SampleCallback)
-        Sidekiq.redis { |r| r.hset("BID-#{batch.bid}", event, true) }
+        Sidekiq.redis { |r| r.hset("BID-#{batch.bid}", event, 'true') }
 
         expect(Sidekiq::Client).not_to receive(:push)
         Sidekiq::Batch.enqueue_callbacks(event, batch.bid)
@@ -290,8 +290,8 @@ def was_performed; end
           expect(Sidekiq::Client).to receive(:push_bulk).with(
             'class' => Sidekiq::Batch::Callback::Worker,
             'args' => [
-              ['SampleCallback2', event.to_s, opts2, batch.bid, nil],
-              ['SampleCallback', event.to_s, opts, batch.bid, nil]
+              ['SampleCallback', event.to_s, opts, batch.bid, nil],
+              ['SampleCallback2', event.to_s, opts2, batch.bid, nil]
             ],
             'queue' => 'default'
           )

spec/spec_helper.rb

@@ -2,11 +2,9 @@
 SimpleCov.start
 
 $LOAD_PATH.unshift File.expand_path('../../lib', __FILE__)
-require 'fakeredis/rspec'
 require 'sidekiq/batch'
 
 redis_opts = { url: "redis://127.0.0.1:6379/1" }
-redis_opts[:driver] = Redis::Connection::Memory if defined?(Redis::Connection::Memory)
 
 Sidekiq.configure_client do |config|
   config.redis = redis_opts
@@ -19,6 +17,12 @@
 RSpec.configure do |config|
   config.filter_run focus: true
   config.run_all_when_everything_filtered = true
+
+  config.before do
+    Sidekiq.redis do |redis|
+      redis.flushdb
+    end
+  end
 end
 
 Dir[File.dirname(__FILE__) + "/support/**/*.rb"].each {|f| require f }