wurk 0.0.1

data/lib/wurk/batch.rb

@@ -0,0 +1,351 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'securerandom'
+require_relative 'lua'
+require_relative 'batch/buffer'
+
+module Wurk
+  # Sidekiq Pro Batches. Group jobs, attach success/complete/death callbacks,
+  # track progress. Spec: docs/target/sidekiq-pro.md §2.
+  #
+  # Lifecycle:
+  #   1. `Batch.new` allocates a fresh BID; the batch is `mutable?` until the
+  #      first `#jobs` block flushes — that flush HSETs the core hash, ZADDs
+  #      to `batches`, and writes tag indexes.
+  #   2. `Batch.new(bid)` reopens an existing batch (legal from inside a job
+  #      or callback only). `mutable?` is false because reopening implies the
+  #      first flush already happened.
+  #   3. `#jobs { ... }` collects `Job.perform_async` calls via the client
+  #      middleware (Thread.current[:wurk_current_batch] is the signal),
+  #      atomically registering each via BATCH_PUSH.
+  #   4. Workers ack on success → BATCH_ACK_SUCCESS → pending--.
+  #      Death handler acks on permanent failure → BATCH_ACK_COMPLETE.
+  #   5. When live jids hits zero → fire `:complete`. When pending also
+  #      hits zero with zero deaths → fire `:success`.
+  #
+  # Nested batches: a job opening its OWN batch (`batch.jobs { ... }`)
+  # increments live counters on the existing batch. A callback opening its
+  # PARENT batch links via parent_bid and adds child BID to b-<bid>-kids.
+  class Batch
+    DEFAULT_EXPIRY_SECONDS = 30 * 24 * 60 * 60
+    POST_SUCCESS_EXPIRY_SECONDS = 24 * 60 * 60
+    CALLBACK_NOTIFY_TTL = 30 * 24 * 60 * 60
+
+    # Bid is URL-safe base64 of 10 random bytes — matches Sidekiq Pro's BID
+    # generator. Length matters: third-party gems that key off bid prefix
+    # (sharded batches in Pro 8) inspect the first character.
+    BID_BYTES = 10
+
+    VALID_EVENTS = %i[success complete death].freeze
+
+    # The 'live' set tracks jobs that have not yet reached a terminal state.
+    # When it's empty, every job has either succeeded or died → `:complete`
+    # is allowed to fire.
+    KEY_SUFFIXES = %w[jids failed died notify cbsucc kids pkids tags].freeze
+
+    THREAD_KEY = :wurk_current_batch
+
+    # Set on the current thread (to a Buffer) only inside an autoflush
+    # `#jobs` block. Client#raw_push reads it: when present, batched pushes
+    # accumulate here instead of round-tripping per job.
+    BUFFER_KEY = :wurk_batch_buffer
+
+    attr_reader :bid, :parent_bid, :linger
+    attr_accessor :description, :callback_queue, :callback_class, :autoflush
+
+    def self.keys_for(bid)
+      base = "b-#{bid}"
+      [base, *KEY_SUFFIXES.map { |s| "#{base}-#{s}" }]
+    end
+
+    def initialize(bid = nil)
+      @bid              = bid || SecureRandom.urlsafe_base64(BID_BYTES)
+      @existing         = !bid.nil?
+      @description      = nil
+      @callback_queue   = 'default'
+      @callback_class   = nil
+      @tags             = []
+      @autoflush        = nil
+      @linger           = nil
+      @parent_bid       = nil
+      @callbacks        = []
+      @expires_in       = DEFAULT_EXPIRY_SECONDS
+      @mutable          = !@existing
+      @flushed_once     = @existing
+      load_existing! if @existing
+    end
+
+    # Hash assignment writes strings — Sidekiq's UI / third-party gems
+    # expect String tags. Array-coercion lets callers pass a String or Set.
+    def tags=(value)
+      @tags = Array(value).map(&:to_s)
+    end
+
+    def tags
+      @tags.dup
+    end
+
+    # Per-batch post-success retention override (seconds). nil falls back to
+    # POST_SUCCESS_EXPIRY_SECONDS when `:success` fires. See §2.8.
+    #
+    # After the first flush the value is persisted to `b-<bid>`; `apply_linger`
+    # reads from Redis, so a setter that only touched memory would silently
+    # ignore the override for any batch reopened by bid.
+    def linger=(duration)
+      @linger = duration&.to_i
+      return unless @flushed_once
+
+      Wurk.redis { |conn| conn.call('HSET', "b-#{@bid}", 'linger', @linger.to_s) }
+    end
+
+    def parent
+      return nil if @parent_bid.nil? || @parent_bid.empty?
+
+      Batch.new(@parent_bid)
+    end
+
+    def mutable?
+      @mutable
+    end
+
+    def include?(jid)
+      Wurk.redis { |conn| conn.call('SISMEMBER', "b-#{@bid}-jids", jid) }.to_i.positive?
+    end
+
+    # Remove jobs from the batch. Decrements pending/total by exactly the
+    # count of jids actually removed (idempotent for repeated calls).
+    def remove_jobs(*jids)
+      return 0 if jids.empty?
+
+      Wurk.redis do |conn|
+        removed = conn.call('SREM', "b-#{@bid}-jids", *jids).to_i
+        if removed.positive?
+          conn.call('HINCRBY', "b-#{@bid}", 'pending', -removed)
+          conn.call('HINCRBY', "b-#{@bid}", 'total', -removed)
+        end
+        removed
+      end
+    end
+
+    # Mark batch invalid. Pending jobs still exist in their queues; the
+    # server middleware short-circuits them when it observes the flag.
+    # Cascades to descendant batches via b-<bid>-kids.
+    def invalidate_all
+      cascade_invalidate(@bid)
+      nil
+    end
+
+    def valid?
+      Wurk.redis { |conn| conn.call('HGET', "b-#{@bid}", 'invalidated') } != '1'
+    end
+
+    def status
+      Status.new(@bid)
+    end
+
+    def expires_in(duration)
+      @expires_in = duration.to_i
+      self
+    end
+
+    # Register a callback. Multiple callbacks of the same event are allowed.
+    # The callback target may be a Class, "Foo#bar" string spec, or anything
+    # responding to `name`. `options` must be JSON-serializable.
+    def on(event, callback, options = {})
+      sym = event.to_sym
+      raise ArgumentError, "invalid event #{event.inspect}" unless VALID_EVENTS.include?(sym)
+      raise ArgumentError, 'callback options must be a Hash' unless options.is_a?(Hash)
+
+      @callbacks << [sym.to_s, callback_target(callback), options]
+      self
+    end
+
+    # Atomic enqueue block. Inside the block, `Job.perform_async` finds
+    # this batch via Thread.current[THREAD_KEY] and stamps `bid` onto the
+    # payload — the client middleware then uses BATCH_PUSH to register and
+    # push atomically. Empty blocks synthesise a Batch::Empty no-op so
+    # callbacks still fire.
+    def jobs(&block)
+      raise ArgumentError, 'jobs requires a block' unless block
+
+      ensure_first_flush!
+      pre_count = job_count
+      collect_jobs(&block)
+      # By the time we check, the buffer (if any) has flushed, so `total`
+      # reflects everything the block pushed — a flat count is reliable.
+      enqueue_empty_marker if job_count == pre_count
+      @mutable = false
+      self
+    end
+
+    private
+
+    # Runs the block with this batch active so the client middleware stamps
+    # the bid. With autoflush on, batched pushes accumulate in a Buffer and
+    # flush once at exit; the per-N flushing happens in Client#raw_push.
+    def collect_jobs
+      previous    = Thread.current[THREAD_KEY]
+      prev_buffer = Thread.current[BUFFER_KEY]
+      buffer      = new_buffer
+      Thread.current[THREAD_KEY] = self
+      Thread.current[BUFFER_KEY] = buffer
+      begin
+        yield
+        flush_buffer(buffer) if buffer
+      ensure
+        Thread.current[THREAD_KEY] = previous
+        Thread.current[BUFFER_KEY] = prev_buffer
+      end
+    end
+
+    # Buffering is opt-in via `autoflush`: `true` buffers the whole block and
+    # flushes once at exit; a positive Integer flushes every N jobs. Anything
+    # falsy keeps the default per-job immediate push.
+    def new_buffer
+      return nil unless @autoflush
+
+      Buffer.new([], autoflush_threshold)
+    end
+
+    # `true` → buffer the whole block (nil threshold, drained at exit);
+    # positive Integer → flush every N. Any other truthy value is a config
+    # typo (`0`, `-1`, `"5"`) — fail fast instead of silently degrading to
+    # "flush at block exit".
+    def autoflush_threshold
+      return nil if @autoflush == true
+      return @autoflush if @autoflush.is_a?(Integer) && @autoflush.positive?
+
+      raise ArgumentError, "autoflush must be true or a positive Integer, got #{@autoflush.inspect}"
+    end
+
+    def flush_buffer(buffer)
+      payloads = buffer.drain
+      Wurk::Client.new.flush_batched(payloads) unless payloads.empty?
+    end
+
+    # First flush writes the core hash, registers in the global `batches`
+    # zset, and links tag indexes. Subsequent #jobs invocations skip this
+    # — `total` is already there and BATCH_PUSH only increments deltas.
+    def ensure_first_flush!
+      return if @flushed_once
+
+      now = ::Process.clock_gettime(::Process::CLOCK_REALTIME)
+      Wurk.redis { |conn| conn.pipelined { |pipe| pipelined_first_flush(pipe, now) } }
+      @parent_bid = current_parent_bid
+      @flushed_once = true
+    end
+
+    def pipelined_first_flush(pipe, now)
+      pipe.call('HSET', "b-#{@bid}", *first_flush_hash(now).flatten)
+      pipe.call('EXPIRE', "b-#{@bid}", @expires_in)
+      pipe.call('ZADD', 'batches', now.to_s, @bid)
+      @tags.each { |t| pipe.call('SADD', "tags:#{t}", @bid) }
+      link_to_parent(pipe) if current_parent_bid
+    end
+
+    def first_flush_hash(now)
+      {
+        'created_at' => now.to_s,
+        'description' => @description.to_s,
+        'callback_queue' => @callback_queue.to_s,
+        'callback_class' => @callback_class.to_s,
+        'parent_bid' => current_parent_bid.to_s,
+        'tags' => @tags.to_json,
+        'linger' => @linger.to_s,
+        'callbacks' => @callbacks.to_json,
+        'total' => '0',
+        'pending' => '0',
+        'failures' => '0'
+      }
+    end
+
+    # If we're flushing from inside another batch's #jobs block, that
+    # parent batch's BID becomes our parent_bid and we get registered into
+    # its kids/pkids sets so cascade success/failure is correctly tracked.
+    def current_parent_bid
+      outer = Thread.current[THREAD_KEY]
+      return nil if outer.nil? || outer.bid == @bid
+
+      outer.bid
+    end
+
+    def link_to_parent(pipe)
+      pipe.call('SADD', "b-#{current_parent_bid}-kids", @bid)
+      pipe.call('SADD', "b-#{current_parent_bid}-pkids", @bid)
+    end
+
+    def job_count
+      Wurk.redis { |conn| conn.call('HGET', "b-#{@bid}", 'total') }.to_i
+    end
+
+    def enqueue_empty_marker
+      require_relative 'batch/empty'
+      previous = Thread.current[THREAD_KEY]
+      Thread.current[THREAD_KEY] = self
+      begin
+        Wurk::Batch::Empty.perform_async
+      ensure
+        Thread.current[THREAD_KEY] = previous
+      end
+    end
+
+    def cascade_invalidate(bid)
+      Wurk.redis do |conn|
+        Wurk::Lua::Loader.eval_cached(conn, :batch_invalidate, keys: ["b-#{bid}", "b-#{bid}-jids"], argv: [])
+        kids = conn.call('SMEMBERS', "b-#{bid}-kids") || []
+        kids.each { |child| cascade_invalidate(child) }
+      end
+    end
+
+    def load_existing!
+      data = fetch_hash
+      load_meta(data)
+      @tags      = parse_json_array(data['tags'])
+      @linger    = nil_if_empty(data['linger'])&.to_i
+      @callbacks = parse_json_callbacks(data['callbacks'])
+    end
+
+    def load_meta(data)
+      @description    = nil_if_empty(data['description'])
+      @callback_queue = nil_if_empty(data['callback_queue']) || 'default'
+      @callback_class = nil_if_empty(data['callback_class'])
+      @parent_bid     = nil_if_empty(data['parent_bid'])
+    end
+
+    def fetch_hash
+      raw = Wurk.redis { |conn| conn.call('HGETALL', "b-#{@bid}") }
+      raw.is_a?(Hash) ? raw : raw.each_slice(2).to_h
+    end
+
+    def callback_target(callback)
+      case callback
+      when Class then callback.name
+      when String, Symbol then callback.to_s
+      else
+        raise ArgumentError, "callback must be Class or String: #{callback.inspect}" unless callback.respond_to?(:name)
+
+        callback.name
+      end
+    end
+
+    def nil_if_empty(val)
+      val.nil? || val.to_s.empty? ? nil : val
+    end
+
+    def parse_json_array(raw)
+      return [] if raw.nil? || raw.empty?
+
+      JSON.parse(raw)
+    rescue JSON::ParserError
+      []
+    end
+
+    def parse_json_callbacks(raw)
+      arr = parse_json_array(raw)
+      arr.is_a?(Array) ? arr : []
+    end
+  end
+end
+
+require_relative 'batch/status'

