joblin 0.1.0

data/lib/joblin/batching/compat/sidekiq.rb

@@ -0,0 +1,149 @@
+begin
+  require 'sidekiq/batch'
+rescue LoadError
+end
+
+module Joblin::Batching
+  module Compat
+    module Sidekiq
+      module WorkerExtension
+        def bid
+          Thread.current[CURRENT_BATCH_THREAD_KEY].bid
+        end
+
+        def batch
+          Thread.current[CURRENT_BATCH_THREAD_KEY]
+        end
+
+        def batch_context
+          batch&.context || {}
+        end
+
+        def valid_within_batch?
+          batch.valid?
+        end
+      end
+
+      class SidekiqCallbackWorker
+        include ::Sidekiq::Worker
+        include WorkerExtension
+        include Batch::Callback::CallbackWorkerCommon
+
+        def self.enqueue_all(args, queue)
+          return if args.empty?
+
+          ::Sidekiq::Client.push_bulk(
+            'class' => self,
+            'args' => args,
+            'queue' => queue
+          )
+        end
+      end
+
+      class ClientMiddleware
+        include ::Sidekiq::ClientMiddleware if defined? ::Sidekiq::ClientMiddleware
+
+        def call(_worker, msg, _queue, _redis_pool = nil)
+          if (batch = Thread.current[CURRENT_BATCH_THREAD_KEY]) && should_handle_batch?(msg)
+            batch.append_jobs(msg['jid']) if (msg['bid'] = batch.bid)
+          end
+          yield
+        end
+
+        def should_handle_batch?(msg)
+          return false if Joblin::Batching::Compat::Sidekiq.is_activejob_job?(msg)
+          true
+        end
+      end
+
+      class ServerMiddleware
+        include ::Sidekiq::ServerMiddleware if defined? ::Sidekiq::ServerMiddleware
+
+        def call(_worker, msg, _queue)
+          if (bid = msg['bid'])
+            prev_batch = Thread.current[CURRENT_BATCH_THREAD_KEY]
+            begin
+              Thread.current[CURRENT_BATCH_THREAD_KEY] = Batch.new(bid)
+              yield
+              Thread.current[CURRENT_BATCH_THREAD_KEY].save_context_changes
+              Batch.process_successful_job(bid, msg['jid'])
+            rescue
+              Batch.process_failed_job(bid, msg['jid'])
+              raise
+            ensure
+              Thread.current[CURRENT_BATCH_THREAD_KEY] = prev_batch
+            end
+          else
+            yield
+          end
+        end
+      end
+
+      def self.is_activejob_job?(msg)
+        return false unless defined?(::ActiveJob)
+
+        msg['class'] == 'ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper' && (msg['wrapped'].to_s).constantize < Joblin::Batching::Compat::ActiveJob::BatchAwareJob
+      end
+
+      def self.switch_tenant(job)
+        if defined?(::Apartment)
+          ::Apartment::Tenant.switch(job['apartment'] || 'public') do
+            yield
+          end
+        else
+          yield
+        end
+      end
+
+      def self.sidekiq_middleware(placement, &blk)
+        install_middleware = ->(config) do
+          config.send("#{placement}_middleware") do |chain|
+            blk.call(chain)
+          end
+        end
+
+        ::Sidekiq.configure_client(&install_middleware) if placement == :client
+        ::Sidekiq.configure_server(&install_middleware)
+      end
+
+      def self.configure
+        return if @already_configured
+        @already_configured = true
+
+        if defined?(::Sidekiq::Batch) && ::Sidekiq::Batch != Joblin::Batching::Batch
+          print "WARNING: Detected Sidekiq Pro or sidekiq-batch. Joblin JobBatches may not be fully compatible!"
+        end
+
+        sidekiq_middleware(:client) do |chain|
+          chain.remove ::Sidekiq::Batch::Middleware::ClientMiddleware if defined?(::Sidekiq::Batch::Middleware::ClientMiddleware)
+          chain.add Joblin::Batching::Compat::Sidekiq::ClientMiddleware
+        end
+
+        sidekiq_middleware(:server) do |chain|
+          chain.remove ::Sidekiq::Batch::Middleware::ServerMiddleware if defined?(::Sidekiq::Batch::Middleware::ServerMiddleware)
+          chain.add Joblin::Batching::Compat::Sidekiq::ServerMiddleware
+        end
+
+        ::Sidekiq.configure_server do |config|
+          config.death_handlers << ->(job, ex) do
+            switch_tenant(job) do
+              if is_activejob_job?(job)
+                Joblin::Batching::Compat::ActiveJob.handle_job_death(job["args"][0], ex)
+              elsif job['bid'].present?
+                ::Sidekiq::Batch.process_dead_job(job['bid'], job['jid'])
+              end
+            end
+          end
+        end
+
+        ::Sidekiq.const_set(:Batch, Joblin::Batching::Batch)
+        # This alias helps apartment-sidekiq set itself up correctly
+        ::Sidekiq::Batch.const_set(:Server, Joblin::Batching::Compat::Sidekiq::ServerMiddleware)
+        ::Sidekiq::Worker.send(:include, Joblin::Batching::Compat::Sidekiq::WorkerExtension)
+        Batch::Callback.worker_class = SidekiqCallbackWorker
+      end
+    end
+  end
+end
+
+require_relative 'sidekiq/web'

