wurk 0.0.1

data/lib/wurk/swarm.rb

@@ -0,0 +1,260 @@
+# frozen_string_literal: true
+
+require_relative 'component'
+require_relative 'launcher'
+require_relative 'fetcher/reliable'
+require_relative 'keys'
+require_relative 'swarm/child_boot'
+
+module Wurk
+  # Parent supervisor. Forks N children per the worker topology, monitors
+  # PIDs, relays signals, respawns crashed children, handles rolling
+  # restart on SIGUSR1, recycles RSS-bloated children.
+  #
+  # Boot ordering (must be exact — see docs/idea/03-process-model.md):
+  #   1. Host app boots fully; eager loads done.
+  #   2. Railtie `after_initialize` fires.
+  #   3. `boot` closes parent-side connections (Redis, ActiveRecord).
+  #   4. `boot` forks N children.
+  #   5. Each child reconnects DB + opens a fresh Redis pool, then
+  #      installs its own signal handlers and starts the Launcher.
+  #   6. Parent calls `supervise` to enter the wait/relay loop.
+  #
+  # Signals (see docs/idea/04-signals.md):
+  #   TERM/INT  → `shutdown`           (graceful drain)
+  #   TSTP      → relay TSTP           (pause fetch)
+  #   CONT      → relay CONT           (resume fetch)
+  #   USR1      → `rolling_restart`    (zero-downtime cycle)
+  class Swarm
+    include Component
+
+    SUPERVISE_TICK = 0.2
+    RESPAWN_BACKOFF = 1.0
+    HEARTBEAT_WAIT = 30
+    MEMORY_CHECK_INTERVAL = 10
+    DEFAULT_SHUTDOWN_TIMEOUT = 25
+
+    attr_reader :topology, :children
+
+    def initialize(topology:, config: Wurk.configuration, memory_limit: nil,
+                   shutdown_timeout: DEFAULT_SHUTDOWN_TIMEOUT)
+      @topology = topology
+      @config = config
+      @memory_limit = memory_limit
+      @shutdown_timeout = shutdown_timeout
+      @children = {}
+      @assignments = []
+      @stopping = false
+      @last_memory_check = 0
+      @signal_queue = ::Thread::Queue.new
+    end
+
+    # `install_signals:` is false in tests so the integration suite can
+    # drive `shutdown` / `rolling_restart` directly without poisoning the
+    # test process's signal handlers.
+    def boot(install_signals: true)
+      raise 'Wurk::Swarm already booted' unless @assignments.empty?
+      raise ArgumentError, 'Topology has no slots' if @topology.empty?
+
+      @assignments = @topology.assignments.freeze
+      close_parent_sockets
+      fork_children
+      install_signal_handlers if install_signals
+      @children.keys
+    end
+
+    def supervise
+      until done?
+        drain_signals
+        reap_one_child
+        check_memory_pressure
+        sleep SUPERVISE_TICK
+      end
+    end
+
+    def shutdown(timeout: @shutdown_timeout)
+      @stopping = true
+      relay_signal('TERM')
+      wait_for_children(timeout)
+      hard_kill_stragglers
+    end
+
+    # SIGUSR1. For each existing child, fork a replacement, wait for its
+    # first heartbeat, then TERM + drain the old one. Long-running jobs
+    # in the old slot get the full shutdown_timeout while the replacement
+    # is already serving new work.
+    def rolling_restart
+      @children.dup.each do |old_pid, meta|
+        replacement = fork_child(meta[:slot], meta[:index])
+        @children[replacement] = meta
+        unless wait_for_heartbeat(replacement)
+          logger.warn do
+            "swarm: replacement #{replacement} heartbeat not seen within #{HEARTBEAT_WAIT}s; proceeding anyway"
+          end
+        end
+        safe_kill(old_pid, 'TERM')
+        wait_pid(old_pid, @shutdown_timeout)
+        @children.delete(old_pid)
+      end
+    end
+
+    private
+
+    # Step 3.
+    def close_parent_sockets
+      @config.reset_redis_pools!
+      close_active_record_pool
+    end
+
+    def close_active_record_pool
+      return unless defined?(::ActiveRecord::Base)
+
+      ::ActiveRecord::Base.connection_handler.clear_active_connections!
+      ::ActiveRecord::Base.connection_handler.flush_idle_connections!
+    rescue StandardError => e
+      logger.warn { "swarm: ActiveRecord close failed: #{e.class}: #{e.message}" }
+    end
+
+    # Step 4.
+    def fork_children
+      @assignments.each_with_index do |slot, idx|
+        @children[fork_child(slot, idx)] = { slot: slot, index: idx }
+      end
+    end
+
+    def fork_child(slot, idx)
+      pid = ::Process.fork
+      return pid if pid
+
+      ChildBoot.new(@config, slot, idx).run
+      exit 0 # unreachable; ChildBoot exits explicitly
+    end
+
+    def install_signal_handlers
+      { 'TERM' => :term, 'INT' => :term, 'TSTP' => :tstp,
+        'CONT' => :cont, 'USR1' => :usr1 }.each do |sig, sym|
+        ::Signal.trap(sig) { @signal_queue << sym }
+      end
+    end
+
+    def drain_signals
+      until @signal_queue.empty?
+        sym = next_signal_symbol
+        next if sym.nil?
+
+        case sym
+        when :term then shutdown
+        when :tstp then relay_signal('TSTP')
+        when :cont then relay_signal('CONT')
+        when :usr1 then rolling_restart
+        end
+      end
+    end
+
+    def next_signal_symbol
+      @signal_queue.pop(true)
+    rescue ThreadError
+      nil
+    end
+
+    def reap_one_child
+      pid, status = ::Process.wait2(-1, ::Process::WNOHANG)
+      on_child_exit(pid, status) if pid
+    rescue Errno::ECHILD
+      @stopping = true
+    end
+
+    def on_child_exit(pid, status)
+      meta = @children.delete(pid)
+      return unless meta
+
+      if @stopping
+        logger.info { "swarm: child #{pid} exited (status=#{status.exitstatus})" }
+      else
+        logger.warn { "swarm: child #{pid} died (status=#{status.exitstatus}); respawning slot #{meta[:index]}" }
+        sleep RESPAWN_BACKOFF
+        @children[fork_child(meta[:slot], meta[:index])] = meta
+      end
+    end
+
+    def check_memory_pressure
+      return unless @memory_limit
+
+      now = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+      return if now - @last_memory_check < MEMORY_CHECK_INTERVAL
+
+      @last_memory_check = now
+      @children.dup.each_key { |pid| recycle_if_bloated(pid) }
+    end
+
+    def recycle_if_bloated(pid)
+      rss = pid_rss_kb(pid)
+      return if rss.nil? || rss < @memory_limit
+
+      logger.warn { "swarm: child #{pid} RSS #{rss}KB >= #{@memory_limit}KB; recycling" }
+      safe_kill(pid, 'TERM')
+    end
+
+    def pid_rss_kb(pid)
+      return nil unless ::File.exist?("/proc/#{pid}/statm")
+
+      ::File.read("/proc/#{pid}/statm").split[1].to_i * 4
+    rescue StandardError
+      nil
+    end
+
+    def relay_signal(sig)
+      @children.each_key { |pid| safe_kill(pid, sig) }
+    end
+
+    def safe_kill(pid, sig)
+      ::Process.kill(sig, pid)
+    rescue Errno::ESRCH, Errno::EPERM
+      nil
+    end
+
+    def wait_pid(pid, timeout)
+      deadline = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) + timeout
+      while ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) < deadline
+        return true if ::Process.wait(pid, ::Process::WNOHANG)
+
+        sleep 0.1
+      end
+      false
+    rescue Errno::ECHILD
+      true
+    end
+
+    def wait_for_children(timeout)
+      deadline = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) + timeout
+      while ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) < deadline && @children.any?
+        reap_one_child
+        sleep 0.1
+      end
+    end
+
+    def hard_kill_stragglers
+      @children.each_key { |pid| safe_kill(pid, 'KILL') }
+      @children.clear
+    end
+
+    # Identity is `<hostname>:<pid>:<nonce>`. PROCESS_NONCE is set when
+    # Component loads in the parent and inherited by every fork — the
+    # parent can compute a child's identity from its PID alone.
+    # Returns true if the heartbeat was observed before the deadline.
+    def wait_for_heartbeat(pid) # rubocop:disable Naming/PredicateMethod
+      identity = "#{hostname}:#{pid}:#{Component::PROCESS_NONCE}"
+      deadline = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) + HEARTBEAT_WAIT
+      while ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) < deadline
+        return true if @config.redis { |c| c.call('SISMEMBER', Keys::PROCESSES, identity) } == 1
+
+        sleep 0.5
+      end
+      false
+    end
+
+    def done?
+      @stopping && @children.empty?
+    end
+  end
+end

