sidekiq-batch 0.1.5 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA1:
-  metadata.gz: 2d61afa130ecc7d339ec1ff9c2b38ce8870ef0e5
-  data.tar.gz: a28ac95c3bb93ebd274f99b253e8f7db7d01ffeb
+SHA256:
+  metadata.gz: '0811ae31d4850eee031831062fbe1aca229d9255f28ee85641fff02f22e4ff49'
+  data.tar.gz: d37febab0ff1087d3b3966f0beae1155246b3535f703e3bd62e9ae04fd83a553
 SHA512:
-  metadata.gz: 3ad9d8f6485798ebe314bd561f040b44742bf892d89d044793c77b3eadb9b548c81dfa5884ca35b89eab044159354d5a46fc9ff99864e0df1b9660c69b16d6dc
-  data.tar.gz: dc66cae976602555b91ce985d6bdc7bf1c5fc3fb9491cc5f14a2debded4799951b1c7eafe7af9940b9986e25d76192969b7140ba58dd647b28a46482627e32f4
+  metadata.gz: dafce26855e716909db76ed31c04f31381d2ca4ab9d9d22b4d85e95b4deda62cf5438d562cce5344596172bbc912dc5995a5b5c435cceadca78004bf8db0a7b3
+  data.tar.gz: e25c54027984ff63d2242037de20795c8e223c4b295504678aa5a4a5f339568f44e290458f1a732d25f4c692dd6b3961256525bed66c19e20b95e0de936bf3be

data/lib/sidekiq/batch.rb

@@ -33,6 +33,11 @@ module Sidekiq
       persist_bid_attr('callback_queue', callback_queue)
     end
 
+    def callback_batch=(callback_batch)
+      @callback_batch = callback_batch
+      persist_bid_attr('callback_batch', callback_batch)
+    end
+
     def on(event, callback, options = {})
       return unless %w(success complete).include?(event.to_s)
       callback_key = "#{@bidkey}-callbacks-#{event}"
@@ -54,7 +59,7 @@ module Sidekiq
 
       begin
         if !@existing && !@initialized
-          parent_bid = Thread.current[:bid].bid if Thread.current[:bid]
+          parent_bid = Thread.current[:batch].bid if Thread.current[:batch]
 
           Sidekiq.redis do |r|
             r.multi do
@@ -70,11 +75,11 @@ module Sidekiq
         @ready_to_queue = []
 
         begin
-          parent = Thread.current[:bid]
-          Thread.current[:bid] = self
+          parent = Thread.current[:batch]
+          Thread.current[:batch] = self
           yield
         ensure
-          Thread.current[:bid] = parent
+          Thread.current[:batch] = parent
         end
 
         return [] if @ready_to_queue.size == 0
@@ -83,6 +88,7 @@ module Sidekiq
           r.multi do
             if parent_bid
               r.hincrby("BID-#{parent_bid}", "children", 1)
+              r.hincrby("BID-#{parent_bid}", "total", @ready_to_queue.size)
               r.expire("BID-#{parent_bid}", BID_EXPIRE_TTL)
             end
 
@@ -141,7 +147,7 @@ module Sidekiq
 
     class << self
       def process_failed_job(bid, jid)
-        _, pending, failed, children, complete = Sidekiq.redis do |r|
+        _, pending, failed, children, complete, parent_bid = Sidekiq.redis do |r|
           r.multi do
             r.sadd("BID-#{bid}-failed", jid)
 
@@ -149,12 +155,26 @@ module Sidekiq
             r.scard("BID-#{bid}-failed")
             r.hincrby("BID-#{bid}", "children", 0)
             r.scard("BID-#{bid}-complete")
+            r.hget("BID-#{bid}", "parent_bid")
 
             r.expire("BID-#{bid}-failed", BID_EXPIRE_TTL)
           end
         end
 
-        enqueue_callbacks(:complete, bid) if pending.to_i == failed.to_i && children == complete
+        # if the batch failed, and has a parent, update the parent to show one pending and failed job
+        if parent_bid
+          Sidekiq.redis do |r|
+            r.multi do
+              r.hincrby("BID-#{parent_bid}", "pending", 1)
+              r.sadd("BID-#{parent_bid}-failed", jid)
+              r.expire("BID-#{parent_bid}-failed", BID_EXPIRE_TTL)
+            end
+          end
+        end
+
+        if pending.to_i == failed.to_i && children == complete
+          enqueue_callbacks(:complete, bid)
+        end
       end
 
       def process_successful_job(bid, jid)