data/lib/joblin/batching/compat.rb

@@ -0,0 +1,20 @@
+
+module Joblin::Batching
+  module Compat
+    def self.load_compat(name)
+      name = name.to_s
+      begin
+        require name
+      rescue LoadError
+      end
+
+      if name.classify.safe_constantize
+        require_relative "./compat/#{name}"
+        "Joblin::Batching::Compat::#{name.classify}".constantize.configure
+      end
+    end
+
+    load_compat(:active_job)
+    load_compat(:sidekiq)
+  end
+end

data/lib/joblin/batching/context_hash.rb

@@ -0,0 +1,157 @@
+module Joblin::Batching
+  class ContextHash
+    delegate_missing_to :flatten
+
+    def initialize(bid, hash = nil)
+      @bid_stack = [bid]
+      @hash_map = {}
+      @dirty = false
+      @flattened = nil
+      @hash_map[bid] = hash.with_indifferent_access if hash
+    end
+
+    # Local is "the nearest batch with a context value"
+    # This allows for, for example, SerialBatchJob to have a modifiable context stored on it's main Batch
+    # that can be accessed transparently from one of it's internal, context-less Batches
+    def local_bid
+      bid = @bid_stack[-1]
+      while bid.present?
+        bhash = resolve_hash(bid)
+        return bid if bhash
+        bid = get_parent_bid(bid)
+      end
+      nil
+    end
+
+    def local
+      @hash_map[local_bid]
+    end
+
+    def own
+      resolve_hash(@bid_stack[-1]) || {}
+    end
+
+    def set_local(new_hash)
+      @dirty = true
+      local.clear.merge!(new_hash)
+    end
+
+    def clear
+      local.clear
+      @flattened = nil
+      @dirty = true
+      self
+    end
+
+    def []=(key, value)
+      @flattened = nil
+      @dirty = true
+      local[key] = value
+    end
+
+    def [](key)
+      bid = @bid_stack[-1]
+      while bid.present?
+        bhash = resolve_hash(bid)
+        return bhash[key] if bhash&.key?(key)
+        bid = get_parent_bid(bid)
+      end
+      nil
+    end
+
+    def reload!
+      @dirty = false
+      @hash_map = {}
+      self
+    end
+
+    def save!(force: false)
+      return unless dirty? || force
+      Batch.redis do |r|
+        r.hset("BID-#{local_bid}", 'context', JSON.unparse(local))
+      end
+    end
+
+    def dirty?
+      @dirty
+    end
+
+    def is_a?(arg)
+      return true if Hash <= arg
+      super
+    end
+
+    def flatten
+      return @flattened if @flattened
+
+      load_all
+      flattened = {}
+      @bid_stack.compact.each do |bid|
+        flattened.merge!(@hash_map[bid]) if @hash_map[bid]
+      end
+      flattened.freeze
+
+      @flattened = flattened.with_indifferent_access
+    end
+
+    def to_h
+      flatten
+    end
+
+    private
+
+    def get_parent_hash(bid)
+      resolve_hash(get_parent_bid(bid)).freeze
+    end
+
+    def get_parent_bid(bid)
+      index = @bid_stack.index(bid)
+      raise "Invalid BID #{bid}" if index.nil? # Sanity Check - this shouldn't happen
+
+      index -= 1
+      if index >= 0
+        @bid_stack[index]
+      else
+        pbid = Batch.redis do |r|
+          callback_params = JSON.parse(r.hget("BID-#{bid}", "callback_params") || "{}")
+          callback_params['for_bid'] || r.hget("BID-#{bid}", "parent_bid")
+        end
+        @bid_stack.unshift(pbid)
+        pbid
+      end
+    end
+
+    def resolve_hash(bid)
+      return nil unless bid.present?
+      return @hash_map[bid] if @hash_map.key?(bid)
+
+      context_json, editable = Batch.redis do |r|
+        r.multi do |r|
+          r.hget("BID-#{bid}", "context")
+          r.hget("BID-#{bid}", "allow_context_changes")
+        end
+      end
+
+      if context_json.present?
+        context_hash = JSON.parse(context_json)
+        context_hash = context_hash.with_indifferent_access
+        context_hash.each do |k, v|
+          v.freeze
+        end
+        context_hash.freeze unless editable
+
+        @hash_map[bid] = context_hash
+      else
+        @hash_map[bid] = nil
+      end
+    end
+
+    def load_all
+      resolve_hash(@bid_stack[0]).freeze
+      while @bid_stack[0].present?
+        get_parent_hash(@bid_stack[0])
+      end
+      @hash_map
+    end
+  end
+end

