joblin 0.1.0

data/lib/joblin/batching/batch.rb

@@ -0,0 +1,537 @@
+
+require_relative "./callback"
+require_relative "./context_hash"
+require_relative "./status"
+require_relative "./pool"
+Dir[File.dirname(__FILE__) + "/jobs/*.rb"].each { |file| require file }
+require_relative "./chain_builder"
+
+# Implement Job Batching similar to Sidekiq::Batch. Supports ActiveJob and Sidekiq, or a mix thereof.
+# Much of this code is modifed/extended from https://github.com/breamware/sidekiq-batch
+
+module Joblin::Batching
+  CURRENT_BATCH_THREAD_KEY = :job_batches_batch
+
+  class Batch
+    include RediConn::RedisModel
+
+    class NoBlockGivenError < StandardError; end
+
+    delegate :redis, to: :class
+
+    BID_EXPIRE_TTL = 90.days.to_i
+    INDEX_ALL_BATCHES = false
+    SCHEDULE_CALLBACK = RediConn::RedisScript.new(Pathname.new(__FILE__) + "../schedule_callback.lua")
+    BID_HIERARCHY = RediConn::RedisScript.new(Pathname.new(__FILE__) + "../hier_batch_ids.lua")
+
+    attr_reader :bid
+
+    def self.current
+      Thread.current[CURRENT_BATCH_THREAD_KEY]
+    end
+
+    def self.current_context
+      self.current&.context
+    end
+
+    def initialize(existing_bid = nil)
+      @bid = existing_bid || SecureRandom.urlsafe_base64(10)
+      @existing = !(!existing_bid || existing_bid.empty?)  # Basically existing_bid.present?
+      @initialized = false
+      @bidkey = "BID-" + @bid.to_s
+      self.created_at = Time.now.utc.to_f unless @existing
+    end
+
+    redis_attr :description
+    redis_attr :created_at
+    redis_attr :callback_queue, read_only: false
+    redis_attr :callback_params, :json
+    redis_attr :allow_context_changes, :bool
+
+    def context
+      return @context if defined?(@context)
+
+      if (@initialized || @existing)
+        @context = ContextHash.new(bid)
+      else
+        @context = ContextHash.new(bid, {})
+      end
+    end
+
+    def context=(value)
+      raise "context is read-only once the batch has been started" if (@initialized || @existing) # && !allow_context_changes
+      raise "context must be a Hash" unless value.is_a?(Hash) || value.nil?
+      return nil if value.nil? && @context.nil?
+
+      value = {} if value.nil?
+      value = value.local if value.is_a?(ContextHash)
+
+      @context ||= ContextHash.new(bid, {})
+      @context.set_local(value)
+      # persist_bid_attr('context', JSON.unparse(@context.local))
+    end
+
+    def save_context_changes
+      @context&.save!
+    end
+
+    # Events:
+    # :complete - triggered once all jobs have been executed, regardless of success or failure.
+    # :success - triggered once all jobs haves successfully executed.
+    # :death - triggered after any job enters the dead state.
+    # :stagnated - triggered when a job dies and no other jobs/batches are active.
+    def on(event, callback, options = {})
+      return unless Callback::VALID_CALLBACKS.include?(event.to_sym)
+
+      callback_key = "#{@bidkey}-callbacks-#{event}"
+      redis.multi do |r|
+        r.sadd(callback_key, JSON.unparse({
+          callback: callback,
+          opts: options
+        }))
+        r.expire(callback_key, BID_EXPIRE_TTL)
+      end
+    end
+
+    def jobs
+      raise NoBlockGivenError unless block_given?
+
+      persist!
+
+      # TODO This keep_open! block is probably desired, but it has some caveats:
+      #  - Old logic didn't auto-clean empty batches
+      #  - Could be an issue if the Thread crashes at a bad moment (do we even need to plan for this?)
+      #  - Technically there could be a race condition without it, but such hasn't been observed
+
+      # keep_open! do
+      begin
+        parent = Thread.current[CURRENT_BATCH_THREAD_KEY]
+        Thread.current[CURRENT_BATCH_THREAD_KEY] = self
+        yield
+      ensure
+        Thread.current[CURRENT_BATCH_THREAD_KEY] = parent
+      end
+      # end
+
+      nil
+    end
+
+    # Mark this Batch as a placeholder. It will be persisted to Redis, but no jobs will be added.
+    # From here, you can either use `.jobs` on the batch to add jobs as usual, or call `.let_close!` to allow cleanup to occur.
+    def placeholder!
+      # TODO Provide a stable `let_close!` token?
+      persist!
+    end
+
+    def invalidate_all
+      redis.setex("invalidated-bid-#{bid}", BID_EXPIRE_TTL, 1)
+    end
+
+    def parent_bid
+      redis.hget(@bidkey, "parent_bid")
+    end
+
+    def parent
+      if parent_bid
+        Batch.new(parent_bid)
+      end
+    end
+
+    def valid?(batch = self)
+      valid = !redis.exists?("invalidated-bid-#{batch.bid}")
+      batch.parent ? valid && valid?(batch.parent) : valid
+    end
+
+    def keep_open!(token = SecureRandom.urlsafe_base64(10))
+      if block_given?
+        begin
+          token = keep_open!(token)
+          yield
+        ensure
+          let_close!(token)
+        end
+      else
+        persist!
+
+        redis.multi do |r|
+          r.sadd("#{@bidkey}-holds", token)
+          r.expire("#{@bidkey}-holds", BID_EXPIRE_TTL)
+        end
+
+        assert_batch_is_open
+
+        token
+      end
+    end
+
+    def let_close!(token = :unset)
+      self.class.with_callback_check(bid, only: %i[complete success]) do |r|
+        if token == :unset # Legacy
+          r.del("#{@bidkey}-holds")
+          r.hset(@bidkey, 'keep_open', "false")
+        else
+          r.srem("#{@bidkey}-holds", token)
+        end
+      end
+    end
+
+    def self.with_batch(batch)
+      batch = self.new(batch) if batch.is_a?(String)
+      parent = Thread.current[CURRENT_BATCH_THREAD_KEY]
+      Thread.current[CURRENT_BATCH_THREAD_KEY] = batch
+      yield
+    ensure
+      Thread.current[CURRENT_BATCH_THREAD_KEY] = parent
+    end
+
+    # Any Batches or Jobs created in the given block won't be assocaiated to the current batch
+    def self.without_batch(&blk)
+      with_batch(nil, &blk)
+    end
+
+    def append_jobs(jids)
+      jids = Array(jids)
+      jids = jids.uniq
+      return unless jids.size > 0
+
+      redis do |r|
+        tme = Time.now.utc.to_f
+        added = r.zadd(@bidkey + "-jids", jids.map{|jid| [tme, jid] }, nx: true)
+        r.multi do |r|
+          r.hincrby(@bidkey, "pending", added)
+          r.hincrby(@bidkey, "job_count", added)
+          r.expire(@bidkey, BID_EXPIRE_TTL)
+          r.expire(@bidkey + "-jids", BID_EXPIRE_TTL)
+        end
+      end
+    end
+
+    protected
+
+    def redis_key
+      @bidkey
+    end
+
+    def flush_pending_attrs
+      super
+      redis.zadd("batches", created_at, bid) if INDEX_ALL_BATCHES
+    end
+
+    def assert_batch_is_open
+      unless defined?(@closed)
+        @closed = redis.hget(@bidkey, 'complete') == 'true'
+      end
+      raise "Cannot add jobs to Batch #{} bid - it has already entered the callback-stage" if @closed
+    end
+
+    def persist!
+      if !@existing && !@initialized
+        parent_bid = Thread.current[CURRENT_BATCH_THREAD_KEY]&.bid
+
+        redis.multi do |r|
+          r.hset(@bidkey, "parent_bid", parent_bid.to_s) if parent_bid
+          r.expire(@bidkey, BID_EXPIRE_TTL)
+
+          if parent_bid
+            r.hincrby("BID-#{parent_bid}", "children", 1)
+            r.expire("BID-#{parent_bid}", BID_EXPIRE_TTL)
+            r.zadd("BID-#{parent_bid}-bids", created_at, bid)
+          else
+            r.zadd("BID-ROOT-bids", created_at, bid)
+          end
+        end
+
+        flush_pending_attrs
+        @context&.save!
+
+        @initialized = true
+      end
+    end
+
+    class << self
+      def current
+        Thread.current[CURRENT_BATCH_THREAD_KEY]
+      end
+
+      def current_context
+        current&.context || {}
+      end
+
+      # Perform a success/failure/complete/etc transaction against Redis, checking
+      #   if any callbacks should be triggered
+      def with_callback_check(bid, except: [], only: Callback::VALID_CALLBACKS, &blk)
+        except = Array(except).map(&:to_sym)
+        only = Array(only).map(&:to_sym)
+
+        precount = 0
+
+        all_results = redis do |r|
+          return unless r.exists?("BID-#{bid}")
+
+          r.multi do |r|
+            # The block probably doesn't _have_ to be part of the same transaction, but it was easier than
+            #   re-assessing possible race-conditions and is technically more performant
+            blk.call(r)
+            futures = r.instance_variable_get(:@futures) || r.instance_variable_get(:@pipeline)&.futures || r.instance_variable_get(:@client)&.futures
+            precount = futures.size
+
+            # Misc
+            r.hget("BID-#{bid}", "parent_bid")
+            r.hget("BID-#{bid}", "keep_open")
+            r.scard("BID-#{bid}-holds")
+
+            # Jobs
+            r.hincrby("BID-#{bid}", "pending", 0)
+            r.scard("BID-#{bid}-failed")
+            r.scard("BID-#{bid}-dead")
+
+            # Batches
+            r.hincrby("BID-#{bid}", "children", 0)
+            r.scard("BID-#{bid}-batches-complete")
+            r.scard("BID-#{bid}-batches-success")
+            r.scard("BID-#{bid}-batches-failed")
+            r.scard("BID-#{bid}-batches-stagnated")
+
+            # Touch Expirations
+            r.expire("BID-#{bid}", BID_EXPIRE_TTL)
+            r.expire("BID-#{bid}-batches-success", BID_EXPIRE_TTL)
+          end
+        end
+
+        # Exclude return values from the passed block
+        actual_results = all_results[precount..-1]
+
+        # "pending" = not successful (yet)
+        # "failed" = dead or retrying
+        # "complete" = successful or failed
+
+        parent_bid, keep_open, holds, \
+          pending_jobs, failed_jobs, dead_jobs, \
+          child_batches, complete_batches, success_batches, failed_batches, stagnated_batches \
+          = actual_results
+
+        pending_batches = child_batches - success_batches
+
+        if keep_open == 'true' || (holds && holds > 0)
+          except << :complete
+          except << :success
+        end
+
+        trigger_callback = ->(callback) {
+          next if except.include?(callback.to_sym)
+          next unless only.include?(callback.to_sym)
+
+          Batch.logger.debug {"Finalize #{callback} bid: #{parent_bid}"}
+          enqueue_callbacks(callback, bid)
+        }
+
+        # Handling all of these cases in one method may be a little less performant than in more specialized methods, but
+        #   I was dealing with duplicate Redis queries and checks being made in up to 4 places and it was getting out of
+        #   hand - this combined method should be easier to maintain and reason about
+
+        all_successful = pending_jobs.zero? && child_batches == success_batches
+
+        if all_successful || (pending_jobs == failed_jobs && child_batches == complete_batches) # All Complete
+          trigger_callback.call(:complete)
+        end
+
+        if all_successful # All Successful
+          trigger_callback.call(:success)
+        elsif pending_jobs == dead_jobs && pending_batches == stagnated_batches # Stagnated
+          trigger_callback.call(:stagnated)
+        end
+
+        all_results[0...precount]
+      end
+
+      def process_failed_job(bid, jid)
+        with_callback_check(bid, except: [:success]) do |r|
+          r.sadd("BID-#{bid}-failed", jid)
+          r.expire("BID-#{bid}-failed", BID_EXPIRE_TTL)
+        end
+      end
+
+      # Dead jobs are a Sidekiq feature.
+      # If this is called for a job, process_failed_job was also called
+      def process_dead_job(bid, jid)
+        enqueue_callbacks(:death, bid)
+
+        with_callback_check(bid) do |r|
+          r.sadd("BID-#{bid}-dead", jid)
+          r.expire("BID-#{bid}-dead", BID_EXPIRE_TTL)
+        end
+      end
+
+      def process_successful_job(bid, jid)
+        with_callback_check(bid) do |r|
+          r.srem("BID-#{bid}-failed", jid)
+          r.hincrby("BID-#{bid}", "pending", -1)
+          r.hincrby("BID-#{bid}", "successful-jobs", 1)
+          r.zrem("BID-#{bid}-jids", jid)
+        end
+      end
+
+      def enqueue_callbacks(event, bid)
+        batch_key = "BID-#{bid}"
+        callback_key = "#{batch_key}-callbacks-#{event}"
+
+        exists, callbacks, queue, parent_bid, callback_params = redis do |r|
+          r.multi do |r|
+            r.exists?(batch_key)
+
+            r.smembers(callback_key)
+            r.hget(batch_key, "callback_queue")
+            r.hget(batch_key, "parent_bid")
+            r.hget(batch_key, "callback_params")
+          end
+        end
+
+        return unless exists
+
+        queue ||= "default"
+        parent_bid = !parent_bid || parent_bid.empty? ? nil : parent_bid # Basically parent_bid.blank?
+
+        # Internal callback params. If this is present, we're trying to enqueue callbacks for a callback, which is a special case that
+        # indicates that the callback completed and we need to close the triggering batch (which is in a done-but-not-cleaned state)
+        callback_params = JSON.parse(callback_params) if callback_params.present?
+
+        # User-configured parameters/arguments to pass to the callback
+        callback_args = callbacks.reduce([]) do |memo, jcb|
+          cb = JSON.load(jcb)
+          memo << [cb['callback'], event.to_s, cb['opts'], bid, parent_bid]
+        end
+
+        opts = {"bid" => bid, "event" => event}
+        should_schedule_batch = callback_args.present? && !callback_params.present?
+        already_processed = redis do |r|
+          SCHEDULE_CALLBACK.call(r, [batch_key], [event.to_s, should_schedule_batch.to_s, BID_EXPIRE_TTL])
+        end
+
+        return if already_processed == 'true'
+
+        if should_schedule_batch
+          logger.debug {"Enqueue callback bid: #{bid} event: #{event} args: #{callback_args.inspect}"}
+
+          # Create a new Batch to handle the callbacks and add it to the _parent_ batch
+          #   (this ensures that the parent's lifecycle status can't change until the child's callbacks are done)
+          with_batch(parent_bid) do
+            cb_batch = self.new
+            cb_batch.callback_params = {
+              for_bid: bid,
+              event: event,
+            }
+            opts['callback_bid'] = cb_batch.bid
+
+            logger.debug {"Adding callback batch: #{cb_batch.bid} for batch: #{bid}"}
+            cb_batch.jobs do
+              push_callbacks(callback_args, queue)
+            end
+          end
+        end
+
+        if callback_params.present?
+          # This is a callback for a callback. Passing `origin` to the Finalizer allows it to also cleanup the original/callback-triggering batch
+          opts['origin'] = callback_params
+        end
+
+        # The Finalizer marks this batch as complete, bumps any necessary counters, cleans up this Batch _if_ no callbacks were scheduled,
+        #   and enqueues parent-Batch callbacks if needed.
+        logger.debug {"Run batch finalizer bid: #{bid} event: #{event} args: #{callback_args.inspect}"}
+        finalizer = Batch::Callback::Finalize.new
+        status = Status.new bid
+        finalizer.dispatch(status, opts)
+      end
+
+      def cleanup_redis(bid)
+        logger.debug {"Cleaning redis of batch #{bid}"}
+        redis do |r|
+          r.zrem("batches", bid)
+          r.zrem("BID-ROOT-bids", bid)
+          r.unlink(
+            "BID-#{bid}",
+            "BID-#{bid}-callbacks-complete",
+            "BID-#{bid}-callbacks-success",
+            "BID-#{bid}-failed",
+            "BID-#{bid}-dead",
+
+            "BID-#{bid}-batches-success",
+            "BID-#{bid}-batches-complete",
+            "BID-#{bid}-batches-failed",
+            "BID-#{bid}-bids",
+            "BID-#{bid}-jids",
+            "BID-#{bid}-pending_callbacks",
+          )
+        end
+      end
+
+      def delete_prematurely!(bid)
+        child_bids = redis do |r|
+          r.zrange("BID-#{bid}-bids", 0, -1)
+        end
+        child_bids.each do |cbid|
+          delete_prematurely!(cbid)
+        end
+        cleanup_redis(bid)
+      end
+
+      # Internal method to cleanup a Redis Hash and related keys
+      def cleanup_redis_index_for(key, suffixes = [""])
+        redis do |r|
+          if r.hget(k, "created_at").present?
+            r.multi do |r|
+              suffixes.each do |suffix|
+                r.expire(key + suffix, BID_EXPIRE_TTL)
+              end
+            end
+            false
+          else
+            r.multi do |r|
+              suffixes.each do |suffix|
+                r.unlink(key + suffix)
+              end
+            end
+            true
+          end
+        end
+      end
+
+      # Administrative/console method to cleanup expired batches from the WebUI
+      def cleanup_redis_index!
+        suffixes = ["", "-callbacks-complete", "-callbacks-success", "-failed", "-dead", "-batches-success", "-batches-complete", "-batches-failed", "-bids", "-jids", "-pending_callbacks"]
+
+        redis do |r|
+          cleanup_index = ->(index) {
+            r.zrangebyscore(index, "0", BID_EXPIRE_TTL.seconds.ago.to_i).each do |bid|
+              r.zrem(index, bid) if cleanup_redis_index_for("BID-#{bid}", suffixes)
+            end
+          }
+
+          cleanup_index.("BID-ROOT-bids")
+          cleanup_index.("batches")
+        end
+      end
+
+      def redis(&blk)
+        ::Joblin.redis(&blk)
+      end
+
+      def logger
+        ::Joblin.logger
+      end
+
+      def push_callbacks(args, queue)
+        Batch::Callback::worker_class.enqueue_all(args, queue)
+      end
+
+      def bid_hierarchy(bid, depth: 4, per_depth: 5, slice: nil)
+        args = [bid, depth, per_depth]
+        args << slice if slice
+        redis do |r|
+          BID_HIERARCHY.call(r, [], args)
+        end
+      end
+    end
+  end
+end
+
+require_relative './compat'