@@ -174,49 +194,96 @@ module Sidekiq
           end
         end
 
-        Sidekiq.logger.info "done: #{jid} in batch #{bid}"
-
-        enqueue_callbacks(:complete, bid) if pending.to_i == failed.to_i && children == complete
-        enqueue_callbacks(:success, bid) if pending.to_i.zero? && children == success
+        # if complete or successfull call complete callback (the complete callback may then call successful)
+        if (pending.to_i == failed.to_i && children == complete) || (pending.to_i.zero? && children == success)
+          enqueue_callbacks(:complete, bid)
+        end
       end
 
       def enqueue_callbacks(event, bid)
         batch_key = "BID-#{bid}"
         callback_key = "#{batch_key}-callbacks-#{event}"
-        needed, _, callbacks, queue, parent_bid = Sidekiq.redis do |r|
+        already_processed, _, callbacks, queue, parent_bid, callback_batch = Sidekiq.redis do |r|
           r.multi do
             r.hget(batch_key, event)
             r.hset(batch_key, event, true)
             r.smembers(callback_key)
             r.hget(batch_key, "callback_queue")
             r.hget(batch_key, "parent_bid")
+            r.hget(batch_key, "callback_batch")
           end
         end
-        return if needed == 'true'
 
-        begin
-          parent_bid = !parent_bid || parent_bid.empty? ? nil : parent_bid    # Basically parent_bid.blank?
-          Sidekiq::Client.push_bulk(
-            'class' => Sidekiq::Batch::Callback::Worker,
-            'args' => callbacks.reduce([]) do |memo, jcb|
-              cb = Sidekiq.load_json(jcb)
-              memo << [cb['callback'], event, cb['opts'], bid, parent_bid]
-            end,
-            'queue' => queue ||= 'default'
-          ) unless callbacks.empty?
-        ensure
-          cleanup_redis(bid) if event == :success
+        return if already_processed == 'true'
+
+        queue ||= "default"
+        parent_bid = !parent_bid || parent_bid.empty? ? nil : parent_bid    # Basically parent_bid.blank?
+        callback_args = callbacks.reduce([]) do |memo, jcb|
+          cb = Sidekiq.load_json(jcb)
+          memo << [cb['callback'], event, cb['opts'], bid, parent_bid]
+        end
+
+        opts = {"bid" => bid, "event" => event}
+
+        # Run callback batch finalize synchronously
+        if callback_batch
+          # Extract opts from cb_args or use current
+          # Pass in stored event as callback finalize is processed on complete event
+          cb_opts = callback_args.first&.at(2) || opts
+
+          Sidekiq.logger.debug {"Run callback batch bid: #{bid} event: #{event} args: #{callback_args.inspect}"}
+          # Finalize now
+          finalizer = Sidekiq::Batch::Callback::Finalize.new
+          status = Status.new bid
+          finalizer.dispatch(status, cb_opts)
+
+          return
+        end
+
+        Sidekiq.logger.debug {"Enqueue callback bid: #{bid} event: #{event} args: #{callback_args.inspect}"}
+
+        if callback_args.empty?
+          # Finalize now
+          finalizer = Sidekiq::Batch::Callback::Finalize.new
+          status = Status.new bid
+          finalizer.dispatch(status, opts)
+        else
+          # Otherwise finalize in sub batch complete callback
+          cb_batch = self.new
+          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
+            push_callbacks callback_args, queue
+          end
         end
       end
 
       def cleanup_redis(bid)
+        Sidekiq.logger.debug {"Cleaning redis of batch #{bid}"}
         Sidekiq.redis do |r|
-          r.del("BID-#{bid}",
-                "BID-#{bid}-callbacks-complete",
-                "BID-#{bid}-callbacks-success",
-                "BID-#{bid}-failed")
+          r.del(
+            "BID-#{bid}",
+            "BID-#{bid}-callbacks-complete",
+            "BID-#{bid}-callbacks-success",
+            "BID-#{bid}-failed",
+
+            "BID-#{bid}-success",
+            "BID-#{bid}-complete",
+            "BID-#{bid}-jids",
+          )
         end
       end
