joblin 0.1.0

data/lib/joblin/batching/pool.rb

@@ -0,0 +1,254 @@
+module Joblin::Batching
+  class Pool
+    include RediConn::RedisModel
+
+    POOL_REFILL = RediConn::RedisScript.new(Pathname.new(__FILE__) + "../pool_refill.lua")
+
+    attr_reader :pid
+    redis_attr :description
+    redis_attr :created_at
+    redis_attr :concurrency, :int
+    redis_attr :complete_count, :int
+    redis_attr :order
+    redis_attr :on_failed_job, :symbol
+    redis_attr :clean_when_empty, :bool
+
+    def initialize(pooolid = nil, **kwargs)
+      if pooolid
+        @existing = true
+        @pid = pooolid
+      else
+        @pid = SecureRandom.urlsafe_base64(10)
+        initialize_new(**kwargs)
+      end
+    end
+
+    def self.from_pid(pid)
+      raise "PID must be given" unless pid.present?
+      new(pid)
+    end
+
+    def <<(job_desc)
+      add_job(job_desc)
+    end
+
+    def add_job(job_desc)
+      add_jobs([job_desc])
+    end
+
+    def add_jobs(job_descs, skip_refill: false)
+      job_descs.each do |job_desc|
+        wrapper = Batch.new
+        wrapper.description = "Pool Job Wrapper (PID: #{pid})"
+        checkin_event = (on_failed_job == :wait) ? :success : :complete
+        wrapper.on(checkin_event, "#{self.class.to_s}.job_checked_in", pool_id: pid)
+        wrapper.placeholder!
+
+        job_desc = job_desc.symbolize_keys
+        job_desc = job_desc.merge!(
+          job: job_desc[:job].to_s,
+          pool_wrapper_batch: wrapper.bid,
+        )
+
+        push_job_to_pool(job_desc)
+      end
+      refill_allotment unless skip_refill
+    end
+
+    def keep_open!(token = SecureRandom.urlsafe_base64(10))
+      if block_given?
+        begin
+          token = keep_open!(token)
+          yield
+        ensure
+          let_close!(token)
+        end
+      else
+        redis.multi do |r|
+          r.sadd("#{redis_key}-holds", token)
+          r.expire("#{redis_key}-holds", Batch::BID_EXPIRE_TTL)
+        end
+        token
+      end
+    end
+
+    def let_close!(token = :unset)
+      if token == :unset # Legacy
+        redis.del("#{redis_key}-holds")
+        redis.hset(redis_key, 'keep_open', 'false')
+      else
+        redis.srem("#{redis_key}-holds", token)
+      end
+
+      cleanup_if_empty
+    end
+
+    def cleanup_redis
+      Batch.logger.debug {"Cleaning redis of pool #{pid}"}
+      redis do |r|
+        r.zrem("pools", pid)
+        r.unlink(
+          "#{redis_key}",
+          "#{redis_key}-jobs",
+        )
+      end
+    end
+
+    def cleanup_if_empty
+      self.order
+
+      activec, pactivec, pendingc, clean_when_empty, keep_open, holds = redis.multi do |r|
+        r.hlen("#{redis_key}-active")
+        r.hget(redis_key, "_active_count")
+        pending_count(r)
+        r.hget(redis_key, 'clean_when_empty')
+        r.hget(redis_key, 'keep_open')
+        r.scard("#{redis_key}-holds")
+      end
+
+      return if keep_open == 'true' || clean_when_empty == 'false' || (holds && holds > 0)
+
+      if activec <= 0 && (pactivec.try(:to_i) || 0) <= 0 && pendingc <= 0
+        cleanup_redis
+      end
+    end
+
+    def active_count(r = redis)
+      r.hlen("#{redis_key}-active") + r.hincrby(redis_key, "_active_count", 0)
+    end
+
+    def active_jobs(r = redis)
+      r.hvals("#{redis_key}-active").map {|desc| JSON.parse(desc)[0] }
+    end
+
+    def pending_count(r = redis)
+      jobs_key = "#{redis_key}-jobs"
+      order = self.order || 'fifo'
+      case order.to_sym
+      when :fifo, :lifo
+        r.llen(jobs_key)
+      when :random
+        r.scard(jobs_key)
+      when :priority
+        r.zcard(jobs_key)
+      end
+    end
+
+    def job_checked_in(status, options)
+      active_count = refill_allotment(status.bid)
+      cleanup_if_empty unless active_count > 0
+    end
+
+    def self.job_checked_in(status, options)
+      pid = options['pool_id']
+      from_pid(pid).job_checked_in(status, options)
+    end
+
+    # Administrative/console method to cleanup expired pools from the WebUI
+    def self.cleanup_redis_index!
+      suffixes = ["", "-active", "-jobs"]
+      r.zrangebyscore("pools", "0", Batch::BID_EXPIRE_TTL.seconds.ago.to_i).each do |pid|
+        r.zrem("pools", pid) if Batch.cleanup_redis_index_for("POOLID-#{pid}", suffixes)
+      end
+    end
+
+    protected
+
+    def redis_key
+      "POOLID-#{pid}"
+    end
+
+    def refill_allotment(checkin_bid = nil)
+      active_count, job_descs = POOL_REFILL.call(redis, [redis_key, "#{redis_key}-jobs", "#{redis_key}-active"], [checkin_bid || ""])
+      return active_count if active_count < 0
+
+      pending_job_descs = job_descs.dup
+
+      added_jobs = {}
+      failed_to_add_jobs = []
+      add_exception = nil
+
+      while pending_job_descs.count > 0
+        begin
+          job_json = pending_job_descs.shift
+          job_desc = ::ActiveJob::Arguments.deserialize(JSON.parse(job_json))[0]&.symbolize_keys
+
+          wbid = job_desc[:pool_wrapper_batch]
+
+          Batch.new(wbid).jobs do
+            ChainBuilder.enqueue_job(job_desc)
+          end
+
+          added_jobs[wbid] = job_json
+        rescue => ex
+          failed_to_add_jobs << job_json
+          add_exception = ex
+        end
+      end
+
+      redis.multi do |r|
+        r.mapped_hmset("#{redis_key}-active", added_jobs) if added_jobs.count > 0
+        # Release reserved slots now that we've added the jobs to `-active`
+        r.hincrby(redis_key, "_active_count", -job_descs.count)
+
+        r.expire(redis_key, Batch::BID_EXPIRE_TTL)
+        r.expire("#{redis_key}-active", Batch::BID_EXPIRE_TTL)
+        r.expire("#{redis_key}-jobs", Batch::BID_EXPIRE_TTL)
+      end
+
+      # If this happens, we end up in a bad state (as we don't try to re-add items to the pool or refill_allotment again), but
+      #   this should be a _really_ rare case that should only occur if we've lost connection to Redis or something, so we're
+      #   operating on the assumption that if we get here, any recovery logic will fail too
+      if add_exception.present?
+        Batch.logger.error {"Error popping jobs from Pool #{pid}: #{add_exception}"}
+        raise add_exception
+      end
+
+      active_count + added_jobs.count
+    end
+
+    def push_job_to_pool(job_desc)
+      jobs_key = "#{redis_key}-jobs"
+      # This allows duplicate jobs when a Redis Set is used
+      job_desc[:_pool_random_key_] = SecureRandom.urlsafe_base64(10)
+      job_json = JSON.unparse(::ActiveJob::Arguments.serialize([job_desc]))
+      order = self.order
+
+      redis.multi do |r|
+        case order.to_sym
+        when :fifo, :lifo
+          r.rpush(jobs_key, job_json)
+        when :random
+          r.sadd(jobs_key, job_json)
+        when :priority
+          r.zadd(jobs_key, job_desc[:priority] || 0, job_json)
+        end
+        r.expire(redis_key, Batch::BID_EXPIRE_TTL)
+        r.expire(jobs_key, Batch::BID_EXPIRE_TTL)
+      end
+    end
+
+    def self.redis(&blk)
+      Batch.redis &blk
+    end
+    delegate :redis, to: :class
+
+    def flush_pending_attrs
+      super
+      redis.expire(redis_key, Batch::BID_EXPIRE_TTL)
+      redis.zadd("pools", created_at, pid)
+    end
+
+    private
+
+    def initialize_new(concurrency: nil, order: :fifo, clean_when_empty: true, on_failed_job: :wait, description: nil)
+      self.created_at = Time.now.utc.to_f
+      self.description = description
+      self.order = order
+      self.concurrency = concurrency
+      self.clean_when_empty = clean_when_empty
+      self.on_failed_job = on_failed_job
+      flush_pending_attrs
+    end
+  end
+end