data/lib/wurk/batch_set.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require_relative 'batch'
+
+module Wurk
+  # Discovery API over the global `batches` sorted set. `size`, `each` (yields
+  # Status), `scan_tags(tag)` for tag-indexed lookup.
+  #
+  # Spec: docs/target/sidekiq-pro.md §2.7.
+  class BatchSet
+    include Enumerable
+
+    PAGE_SIZE = 100
+
+    def initialize(key: 'batches')
+      @key = key
+    end
+
+    def size
+      Wurk.redis { |conn| conn.call('ZCARD', @key) }.to_i
+    end
+
+    # Newest-first iteration of every batch indexed in the set. Yields
+    # `Wurk::Batch::Status` instances — they HGETALL on construction so
+    # the caller pays one Redis round-trip per batch.
+    def each
+      return enum_for(:each) unless block_given?
+
+      page = 0
+      loop do
+        start = page * PAGE_SIZE
+        stop  = start + PAGE_SIZE - 1
+        bids  = Wurk.redis { |conn| conn.call('ZRANGE', @key, start, stop, 'REV') }
+        bids.each { |bid| yield Wurk::Batch::Status.new(bid) }
+        break if bids.size < PAGE_SIZE
+
+        page += 1
+      end
+    end
+
+    # Yields each BID indexed under the given tag (set at `tags:<tag>`).
+    # No JSON parsing; cheap iteration. Caller wraps in `Status.new(bid)`
+    # if it needs full state.
+    def scan_tags(tag, &block)
+      return enum_for(:scan_tags, tag) unless block_given?
+
+      cursor = '0'
+      Wurk.redis do |conn|
+        loop do
+          cursor, members = conn.call('SSCAN', "tags:#{tag}", cursor, 'COUNT', PAGE_SIZE)
+          members.each(&block)
+          break if cursor == '0'
+        end
+      end
+    end
+  end
+
+  class Batch
+    # Discovery view over batches that triggered `:death`. Iterates Status
+    # instances newest-first.
+    class DeadSet < ::Wurk::BatchSet
+      def initialize
+        super(key: 'dead-batches')
+      end
+    end
+  end
+end