+
+    private
+
+      def push_callbacks args, queue
+        Sidekiq::Client.push_bulk(
+          'class' => Sidekiq::Batch::Callback::Worker,
+          'args' => args,
+          'queue' => queue
+        ) unless args.empty?
+      end
     end
   end
 end

data/lib/sidekiq/batch/callback.rb

@@ -6,66 +6,98 @@ module Sidekiq
 
         def perform(clazz, event, opts, bid, parent_bid)
           return unless %w(success complete).include?(event)
-          clazz, method = clazz.split("#") if (clazz.class == String && clazz.include?("#"))
+          clazz, method = clazz.split("#") if (clazz && clazz.class == String && clazz.include?("#"))
           method = "on_#{event}" if method.nil?
           status = Sidekiq::Batch::Status.new(bid)
 
-          if object = Object.const_get(clazz)
+          if clazz && object = Object.const_get(clazz)
             instance = object.new
             instance.send(method, status, opts) if instance.respond_to?(method)
           end
-
-          send(event.to_sym, bid, status, parent_bid)
         end
+      end
+
+      class Finalize
+        def dispatch status, opts
+          bid = opts["bid"]
+          callback_bid = status.bid
+          event = opts["event"].to_sym
+          callback_batch = bid != callback_bid
+
+          Sidekiq.logger.debug {"Finalize #{event} batch id: #{opts["bid"]}, callback batch id: #{callback_bid} callback_batch #{callback_batch}"}
+
+          batch_status = Status.new bid
+          send(event, bid, batch_status, batch_status.parent_bid)
 
 
+          # Different events are run in different callback batches
+          Sidekiq::Batch.cleanup_redis callback_bid if callback_batch
+          Sidekiq::Batch.cleanup_redis bid if event == :success
+        end
+
         def success(bid, status, parent_bid)
-          if (parent_bid)
-            _, _, success, pending, children = Sidekiq.redis do |r|
-              r.multi do
-                r.sadd("BID-#{parent_bid}-success", bid)
-                r.expire("BID-#{parent_bid}-success", Sidekiq::Batch::BID_EXPIRE_TTL)
-                r.scard("BID-#{parent_bid}-success")
-                r.hincrby("BID-#{parent_bid}", "pending", 0)
-                r.hincrby("BID-#{parent_bid}", "children", 0)
-              end
-            end
+          return unless parent_bid
 
-            Batch.enqueue_callbacks(:success, parent_bid) if pending.to_i.zero? && children == success
+          _, _, success, _, complete, pending, children, failure = Sidekiq.redis do |r|
+            r.multi do
+              r.sadd("BID-#{parent_bid}-success", bid)
+              r.expire("BID-#{parent_bid}-success", Sidekiq::Batch::BID_EXPIRE_TTL)
+              r.scard("BID-#{parent_bid}-success")
+              r.sadd("BID-#{parent_bid}-complete", bid)
+              r.scard("BID-#{parent_bid}-complete")
+              r.hincrby("BID-#{parent_bid}", "pending", 0)
+              r.hincrby("BID-#{parent_bid}", "children", 0)
+              r.scard("BID-#{parent_bid}-failed")
+            end
           end
-
-          Sidekiq.redis do |r|
-            r.del "BID-#{bid}-success", "BID-#{bid}-complete", "BID-#{bid}-jids", "BID-#{bid}-failed"
+          # if job finished successfully and parent batch completed call parent complete callback
+          # Success callback is called after complete callback
+          if complete == children && pending == failure
+            Sidekiq.logger.debug {"Finalize parent complete bid: #{parent_bid}"}
+            Batch.enqueue_callbacks(:complete, parent_bid)
           end
+
         end
 
         def complete(bid, status, parent_bid)