data/lib/joblin/batching/pool_refill.lua

@@ -0,0 +1,47 @@
+
+local poolkey = KEYS[1]
+local qkey = KEYS[2]
+local activekey = KEYS[3]
+
+local checkin_item = ARGV[1]
+
+if redis.call('EXISTS', poolkey) == 0 then
+    return { -1, {} } -- pool doesn't exist
+end
+
+if checkin_item ~= "" then
+    redis.call("HDEL", activekey, checkin_item)
+    redis.call("HINCRBY", poolkey, "complete_count", 1)
+end
+
+local pool_type = redis.call('HGET', poolkey, "order")
+local allotment = tonumber(redis.call("HGET", poolkey, "concurrency"))
+local active = redis.call("HLEN", activekey) + (redis.call("HGET", poolkey, "_active_count") or 0)
+
+local pop_count = allotment - active
+
+local popped_items = {}
+
+if pop_count > 0 then
+    if pool_type == "fifo" then
+        popped_items = redis.call("LPOP", qkey, pop_count) or {}
+    elseif pool_type == "lifo" then
+        popped_items = redis.call("RPOP", qkey, pop_count) or {}
+    elseif pool_type == "random" then
+        popped_items = redis.call("SPOP", qkey, pop_count) or {}
+    elseif pool_type == "priority" then
+        local temp_items = redis.call("ZPOPMAX", qkey, pop_count) or {}
+        for i,v in ipairs(temp_items) do
+            if i % 2 == 1 then
+                table.insert(popped_items, v)
+            end
+        end
+    end
+end
+
+-- Reserve slots for these jobs while we return to Ruby and deserialize them
+--  This could also be inlined by just storing a key in the queue and storing parameters
+--  in a Hash, but this seems more efficient.
+redis.call('HINCRBY', poolkey, "_active_count", #popped_items)
+
+return { active, popped_items }

data/lib/joblin/batching/schedule_callback.lua

@@ -0,0 +1,14 @@
+
+local previously_scheduled = redis.call('HGET', KEYS[1], ARGV[1])
+redis.call('HSET', KEYS[1], ARGV[1], 'true')
+
+if previously_scheduled ~= 'true' then
+    local pcb_key = KEYS[1] .. '-pending_callbacks'
+    redis.call('SADD', pcb_key, ARGV[1] .. '-finalize')
+    if ARGV[2] == 'true' then
+        redis.call('SADD', pcb_key, ARGV[1])
+    end
+    redis.call('EXPIRE', pcb_key, ARGV[3])
+end
+
+return previously_scheduled

data/lib/joblin/batching/status.rb

@@ -0,0 +1,89 @@
+module Joblin::Batching
+  class Batch
+    class Status
+      attr_reader :bid
+
+      def initialize(bid)
+        bid = bid.bid if bid.is_a?(Batch)
+        @bid = bid
+      end
+
+      def join
+        raise "Not supported"
+      end
+
+      def pending
+        Batch.redis { |r| r.hget("BID-#{bid}", 'pending') }.to_i
+      end
+
+      def failures
+        Batch.redis { |r| r.scard("BID-#{bid}-failed") }.to_i
+      end
+
+      def dead
+        Batch.redis { |r| r.scard("BID-#{bid}-dead") }.to_i
+      end
+
+      def completed_count
+        job_count - pending
+      end
+
+      def job_count
+        Batch.redis { |r| r.hget("BID-#{bid}", "job_count") }.to_i
+      end
+
+      def created_at
+        Batch.redis { |r| r.hget("BID-#{bid}", 'created_at') }
+      end
+
+      def parent_bid
+        Batch.redis { |r| r.hget("BID-#{bid}", "parent_bid") }
+      end
+
+      def failure_info
+        Batch.redis { |r| r.smembers("BID-#{bid}-failed") } || []
+      end
+
+      def complete?
+        'true' == Batch.redis { |r| r.hget("BID-#{bid}", 'complete') }
+      end
+
+      def success?
+        'true' == Batch.redis { |r| r.hget("BID-#{bid}", 'success') }
+      end
+
+      def child_count
+        Batch.redis { |r| r.hget("BID-#{bid}", 'children') }.to_i
+      end
+
+      def completed_children_count
+        Batch.redis { |r| r.scard("BID-#{bid}-batches-complete") }.to_i
+      end
+
+      def successful_children_count
+        Batch.redis { |r| r.scard("BID-#{bid}-batches-success") }.to_i
+      end
+
+      def failed_children_count
+        Batch.redis { |r| r.scard("BID-#{bid}-batches-failed") }.to_i
+      end
+
+      def data
+        {
+          bid: bid,
+          failures: failures,
+          pending: pending,
+          created_at: created_at,
+          complete: complete?,
+          success: success?,
+          failure_info: failure_info,
+          parent_bid: parent_bid,
+          child_count: child_count,
+          completed_children_count: completed_children_count,
+          successful_children_count: successful_children_count,
+          failed_children_count: failed_children_count,
+        }
+      end
+    end
+  end
+end

data/lib/joblin/engine.rb

@@ -0,0 +1,15 @@
+require "rails"
+
+module Joblin
+  class Engine < ::Rails::Engine
+    isolate_namespace Joblin
+
+    initializer :append_migrations do |app|
+      config.paths["db/migrate"].expanded.each do |expanded_path|
+        app.config.paths["db/migrate"] << expanded_path
+      end
+      # Apartment will modify this, but it doesn't fully support engine migrations, so we'll reset it here
+      ActiveRecord::Migrator.migrations_paths = Rails.application.paths["db/migrate"].to_a
+    end
+  end
+end

data/lib/joblin/lazy_access.rb

@@ -0,0 +1,72 @@
+module Joblin
+  module LazyAccess
+    extend ActiveSupport::Concern
+
+    def initialize(...)
+      super
+      @access_cache = {}
+    end
+
+    def self.raw
+      current = Thread.current[:lazy_access_read_raw]
+      Thread.current[:lazy_access_read_raw] = true
+      yield
+    ensure
+      Thread.current[:lazy_access_read_raw] = current
+    end
+
+    def self.load(val)
+      case val
+      when String
+        if val.start_with?("gid://")
+          val = GlobalID::Locator.locate(val)
+        end
+      when Hash
+        val = LazyAccessHash.new(val)
+      when Array
+        val = LazyAccessArray.new(val)
+      end
+      val
+    end
+
+    def self.dump(val)
+      case val
+      when Array
+        val.to_a.map {|x| dump(x) }
+      when Hash
+        val.to_h.transform_values{|x| dump(x) }
+      when ActiveRecord::Base
+        val.to_gid
+      else
+        val
+      end
+    end
+
+    def [](key)
+      if Thread.current[:lazy_access_read_raw]
+        super
+      else
+        key = key.to_s if key.is_a?(Symbol)
+        unless @access_cache.key?(key)
+          val = super
+          @access_cache[key] = LazyAccess.load(val)
+        end
+        @access_cache[key]
+      end
+    end
+
+    def []=(key, value)
+      key = key.to_s if key.is_a?(Symbol)
+      @access_cache[key] = value
+      super
+    end
+  end
+
+  class LazyAccessHash < HashWithIndifferentAccess
+    include LazyAccess
+  end
+
+  class LazyAccessArray < Array
+    include LazyAccess
+  end
+end

data/lib/joblin/uniqueness/compat/active_job.rb

@@ -0,0 +1,75 @@
+
+module Joblin::Uniqueness
+  module Compat
+    module ActiveJob
+
+      class ActiveJobLockContext < LockContext
+        def job_scheduled_at
+          job_instance&.scheduled_at
+        end
+
+        def reenqueue(schedule_in:)
+          job_class.set(
+            queue: job_queue.to_sym,
+            wait: schedule_in,
+            priortity: job_instance.priority,
+          ).perform_later(*job_instance.arguments)
+        end
+      end
+
+      module UniqueJobExtension
+        extend ActiveSupport::Concern
+
+        included do
+          set_callback(:enqueue, :around, prepend: true) do |job, block|
+            ctx = uniqueness_lock_context
+            @uniqueness_cache_data = ctx.cache_data
+            ctx.handle_lifecycle!(:enqueue, &block)
+          end
+
+          around_perform do |job, block|
+            ctx = uniqueness_lock_context
+            ctx.handle_lifecycle!(:perform, &block)
+          end
+        end
+
+        def serialize
+          super.tap do |data|
+            data['uniqueness_cache_data'] = @uniqueness_cache_data.stringify_keys
+          end
+        end
+
+        def deserialize(data)
+          super
+          @uniqueness_cache_data = data['uniqueness_cache_data']&.symbolize_keys
+        end
+
+        def uniqueness_lock_context
+          ActiveJobLockContext.new({
+            job_clazz: self.class,
+            jid: self.job_id,
+            args: self.arguments,
+            queue: self.queue_name,
+            **(@uniqueness_cache_data || {})
+          }, job_instance: self)
+        end
+      end
+
+      module JobExtension
+        extend ActiveSupport::Concern
+        include UniqueJobCommon
+
+        class_methods do
+          def ensure_uniqueness(**kwargs)
+            super(**kwargs)
+            include UniqueJobExtension
+          end
+        end
+      end
+
+      def self.configure
+        ::ActiveJob::Base.include JobExtension
+      end
+    end
+  end
+end

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

@@ -0,0 +1,135 @@
+
+module Joblin::Uniqueness
+  module Compat
+    module Sidekiq
+      module WorkerExtension
+        extend ActiveSupport::Concern
+        include UniqueJobCommon
+
+        class_methods do
+          def ensure_uniqueness(**kwargs)
+            super(**kwargs)
+            if !(defined?(@@validated_config) && @@validated_config)
+              Compat::Sidekiq.validate_middleware_placement!()
+              @@validated_config = true
+            end
+          end
+        end
+      end
+
+      class SidekiqLockContext < LockContext
+        def job_scheduled_at
+          @job_instance&.[]("at")
+        end
+
+        def reenqueue(schedule_in:)
+          job_class.set(queue: job_queue.to_sym).perform_in(schedule_in, *@job_instance["args"])
+        end
+      end
+
+      class CommonMiddleware
+        def lock_context(msg)
+          opts = worker_uniqueness(msg)
+          return nil unless opts
+
+          SidekiqLockContext.new({
+            job_clazz: msg['class'],
+            jid: msg['jid'],
+            queue: msg['queue'],
+            args: msg['args'],
+            # kwargs: msg['kwargs'],
+            **(msg['uniqueness_cache_data']&.symbolize_keys || {}),
+          }, job_instance: msg)
+        end
+
+        def worker_uniqueness(msg)
+          return nil if Compat::Sidekiq.is_activejob_job?(msg)
+
+          worker_class = msg['class'].constantize
+          return nil unless worker_class.respond_to?(:unique_job_options)
+
+          worker_class.unique_job_options
+        end
+      end
+
+      class ClientMiddleware < CommonMiddleware
+        include ::Sidekiq::ClientMiddleware if defined? ::Sidekiq::ClientMiddleware
+
+        def call(_worker, msg, _queue, _redis_pool = nil, &blk)
+          ctx = lock_context(msg)
+          return blk.call unless ctx
+          msg['uniqueness_cache_data'] = ctx.cache_data.stringify_keys
+          ctx.handle_lifecycle!(:enqueue, &blk)
+        end
+      end
+
+      class ServerMiddleware < CommonMiddleware
+        include ::Sidekiq::ServerMiddleware if defined? ::Sidekiq::ServerMiddleware
+
+        def call(_worker, msg, _queue, &blk)
+          ctx = lock_context(msg)
+          return blk.call unless ctx
+          ctx.handle_lifecycle!(:perform, &blk)
+        end
+      end
+
+      def self.is_activejob_job?(msg)
+        return false unless defined?(::ActiveJob)
+
+        msg['class'] == 'ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper' && (msg['wrapped'].to_s).constantize < Compat::ActiveJob::UniqueJobExtension
+      end
+
+      def self.validate_middleware_order(chain, order)
+        chain_classes = chain.entries.map(&:klass)
+        filtered = chain_classes.select { |klass| order.include?(klass) }
+        raise "Middleware chain does not contain all required middleware: #{order - filtered}" unless order.all? { |klass| filtered.include?(klass) }
+        raise "Middleware must be in order: #{order.inspect}" if filtered != order
+      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.validate_middleware_placement!
+        sidekiq_middleware(:client) do |chain|
+          # Unique middleware must come _before_ the Batch middleware so that the uniqueness middleware can wrap the job in a batch
+          validate_middleware_order(chain, [
+            Joblin::Uniqueness::Compat::Sidekiq::ClientMiddleware,
+            Joblin::Batching::Compat::Sidekiq::ClientMiddleware,
+          ])
+        end
+
+        sidekiq_middleware(:server) do |chain|
+          # Unique middleware must com _after_ the Batch middleware so that the Batch is loaded before reaching the uniqueness middleware
+          validate_middleware_order(chain, [
+            Joblin::Batching::Compat::Sidekiq::ServerMiddleware,
+            Joblin::Uniqueness::Compat::Sidekiq::ServerMiddleware,
+          ])
+        end
+      end
+
+      def self.configure
+        sidekiq_middleware(:client) do |chain|
+          chain.insert_before Joblin::Batching::Compat::Sidekiq::ClientMiddleware, Compat::Sidekiq::ClientMiddleware
+        end
+
+        sidekiq_middleware(:server) do |chain|
+          chain.insert_after Joblin::Batching::Compat::Sidekiq::ServerMiddleware, Compat::Sidekiq::ServerMiddleware
+        end
+
+        ::Sidekiq::Worker.extend(ActiveSupport::Concern) unless ::Sidekiq::Worker < ActiveSupport::Concern
+
+        ::Sidekiq::Worker.send(:include, Compat::Sidekiq::WorkerExtension)
+      end
+    end
+  end
+end
+
+# require_relative 'sidekiq/web'