data/lib/wurk/testing.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require_relative 'queues'
+require_relative 'middleware/chain'
+
+module Wurk
+  # Sidekiq::Testing-compatible test harness (aliased to Sidekiq::Testing).
+  # Three modes control how `Wurk::Client#raw_push` behaves:
+  #
+  #   :disable — real Redis push (the default; production behavior)
+  #   :fake    — payloads collected in the in-memory Wurk::Queues store
+  #   :inline  — jobs executed synchronously the instant they're pushed
+  #
+  # A block form switches the mode for the duration of the block on the current
+  # thread only (`fake! { ... }`); the no-block form sets it process-globally.
+  #
+  # Spec: docs/target/sidekiq-free.md §24.
+  module Testing
+    class TestModeAlreadySetError < ::RuntimeError; end
+    # Raised by `Worker.perform_one` / `drain` when no fake job is available.
+    class EmptyQueueError < ::RuntimeError; end
+
+    THREAD_KEY = :__wurk_testing_mode
+
+    class << self
+      def disable!(&) = __set_test_mode(:disable, &)
+      def fake!(&)    = __set_test_mode(:fake, &)
+      def inline!(&)  = __set_test_mode(:inline, &)
+
+      def disabled? = mode == :disable
+      def enabled?  = !disabled?
+      def fake?     = mode == :fake
+      def inline?   = mode == :inline
+
+      # Thread-local override (set by a block) wins over the global mode, so a
+      # `fake! { ... }` block is isolated to the calling thread.
+      def mode
+        ::Thread.current[THREAD_KEY] || @mode || :disable
+      end
+
+      # Block → thread-local for the block's duration; no block → global.
+      # Nesting block forms (`fake! { inline! { ... } }`) is rejected, matching
+      # Sidekiq 8.
+      def __set_test_mode(new_mode, &block)
+        return @mode = new_mode unless block
+
+        if ::Thread.current[THREAD_KEY]
+          raise TestModeAlreadySetError, 'Nested Sidekiq::Testing block modes are not allowed'
+        end
+
+        ::Thread.current[THREAD_KEY] = new_mode
+        begin
+          block.call
+        ensure
+          ::Thread.current[THREAD_KEY] = nil
+        end
+      end
+
+      # In-process server-middleware chain used for inline execution. Empty by
+      # default — configure with `Sidekiq::Testing.server_middleware { |c| ... }`.
+      def server_middleware
+        @server_middleware ||= ::Wurk::Middleware::Chain.new(::Wurk.configuration)
+        yield @server_middleware if block_given?
+        @server_middleware
+      end
+
+      # --- push hooks invoked by Wurk::Client#raw_push -------------------
+
+      # Route a push through the active test mode (only called when enabled?).
+      def dispatch_push(payloads)
+        inline? ? inline_push(payloads) : fake_push(payloads)
+      end
+
+      # Collect payloads into the in-memory store. `enqueued_at` is stamped now
+      # unless the job is scheduled (`at`), mirroring the real client.
+      def fake_push(payloads)
+        now = ::Process.clock_gettime(::Process::CLOCK_REALTIME, :millisecond)
+        payloads.each do |payload|
+          payload['enqueued_at'] = now unless payload['at']
+          ::Wurk::Queues.push(payload['queue'], payload['class'], payload)
+        end
+        payloads.last['jid']
+      end
+
+      # Execute each payload immediately through the inline server chain.
+      def inline_push(payloads)
+        payloads.each { |payload| ::Object.const_get(payload['class'].to_s).process_job(payload) }
+        payloads.last['jid']
+      end
+
+      # Run every fake job across all classes until the store is empty.
+      def drain_all
+        count = 0
+        while (job = ::Wurk::Queues.shift_any)
+          ::Object.const_get(job['class'].to_s).process_job(job)
+          count += 1
+        end
+        count
+      end
+    end
+  end
+end