-          if (parent_bid)
+          pending, children, success = Sidekiq.redis do |r|
+            r.multi do
+              r.hincrby("BID-#{bid}", "pending", 0)
+              r.hincrby("BID-#{bid}", "children", 0)
+              r.scard("BID-#{bid}-success")
+            end
+          end
+
+          # if we batch was successful run success callback
+          if pending.to_i.zero? && children == success
+            Batch.enqueue_callbacks(:success, bid)
+
+          elsif parent_bid
+            # if batch was not successfull check and see if its parent is complete
+            # if the parent is complete we trigger the complete callback
+            # We don't want to run this if the batch was successfull because the success
+            # callback may add more jobs to the parent batch
+
+            Sidekiq.logger.debug {"Finalize parent complete bid: #{parent_bid}"}
             _, complete, pending, children, failure = Sidekiq.redis do |r|
               r.multi do
                 r.sadd("BID-#{parent_bid}-complete", bid)
                 r.scard("BID-#{parent_bid}-complete")
                 r.hincrby("BID-#{parent_bid}", "pending", 0)
                 r.hincrby("BID-#{parent_bid}", "children", 0)
-                r.hlen("BID-#{parent_bid}-failed")
+                r.scard("BID-#{parent_bid}-failed")
               end
             end
-
-            Batch.enqueue_callbacks(:complete, parent_bid) if complete == children && pending == failure
-          end
-
-          pending, children, success = Sidekiq.redis do |r|
-            r.multi do
-              r.hincrby("BID-#{bid}", "pending", 0)
-              r.hincrby("BID-#{bid}", "children", 0)
-              r.scard("BID-#{bid}-success")
+            if complete == children && pending == failure
+              Batch.enqueue_callbacks(:complete, parent_bid)
             end
           end
-
-          Batch.enqueue_callbacks(:success, bid) if pending.to_i.zero? && children == success
-
         end
 
+        def cleanup_redis bid, callback_bid=nil
+          Sidekiq::Batch.cleanup_redis bid
+          Sidekiq::Batch.cleanup_redis callback_bid if callback_bid
+        end
       end
     end
   end

data/lib/sidekiq/batch/extension/worker.rb

@@ -1,11 +1,11 @@
 module Sidekiq::Batch::Extension
   module Worker
     def bid
-      Thread.current[:bid]
+      Thread.current[:batch].bid
     end
 
     def batch
-      Sidekiq::Batch.new(Thread.current[:bid].bid) if Thread.current[:bid]
+      Thread.current[:batch]
     end
 
     def valid_within_batch?

data/lib/sidekiq/batch/middleware.rb

@@ -5,7 +5,7 @@ module Sidekiq
     module Middleware
       class ClientMiddleware
         def call(_worker, msg, _queue, _redis_pool = nil)
-          if (batch = Thread.current[:bid])
+          if (batch = Thread.current[:batch])
             batch.increment_job_queue(msg['jid']) if (msg[:bid] = batch.bid)
           end
           yield
@@ -16,15 +16,15 @@ module Sidekiq
         def call(_worker, msg, _queue)
           if (bid = msg['bid'])
             begin
-              Thread.current[:bid] = Sidekiq::Batch.new(bid)
+              Thread.current[:batch] = Sidekiq::Batch.new(bid)
               yield
-              Thread.current[:bid] = nil
+              Thread.current[:batch] = nil
               Batch.process_successful_job(bid, msg['jid'])
             rescue
               Batch.process_failed_job(bid, msg['jid'])
               raise
             ensure
-              Thread.current[:bid] = nil
+              Thread.current[:batch] = nil
             end
           else
             yield

data/lib/sidekiq/batch/status.rb

@@ -45,13 +45,15 @@ module Sidekiq
 
       def data
         {
+          bid: bid,
           total: total,
           failures: failures,
           pending: pending,
           created_at: created_at,
           complete: complete?,
           failure_info: failure_info,
-          parent_bid: parent_bid
+          parent_bid: parent_bid,
+          child_count: child_count
         }
       end
     end

data/lib/sidekiq/batch/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: sidekiq-batch
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.1.6
 platform: ruby
 authors:
 - Marcin Naglik
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2018-10-16 00:00:00.000000000 Z
+date: 2019-08-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: sidekiq
@@ -121,8 +121,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubyforge_project: 
-rubygems_version: 2.5.2.3
+rubygems_version: 3.0.3
 signing_key: 
 specification_version: 4
 summary: Sidekiq Batch Jobs