data/lib/joblin/batching/callback.rb

@@ -0,0 +1,135 @@
+module Joblin::Batching
+  class Batch
+    module Callback
+      mattr_accessor :worker_class
+
+      VALID_CALLBACKS = %i[success complete death stagnated].freeze
+
+      module CallbackWorkerCommon
+        def perform(definition, event, opts, bid, parent_bid)
+          return unless VALID_CALLBACKS.include?(event.to_sym)
+
+          method = nil
+          target = :instance
+          clazz = definition
+          if clazz.is_a?(String)
+            if clazz.include?('#')
+              clazz, method = clazz.split("#")
+            elsif clazz.include?('.')
+              clazz, method = clazz.split(".")
+              target = :class
+            end
+          end
+
+          method ||= "on_#{event}"
+          status = Batch::Status.new(bid)
+
+          if clazz && object = Object.const_get(clazz)
+            target = target == :instance ? object.new : object
+            if target.respond_to?(method, true)
+              target.send(method, status, opts.with_indifferent_access)
+            else
+              Batch.logger.warn("Invalid callback method #{definition} - #{target.to_s} does not respond to #{method}")
+            end
+          else
+            Batch.logger.warn("Invalid callback method #{definition} - Class #{clazz} not found")
+          end
+        end
+      end
+
+      class Finalize
+        # The methods in this class are called after all same-named callbacks have been
+        # completed for the passed Batch.
+        # These methods mainly handle bubbling events up to the parent Batch
+        # You could say that they are the callbacks for callbacks.
+
+        def dispatch(status, opts)
+          bid = opts["bid"]
+          event = opts["event"].to_sym
+
+          Batch.logger.debug {"Finalize #{event} batch id: #{opts["bid"]}"}
+
+          batch_status = Status.new bid
+          send(event, bid, batch_status, batch_status.parent_bid)
+
+          Batch.redis do |r|
+            r.srem("BID-#{bid}-pending_callbacks", "#{event}-finalize")
+          end
+
+          if event == :success
+            if opts['origin'].present?
+              # This is a callback for a callback. In this case we need to check if we should cleanup the original bid.
+              origin_bid = opts['origin']['for_bid']
+              _, pending, success_ran = Batch.redis do |r|
+                r.multi do |r|
+                  r.srem("BID-#{origin_bid}-pending_callbacks", opts['origin']['event'])
+                  r.scard("BID-#{origin_bid}-pending_callbacks")
+                  r.hget("BID-#{origin_bid}", "success")
+                end
+              end
+              Batch.cleanup_redis(origin_bid) if pending == 0 && success_ran == 'true'
+            end
+
+            if (Batch.redis {|r| r.scard("BID-#{bid}-pending_callbacks") }) == 0
+              Batch.cleanup_redis(bid)
+            end
+          end
+        end
+
+        def success(bid, status, parent_bid)
+          return unless parent_bid
+
+          Batch.with_callback_check(parent_bid) do |r|
+            r.sadd("BID-#{parent_bid}-batches-success", bid)
+            r.srem("BID-#{parent_bid}-batches-failed", bid)
+            r.sadd("BID-#{parent_bid}-batches-complete", bid)
+          end
+        end
+
+        def complete(bid, status, parent_bid)
+          return unless parent_bid
+
+          pending, children, success = Batch.redis do |r|
+            r.multi do |r|
+              r.hincrby("BID-#{bid}", "pending", 0)
+              r.hincrby("BID-#{bid}", "children", 0)
+              r.scard("BID-#{bid}-batches-success")
+            end
+          end
+
+          if !(pending.to_i.zero? && children == success)
+            # If batch was not successfull check and see if its parent is complete
+            # if the parent is complete we can trigger its complete callback.
+            #
+            # Otherwise, we don't want to to trigger the parent's :complete here (and
+            # we instead opt to have success tigger parent :complete) - this
+            # allows the success callback to add additional jobs to the parent batch
+            # before triggering :complete.
+
+            Batch.with_callback_check(parent_bid, except: [:success]) do |r|
+              r.sadd("BID-#{parent_bid}-batches-complete", bid)
+              r.sadd("BID-#{parent_bid}-batches-failed", bid)
+            end
+          end
+        end
+
+        def death(bid, status, parent_bid)
+          return unless parent_bid
+
+          # We only need to bubble the event here - other events (eg stagnation) will be checked and bubbled elsewhere.
+
+          Batch.enqueue_callbacks(:death, parent_bid)
+        end
+
+        def stagnated(bid, status, parent_bid)
+          return unless parent_bid
+
+          Batch.with_callback_check(parent_bid) do |r|
+            r.sadd("BID-#{parent_bid}-batches-stagnated", bid)
+            r.expire("BID-#{parent_bid}-batches-stagnated", BID_EXPIRE_TTL)
+          end
+        end
+      end
+    end
+  end
+end