data/lib/wurk/topology.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module Wurk
+  # Worker topology DSL (Wurk extension on top of Ent's flat swarm).
+  # Lets users declare specialized slots: e.g. 2 forks dedicated to the
+  # critical queue with low concurrency, 2 forks for bulk + low with high
+  # concurrency. Stronger queue isolation than a flat swarm.
+  #
+  # Each Slot describes a *kind* of fork; `count` is how many identical
+  # forks of that kind to spawn. Swarm consumes `assignments` (the flat
+  # list of forks to spawn, in order) so a slot with count=2 yields two
+  # assignment entries pointing at the same Slot.
+  #
+  # See docs/idea/03-process-model.md §Worker topology.
+  class Topology
+    # `:count` shadows Struct#count by design — Slot is a kw-init data
+    # carrier and the slot's child-count is the field users read.
+    Slot = Struct.new(:count, :queues, :concurrency, keyword_init: true) do # rubocop:disable Lint/StructNewOverride
+      def to_h
+        { count: count, queues: queues, concurrency: concurrency }
+      end
+    end
+
+    def initialize
+      @slots = []
+    end
+
+    # Declare one slot kind. Returns self so calls chain.
+    def slot(count:, queues:, concurrency:)
+      queue_list = validate_slot!(count, queues, concurrency)
+      @slots << Slot.new(count: count, queues: queue_list, concurrency: concurrency).freeze
+      self
+    end
+
+    def slots
+      @slots.dup.freeze
+    end
+
+    def empty?
+      @slots.empty?
+    end
+
+    # Flat ordered list of slots to fork, one per child process. A slot
+    # with count=N contributes N entries.
+    def assignments
+      @slots.flat_map { |s| Array.new(s.count, s) }
+    end
+
+    def total_processes
+      @slots.sum(&:count)
+    end
+
+    # Convenience: build a flat topology of `count` identical forks
+    # consuming `queues` with `concurrency` threads each. Used by the
+    # railtie when the host hasn't declared a custom topology.
+    def self.flat(count:, queues:, concurrency:)
+      new.slot(count: count, queues: queues, concurrency: concurrency)
+    end
+
+    private
+
+    def validate_slot!(count, queues, concurrency)
+      raise ArgumentError, "count must be > 0 (got #{count.inspect})" unless count.is_a?(Integer) && count.positive?
+      unless concurrency.is_a?(Integer) && concurrency.positive?
+        raise ArgumentError, "concurrency must be > 0 (got #{concurrency.inspect})"
+      end
+
+      queue_list = Array(queues).map(&:to_s)
+      raise ArgumentError, 'queues cannot be empty' if queue_list.empty?
+
+      queue_list.freeze
+    end
+  end
+end