data/lib/joblin/batching/hier_batch_ids.lua

@@ -0,0 +1,25 @@
+
+local function add_bids(root, depth)
+    local result_data = {}
+
+    if depth > 0 then
+        local sbids
+        if depth == tonumber(ARGV[2]) and ARGV[4] then
+            local min, max = ARGV[4]:match('(%d+):(%d+)')
+            sbids = redis.call('ZRANGE', 'BID-' .. root .. '-bids', min, max)
+        else
+            sbids = redis.call('ZRANGE', 'BID-' .. root .. '-bids', 0, tonumber(ARGV[3]) - 1)
+        end
+
+        local sub_data = {}
+        for _,v in ipairs(sbids) do
+            table.insert(sub_data, add_bids(v, depth - 1))
+        end
+
+        return { root, sub_data }
+    end
+
+    return { root, result_data}
+end
+
+return add_bids(ARGV[1], tonumber(ARGV[2]))

data/lib/joblin/batching/jobs/base_job.rb

@@ -0,0 +1,7 @@
+require 'active_job'
+
+module Joblin::Batching
+  class BaseJob < ::ActiveJob::Base
+
+  end
+end

data/lib/joblin/batching/jobs/concurrent_batch_job.rb

@@ -0,0 +1,20 @@
+require_relative './base_job'
+
+module Joblin::Batching
+  class ConcurrentBatchJob < BaseJob
+    def self.make_batch(sub_jobs, **kwargs, &blk)
+      ManagedBatchJob.make_batch(
+        sub_jobs,
+        concurrency: true,
+        **kwargs,
+        ordered: false,
+        desc_prefix: 'ConcurrentBatchJob: ',
+        &blk
+      )
+    end
+
+    def perform(sub_jobs, **kwargs)
+      self.class.make_batch(sub_jobs, **kwargs)
+    end
+  end
+end

data/lib/joblin/batching/jobs/managed_batch_job.rb