data/lib/wurk/capsule.rb

@@ -0,0 +1,176 @@
+# frozen_string_literal: true
+
+require_relative 'redis_pool'
+require_relative 'middleware/chain'
+require_relative 'fetcher/reliable'
+
+module Wurk
+  # One processing unit: a set of threads + queues sharing a fetcher and a
+  # Redis pool. Configurations can hold many capsules; each maps to its own
+  # Manager and Processors.
+  #
+  # Spec: docs/target/sidekiq-free.md §5 (Sidekiq::Capsule).
+  class Capsule
+    MODES = %i[strict weighted random].freeze
+
+    attr_reader :name, :queues, :mode, :weights, :config
+    attr_accessor :concurrency, :fetcher
+
+    def initialize(name, config)
+      @name = name.to_s
+      @config = config
+      @concurrency = config[:concurrency] || 5
+      @queues = ['default']
+      @mode = :strict
+      @weights = { 'default' => 0 }
+      @fetcher = nil
+      @redis_pool = nil
+      @local_redis_pool = nil
+      @client_chain = nil
+      @server_chain = nil
+    end
+
+    def to_h
+      { concurrency: @concurrency, mode: @mode, weights: @weights }
+    end
+
+    # Parses queue specs Sidekiq-style:
+    #   %w[high default low]        → mode :strict,   weights all 0
+    #   %w[high,3 default,2 low,1]  → mode :weighted, weights {q=>w}
+    #   %w[a,1 b,1 c,1]             → mode :random,   weights all 1
+    # @queues is expanded by weight so a uniform shuffle gives weighted
+    # fairness (e.g. ["high","high","high","default","default","low"]).
+    def queues=(val)
+      parsed = Array(val).map { |entry| parse_queue_entry(entry) }
+      raise ArgumentError, 'queues cannot be empty' if parsed.empty?
+
+      @weights = parsed.to_h
+      @mode = detect_mode(parsed)
+      @queues = expand_by_weight(parsed, @mode)
+    end
+
+    # Lossless `name[,weight]` specs (unlike `queues`, which is the
+    # weight-expanded list). Lets Configuration#topology rebuild a slot that
+    # round-trips back through `queues=` without flattening weights.
+    def queue_specs
+      @weights.map { |q, w| w.positive? ? "#{q},#{w}" : q }
+    end
+
+    # Materialize everything that lazy-inits via `||=` and default the fetcher,
+    # BEFORE Configuration#freeze! freezes the capsule — otherwise the first
+    # post-freeze access (a fetch tick, a middleware call) hits a nil fetcher
+    # or FrozenErrors building a pool. The swarm's ChildBoot used to do this
+    # by hand; centralizing it here covers the standalone CLI and embedded
+    # paths too (the bug behind a nil `fetcher` in `exe/wurk`). Idempotent.
+    def prepare!
+      @fetcher ||= Wurk::Fetcher::Reliable.new(self)
+      redis_pool
+      local_redis_pool
+      client_middleware
+      server_middleware
+      self
+    end
+
+    def client_middleware
+      chain = (@client_chain ||= @config.client_middleware.copy_for(self))
+      yield chain if block_given?
+      chain
+    end
+
+    def server_middleware
+      # copy_for(self) — not dup — binds the chain's `@config` to this capsule,
+      # so middleware that reach for `redis_pool`/`redis`/`logger` resolve them
+      # instead of hitting `nil` (a plain dup leaves @config nil).
+      chain = (@server_chain ||= @config.server_middleware.copy_for(self))
+      yield chain if block_given?
+      chain
+    end
+
+    # Pool size = concurrency + POOL_OVERHEAD. Every processor thread can
+    # be parked in a blocking BLMOVE (reliable fetch), holding its slot
+    # for ~3s. POOL_OVERHEAD reserves connections for the Launcher's
+    # heartbeat thread + the Scheduled::Poller so they can never be
+    # starved by busy fetchers. Without this, concurrency=1 deadlocks
+    # immediately (worker holds the only slot, heartbeat times out).
+    POOL_OVERHEAD = 2
+
+    def redis_pool
+      @redis_pool ||= build_pool(size: @concurrency + POOL_OVERHEAD, name: "#{@name}-main")
+    end
+
+    def local_redis_pool
+      @local_redis_pool ||= build_pool(size: @concurrency, name: "#{@name}-local")
+    end
+
+    # Disconnect and drop cached pools. Called by Wurk::Swarm just before
+    # fork (parent side: close inherited sockets) and just after fork
+    # (child side: rebuild lazily). Connection_pool#shutdown is terminal,
+    # so dropping the reference is required — `redis_pool` will rebuild.
+    def reset_redis_pools!
+      @redis_pool&.disconnect!
+      @redis_pool = nil
+      @local_redis_pool&.disconnect!
+      @local_redis_pool = nil
+    end
+
+    def redis(&)
+      redis_pool.with(&)
+    end
+
+    def lookup(name)
+      @config.lookup(name)
+    end
+
+    def logger
+      @config.logger
+    end
+
+    # No-op in OSS. Reserved for Pro/Ent hooks that need to flush per-capsule
+    # state on shutdown. Manager#stop invokes this in `ensure` so the contract
+    # is honored regardless of how stop unwinds.
+    #
+    # Spec: docs/target/sidekiq-free.md §5.
+    def stop; end
+
+    private
+
+    def parse_queue_entry(entry)
+      qname, weight_str = entry.to_s.split(',', 2)
+      qname = qname.to_s.strip
+      raise ArgumentError, "queue name cannot be empty: `#{entry}`" if qname.empty?
+      return [qname, 0] if weight_str.nil?
+
+      weight = Integer(weight_str)
+      raise ArgumentError, "queue weight must be > 0: `#{entry}`" if weight <= 0
+
+      [qname, weight]
+    end
+
+    def detect_mode(parsed)
+      weights = parsed.map(&:last)
+      if weights.all?(&:zero?)
+        :strict
+      elsif weights.uniq.size == 1
+        :random
+      else
+        :weighted
+      end
+    end
+
+    def expand_by_weight(parsed, mode)
+      return parsed.map(&:first) if %i[strict random].include?(mode)
+
+      parsed.flat_map { |q, w| [q] * w }
+    end
+
+    def build_pool(size:, name:)
+      cfg = @config.redis_config
+      RedisPool.new(
+        size: size,
+        url: cfg[:url] || RedisPool::DEFAULT_URL,
+        timeout: cfg[:timeout] || RedisPool::DEFAULT_TIMEOUT,
+        name: name
+      )
+    end
+  end
+end