data/lib/wurk/unique.rb

@@ -0,0 +1,240 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'digest'
+require_relative 'middleware'
+
+module Wurk
+  # Sidekiq Enterprise unique jobs. Best-effort dedup at enqueue time keyed
+  # by a SHA256 digest of `[class, queue, args]` (overridable via
+  # `sidekiq_unique_context`). Three lock-release strategies:
+  #
+  #   * `unique_until: :success` (default) — lock retained through retries;
+  #     server middleware DELs it on successful perform. Surviving across
+  #     a process crash is bounded by `unique_for` TTL.
+  #   * `unique_until: :start` — server middleware DELs the lock right
+  #     *before* invoking perform; a duplicate can be enqueued while the
+  #     first is running.
+  #
+  # Wire-compat (§3.9): single-key Redis layout — `unique:<sha256>` STRING
+  # holding the owning JID. Scheduled jobs extend the TTL by the delay so
+  # the lock covers the entire wait+execution window (§3.4).
+  #
+  # Spec: docs/target/sidekiq-ent.md §3.
+  module Unique
+    KEY_PREFIX = 'unique:'
+    DEFAULT_UNTIL = :success
+    VALID_UNTIL = %i[success start].freeze
+
+    # `Sidekiq::Enterprise.unique!` flips this on. The middleware pair is
+    # always loaded (so worker `sidekiq_options unique_for:` is a no-op
+    # without `unique!`) — only when the flag is set does the client
+    # middleware actually compute and SETNX the digest.
+    class << self
+      def enabled?
+        @enabled == true
+      end
+
+      def enable! # rubocop:disable Naming/PredicateMethod
+        @enabled = true
+        register_middleware!
+        true
+      end
+
+      # Test helper — not part of the public Sidekiq surface. Clears the
+      # flag so per-test enable!/disable! does not leak across runs.
+      def disable!
+        @enabled = false
+        nil
+      end
+
+      # Compute the lock key for an arbitrary `(queue, klass, args)` triple.
+      # Used by both the client middleware and the public `locked?` probe so
+      # they cannot drift.
+      def lock_key(klass, queue, args)
+        context = [klass.to_s, queue.to_s, args]
+        "#{KEY_PREFIX}#{Digest::SHA256.hexdigest(JSON.dump(context))}"
+      end
+
+      # Compute the lock key from a job payload, honoring
+      # `sidekiq_unique_context` when the worker class is loaded and
+      # defines it.
+      def lock_key_for(job)
+        context = unique_context(job)
+        "#{KEY_PREFIX}#{Digest::SHA256.hexdigest(JSON.dump(context))}"
+      end
+
+      # Default: `[class, queue, args]`. Workers may override by defining
+      # `self.sidekiq_unique_context(job)` returning any JSON-serializable
+      # value (e.g. a subset of args). Spec §3.5.
+      def unique_context(job)
+        klass = resolve_class(job['class'])
+        if klass.respond_to?(:sidekiq_unique_context)
+          klass.sidekiq_unique_context(job)
+        else
+          [job['class'], job['queue'], job['args']]
+        end
+      end
+
+      private
+
+      def resolve_class(name)
+        return nil if name.nil? || name.to_s.empty?
+
+        ::Object.const_get(name.to_s)
+      rescue ::NameError
+        nil
+      end
+
+      def register_middleware!
+        Wurk.configuration.client_middleware.add(ClientMiddleware) \
+          unless Wurk.configuration.client_middleware.exists?(ClientMiddleware)
+        Wurk.configuration.server_middleware.add(ServerMiddleware) \
+          unless Wurk.configuration.server_middleware.exists?(ServerMiddleware)
+      end
+    end
+
+    # Coerce `unique_for` to a numeric seconds value. Accepts Integer,
+    # Numeric, ActiveSupport::Duration (any `to_i`-respondent), or `false`
+    # (skip). Returns nil when uniqueness should be skipped.
+    def self.coerce_ttl(value)
+      return nil if value.nil? || value == false
+      return value if value.is_a?(Integer) && value.positive?
+      return value.to_i if value.is_a?(Numeric)
+      return value.to_i if duration_like?(value)
+
+      nil
+    end
+
+    def self.duration_like?(value)
+      return false unless value.respond_to?(:to_i)
+
+      value.respond_to?(:since) || value.class.name.to_s.include?('Duration')
+    end
+
+    # ------------------------------------------------------------------
+    # Introspection — `Sidekiq::Enterprise::Unique.locked?`
+    # ------------------------------------------------------------------
+
+    # @return [String, nil] owning jid, or nil when the lock is free.
+    def self.locked?(queue_or_klass, klass_or_args = nil, args = nil)
+      queue, klass, payload = normalize_locked_args(queue_or_klass, klass_or_args, args)
+      key = lock_key(klass, queue, payload)
+      Wurk.redis { |c| c.call('GET', key) }
+    end
+
+    # Accepts either `(klass, args)` or `(queue, klass, args)`. Without a
+    # queue the default Wurk job queue is assumed — matches the Sidekiq
+    # Ent docs §3.6.
+    def self.normalize_locked_args(first, second, third)
+      if third.nil?
+        [Wurk.default_job_options['queue'] || 'default', first, Array(second)]
+      else
+        [first.to_s, second, Array(third)]
+      end
+    end
+    private_class_method :normalize_locked_args
+
+    # ------------------------------------------------------------------
+    # Client middleware — SETNX lock at push time.
+    # ------------------------------------------------------------------
+    #
+    # Drops the duplicate by returning nil from the chain (Wurk::Client
+    # treats nil as "halted"; the caller's `perform_async` returns nil
+    # JID). Logs the holder JID for debuggability.
+    class ClientMiddleware
+      include Wurk::Middleware::ClientMiddleware
+
+      def call(_worker, job, _queue, redis_pool, &)
+        return yield unless Wurk::Unique.enabled?
+
+        ttl = effective_ttl(job)
+        return yield if ttl.nil?
+
+        acquire_or_drop(redis_pool, job, Wurk::Unique.lock_key_for(job), ttl, &)
+      end
+
+      private
+
+      # Add `at - now` delay to the base TTL so a scheduled job's lock
+      # spans the wait + execution window (§3.4). Returns nil when the
+      # job opts out (`unique_for: false` / missing).
+      def effective_ttl(job)
+        base = Wurk::Unique.coerce_ttl(job['unique_for'])
+        return nil if base.nil?
+        return base unless job['at']
+
+        delay = (job['at'].to_f - ::Time.now.to_f).ceil
+        delay.positive? ? base + delay : base
+      end
+
+      def acquire_or_drop(pool, job, key, ttl)
+        pool.with do |conn|
+          return yield if conn.call('SET', key, job['jid'], 'NX', 'EX', ttl) == 'OK'
+
+          log_duplicate(job, conn.call('GET', key))
+        end
+        nil
+      end
+
+      def log_duplicate(job, holder)
+        return unless Wurk.logger
+
+        msg = "Wurk::Unique: duplicate #{job['class']} dropped " \
+              "(jid=#{job['jid']} blocked by jid=#{holder || '?'})"
+        Wurk.logger.info { msg }
+      end
+    end
+
+    # ------------------------------------------------------------------
+    # Server middleware — release lock per `unique_until` strategy.
+    # ------------------------------------------------------------------
+    #
+    # `:start`   → DEL before perform. Lock-after-this-point not held; a
+    #              duplicate can be re-enqueued while the first runs.
+    # `:success` → DEL only on successful return. Retries keep the lock.
+    #              Spec §3.7: a raise during perform leaves the lock so
+    #              the retry can proceed; the TTL bounds the worst case.
+    class ServerMiddleware
+      include Wurk::Middleware::ServerMiddleware
+
+      def call(_worker, job, _queue)
+        return yield unless Wurk::Unique.enabled? && Wurk::Unique.coerce_ttl(job['unique_for'])
+
+        mode = unique_until(job)
+        key = Wurk::Unique.lock_key_for(job)
+
+        if mode == :start
+          release(key, job['jid'])
+          yield
+        else
+          result = yield
+          release(key, job['jid'])
+          result
+        end
+      end
+
+      private
+
+      # Honor `unique_until: :start | :success`, fall back to default.
+      def unique_until(job)
+        raw = job['unique_until']
+        return DEFAULT_UNTIL if raw.nil?
+
+        sym = raw.to_sym
+        VALID_UNTIL.include?(sym) ? sym : DEFAULT_UNTIL
+      end
+
+      # CAS DEL: only drop the key if the owning JID still matches ours.
+      # Prevents a long-overdue retry from releasing a fresh lock held by
+      # a re-enqueued duplicate after the original TTL expired.
+      def release(key, jid)
+        redis_pool.with do |conn|
+          conn.call('DEL', key) if conn.call('GET', key) == jid
+        end
+      rescue StandardError => e
+        Wurk.logger&.warn { "Wurk::Unique release failed: #{e.class}: #{e.message}" }
+      end
+    end
+  end
+end

data/lib/wurk/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Wurk
+  VERSION = "0.0.1"
+end