@@ -0,0 +1,175 @@
+require_relative './base_job'
+
+module Joblin::Batching
+  class ManagedBatchJob < BaseJob
+    def self.make_batch(sub_jobs, ordered: true, concurrency: nil, context: nil, preflight_check: nil, desc_prefix: nil, &blk)
+      desc_prefix ||= ''
+
+      if concurrency == 0 || concurrency == nil || concurrency == true
+        concurrency = sub_jobs.count
+      elsif concurrency == false
+        concurrency = 1
+      end
+
+      root_batch = Batch.new
+      man_batch_id = nil
+
+      if concurrency < sub_jobs.count
+        man_batch_id = SecureRandom.urlsafe_base64(10)
+
+        Batch.redis do |r|
+          r.multi do |r|
+            r.hset("MNGBID-#{man_batch_id}", "root_bid", root_batch.bid)
+            r.hset("MNGBID-#{man_batch_id}", "ordered", ordered ? 1 : 0)
+            r.hset("MNGBID-#{man_batch_id}", "concurrency", concurrency)
+            r.hset("MNGBID-#{man_batch_id}", "preflight_check", preflight_check) if preflight_check.present?
+            r.expire("MNGBID-#{man_batch_id}", Batch::BID_EXPIRE_TTL)
+
+            mapped_sub_jobs = sub_jobs.each_with_index.map do |j, i|
+              j['_mngbid_index_'] = i # This allows duplicate jobs when a Redis Set is used
+              j = ::ActiveJob::Arguments.serialize([j])
+              JSON.unparse(j)
+            end
+            if ordered
+              r.rpush("MNGBID-#{man_batch_id}-jobs", mapped_sub_jobs)
+            else
+              r.sadd("MNGBID-#{man_batch_id}-jobs", mapped_sub_jobs)
+            end
+            r.expire("MNGBID-#{man_batch_id}-jobs", Batch::BID_EXPIRE_TTL)
+          end
+        end
+
+        root_batch.allow_context_changes = (concurrency == 1)
+        root_batch.on(:success, "#{to_s}.cleanup_redis", managed_batch_id: man_batch_id)
+
+        desc_prefix = "MGD(#{man_batch_id}): #{desc_prefix}"
+      end
+
+      root_batch.context = context
+
+      blk.call(ManagedBatchProxy.new(root_batch)) if blk.present?
+
+      root_batch.description = "#{desc_prefix}#{root_batch.description || 'Root'}"
+
+      root_batch.context["managed_batch_bid"] = man_batch_id if man_batch_id
+
+      if concurrency < sub_jobs.count
+        root_batch.placeholder!
+        concurrency.times do
+          perform_next_sequence_job(man_batch_id, skip_preflight: true)
+        end
+      else
+        root_batch.jobs do
+          sub_jobs.each do |j|
+            ChainBuilder.enqueue_job(j)
+          end
+        end
+      end
+
+      root_batch
+    end
+
+    def perform(sub_jobs, **kwargs)
+      self.class.make_batch(sub_jobs, **kwargs)
+    end
+
+    def self.cleanup_redis(status, options)
+      man_batch_id = options['managed_batch_id']
+      Batch.redis do |r|
+        r.del(
+          "MNGBID-#{man_batch_id}",
+          "MNGBID-#{man_batch_id}-jobs",
+        )
+      end
+    end
+
+    def self.job_succeeded_callback(status, options)
+      man_batch_id = options['managed_batch_id']
+      perform_next_sequence_job(man_batch_id)
+    end
+
+    protected
+
+    def self.perform_next_sequence_job(man_batch_id, skip_preflight: false)
+      root_bid, ordered, preflight_check = Batch.redis do |r|
+        r.multi do |r|
+          r.hget("MNGBID-#{man_batch_id}", "root_bid")
+          r.hget("MNGBID-#{man_batch_id}", "ordered")
+          r.hget("MNGBID-#{man_batch_id}", "preflight_check")
+        end
+      end
+
+      if !skip_preflight && preflight_check.present?
+        if preflight_check.include?(".")
+          clazz, method_name = preflight_check.split('.')
+          clazz = clazz.constantize
+        else
+          clazz = Object
+          method_name = preflight_check
+        end
+        preflight_check = ->(*args) { clazz.send(method_name, *args) }
+      else
+        preflight_check = ->(*args) { true }
+      end
+
+      ordered = Joblin::MiscHelper.to_boolean(ordered)
+
+      loop do
+        next_job_json = Batch.redis do |r|
+          if ordered
+            r.lpop("MNGBID-#{man_batch_id}-jobs")
+          else
+            r.spop("MNGBID-#{man_batch_id}-jobs")
+          end
+        end
+
+        break unless next_job_json.present?
+
+        next_job = JSON.parse(next_job_json)
+        next_job = ::ActiveJob::Arguments.deserialize(next_job)[0]
+
+        preflight_result = preflight_check.call(next_job)
+        if preflight_result == :abort
+          cleanup_redis(nil, { "managed_batch_id" => man_batch_id })
+          break
+        elsif !preflight_check
+          next
+        end
+
+        Batch.new(root_bid).jobs do
+          Batch.new.tap do |batch|
+            batch.description = "Managed Batch Fiber (#{man_batch_id})"
+            batch.on(:success, "#{self.to_s}.job_succeeded_callback", managed_batch_id: man_batch_id)
+
+            if next_job[:chain_link].present?
+              # Annotate Batch with chain-step info
+              batch.context["csb:chain_link"] = next_job[:chain_link]
+              batch.on(:complete, "#{ChainBuilder.to_s}.chain_step_complete", chain_link: next_job[:chain_link])
+            end
+
+            batch.jobs do
+              ChainBuilder.enqueue_job(next_job)
+            end
+          end
+        end
+
+        break
+      end
+    end
+
+    class ManagedBatchProxy
+      def initialize(real_batch)
+        @real_batch = real_batch
+      end
+
+      delegate_missing_to :real_batch
+
+      def jobs
+        raise "Managed Batches do not support calling .jobs directly!"
+      end
+
+      private
+      attr_reader :real_batch
+    end
+  end
+end

data/lib/joblin/batching/jobs/serial_batch_job.rb

@@ -0,0 +1,20 @@
+require_relative './base_job'
+
+module Joblin::Batching
+  class SerialBatchJob < BaseJob
+    def self.make_batch(sub_jobs, **kwargs, &blk)
+      ManagedBatchJob.make_batch(
+        sub_jobs,
+        **kwargs,
+        ordered: true,
+        concurrency: false,
+        desc_prefix: 'SerialBatchJob: ',
+        &blk
+      )
+    end
+
+    def perform(sub_jobs, **kwargs)
+      self.class.make_batch(sub_jobs, **kwargs)
+    end
+  end
+end