pgbus 0.1.4 → 0.1.6

data/lib/pgbus/process/worker.rb

@@ -9,35 +9,56 @@ module Pgbus
 
       attr_reader :queues, :threads, :config
 
-      def initialize(queues:, threads: 5, config: Pgbus.configuration)
+      def initialize(queues:, threads: 5, config: Pgbus.configuration,
+                     single_active_consumer: false, consumer_priority: 0)
         @queues = Array(queues)
         @threads = threads
         @config = config
-        @shutting_down = false
+        @single_active_consumer = single_active_consumer
+        @consumer_priority = consumer_priority
+        @lifecycle = Lifecycle.new
         @jobs_processed = Concurrent::AtomicFixnum.new(0)
         @jobs_failed = Concurrent::AtomicFixnum.new(0)
+        @in_flight = Concurrent::AtomicFixnum.new(0)
+        @rate_counter = RateCounter.new(:processed, :failed, :dequeued)
         @started_at = Time.now
         @executor = Pgbus::ActiveJob::Executor.new
         @pool = Concurrent::FixedThreadPool.new(threads)
+        @circuit_breaker = Pgbus::CircuitBreaker.new(config: config)
+        @queue_lock = QueueLock.new if @single_active_consumer
+        @wake_signal = WakeSignal.new
       end
 
       def stats
-        { jobs_processed: @jobs_processed.value, jobs_failed: @jobs_failed.value, started_at: @started_at }
+        {
+          jobs_processed: @jobs_processed.value,
+          jobs_failed: @jobs_failed.value,
+          in_flight: @in_flight.value,
+          state: @lifecycle.state,
+          consumer_priority: @consumer_priority,
+          single_active_consumer: @single_active_consumer,
+          locked_queues: @queue_lock&.held_queues || [],
+          rates: @rate_counter.rates,
+          started_at: @started_at
+        }
       end
 
       def run
         setup_signals
         start_heartbeat
         resolve_wildcard_queues
+        @lifecycle.transition_to!(:running)
         Pgbus.logger.info { "[Pgbus] Worker started: queues=#{queues.join(",")} threads=#{threads} pid=#{::Process.pid}" }
 
         loop do
           process_signals
-          break if @shutting_down && @pool.queue_length.zero?
-          break if recycle_needed? && @pool.queue_length.zero?
+          check_recycle
 
-          claim_and_execute unless @shutting_down || recycle_needed?
-          interruptible_sleep(config.polling_interval) if (@shutting_down || recycle_needed?) && !@pool.queue_length.zero?
+          break if @lifecycle.stopped?
+          break if @lifecycle.draining? && @pool.queue_length.zero?
+
+          claim_and_execute if @lifecycle.can_process?
+          @wake_signal.wait(timeout: config.polling_interval) if @lifecycle.draining? || @lifecycle.paused?
         end
 
         shutdown
@@ -45,43 +66,62 @@ module Pgbus
 
       def graceful_shutdown
         Pgbus.logger.info { "[Pgbus] Worker shutting down gracefully..." }
-        @shutting_down = true
+        @lifecycle.transition_to(:draining)
+        @wake_signal.notify!
       end
 
       def immediate_shutdown
         Pgbus.logger.warn { "[Pgbus] Worker shutting down immediately!" }
-        @shutting_down = true
+        @lifecycle.transition_to!(:stopped)
+        @wake_signal.notify!
         @pool.kill
       end
 
       private
 
       def claim_and_execute
+        poll_interval = effective_polling_interval
+
         idle = @pool.max_length - @pool.queue_length
-        return interruptible_sleep(config.polling_interval) if idle <= 0
+        return @wake_signal.wait(timeout: poll_interval) if idle <= 0
+
+        if config.prefetch_limit
+          available = config.prefetch_limit - @in_flight.value
+          return @wake_signal.wait(timeout: poll_interval) if available <= 0
+
+          idle = [idle, available].min
+        end
 
         tagged_messages = fetch_messages(idle)
 
         if tagged_messages.empty?
-          interruptible_sleep(config.polling_interval)
+          @wake_signal.wait(timeout: poll_interval)
           return
         end
 
-        tagged_messages.each do |queue_name, message|
-          @pool.post { process_message(message, queue_name) }
+        @rate_counter.increment(:dequeued, tagged_messages.size)
+        tagged_messages.each do |queue_name, message, source_queue|
+          @in_flight.increment
+          @pool.post { process_message(message, queue_name, source_queue: source_queue) }
         end
       end
 
       # Returns an array of [queue_name, message] pairs so we always know
       # which queue each message came from (PGMQ messages don't carry this).
       def fetch_messages(qty)
-        if queues.size == 1
-          queue = queues.first
+        active_queues = queues.reject { |q| @circuit_breaker.paused?(q) }
+        active_queues = active_queues.select { |q| @queue_lock.try_lock(q) } if @single_active_consumer
+        return [] if active_queues.empty?
+
+        if priority_enabled?
+          fetch_prioritized(active_queues, qty)
+        elsif active_queues.size == 1
+          queue = active_queues.first
           messages = Pgbus.client.read_batch(queue, qty: qty) || []
           messages.map { |m| [queue, m] }
         else
-          per_queue = [(qty / queues.size.to_f).ceil, 1].max
-          queues.flat_map do |q|
+          per_queue = [(qty / active_queues.size.to_f).ceil, 1].max
+          active_queues.flat_map do |q|
             (Pgbus.client.read_batch(q, qty: per_queue) || []).map { |m| [q, m] }
           end.first(qty)
         end
@@ -90,13 +130,45 @@ module Pgbus
         []
       end
 
-      def process_message(message, queue_name)
-        result = @executor.execute(message, queue_name)
+      def fetch_prioritized(active_queues, qty)
+        remaining = qty
+        results = []
+
+        active_queues.each do |q|
+          break if remaining <= 0
+
+          batch = Pgbus.client.read_batch_prioritized(q, qty: remaining)
+          batch.each do |physical_queue, message|
+            results << [q, message, physical_queue]
+          end
+          remaining -= batch.size
+        end
+
+        results
+      end
+
+      def priority_enabled?
+        config.priority_levels && config.priority_levels > 1
+      end
+
+      def process_message(message, queue_name, source_queue: nil)
+        result = @executor.execute(message, queue_name, source_queue: source_queue)
         @jobs_processed.increment
-        @jobs_failed.increment if result == :failed
+        @rate_counter.increment(:processed)
+        if result == :failed
+          @jobs_failed.increment
+          @rate_counter.increment(:failed)
+          @circuit_breaker.record_failure(queue_name)
+        else
+          @circuit_breaker.record_success(queue_name)
+        end
       rescue StandardError => e
         @jobs_failed.increment
+        @rate_counter.increment(:failed)
+        @circuit_breaker.record_failure(queue_name)
         Pgbus.logger.error { "[Pgbus] Unhandled error processing message: #{e.message}" }
+      ensure
+        @in_flight.decrement
       end
 
       # Resolve "*" to all non-DLQ queues from pgmq.meta, stripping the prefix.
@@ -125,6 +197,13 @@ module Pgbus
         @queues = [config.default_queue]
       end
 
+      def check_recycle
+        return unless @lifecycle.running? && recycle_needed?
+
+        @lifecycle.transition_to(:draining)
+        @wake_signal.notify!
+      end
+
       def recycle_needed?
         exceeded_max_jobs? || exceeded_max_memory? || exceeded_max_lifetime?
       end
@@ -162,10 +241,26 @@ module Pgbus
         end
       end
 
+      def effective_polling_interval
+        return config.polling_interval if @consumer_priority.zero?
+
+        ConsumerPriority.effective_polling_interval(
+          base_interval: config.polling_interval,
+          my_priority: @consumer_priority,
+          max_priority: ConsumerPriority.max_active_priority(queues, ::Process.pid)
+        )
+      rescue StandardError
+        config.polling_interval
+      end
+
       def start_heartbeat
         @heartbeat = Heartbeat.new(
           kind: "worker",
-          metadata: { queues: queues, threads: threads, pid: ::Process.pid }
+          metadata: {
+            queues: queues, threads: threads, pid: ::Process.pid,
+            consumer_priority: @consumer_priority
+          },
+          on_beat: -> { @rate_counter.snapshot! }
         )
         @heartbeat.start
       end
@@ -174,6 +269,7 @@ module Pgbus
         Pgbus.logger.info { "[Pgbus] Worker draining thread pool..." }
         @pool.shutdown
         @pool.wait_for_termination(30)
+        @queue_lock&.unlock_all
         @heartbeat&.stop
         restore_signals
         Pgbus.logger.info { "[Pgbus] Worker stopped. Processed: #{@jobs_processed.value}, Failed: #{@jobs_failed.value}" }

data/lib/pgbus/queue_factory.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module Pgbus
+  # Dispatches queue operations based on queue type (standard vs priority).
+  # Replaces conditional `priority_enabled?` checks scattered through Client
+  # with a single strategy object selected at initialization.
+  #
+  # Inspired by LavinMQ's QueueFactory which dispatches queue creation by
+  # type: standard, durable, priority, stream, delayed.
+  module QueueFactory
+    def self.for(config)
+      if config.priority_levels && config.priority_levels > 1
+        PriorityStrategy.new(config)
+      else
+        StandardStrategy.new(config)
+      end
+    end
+
+    # Standard single-queue strategy: one PGMQ queue per logical name.
+    class StandardStrategy
+      def initialize(config)
+        @config = config
+      end
+
+      def physical_queue_names(name)
+        [@config.queue_name(name)]
+      end
+
+      def target_queue(name, _priority)
+        @config.queue_name(name)
+      end
+
+      def priority?
+        false
+      end
+    end
+
+    # Priority sub-queue strategy: N PGMQ queues per logical name (_p0.._pN).
+    class PriorityStrategy
+      def initialize(config)
+        @config = config
+      end
+
+      def physical_queue_names(name)
+        @config.priority_queue_names(name)
+      end
+
+      def target_queue(name, priority)
+        level = if priority
+                  priority.clamp(0, @config.priority_levels - 1)
+                else
+                  @config.default_priority.clamp(0, @config.priority_levels - 1)
+                end
+        @config.priority_queue_name(name, level)
+      end
+
+      def priority?
+        true
+      end
+    end
+  end
+end

data/lib/pgbus/rate_counter.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+module Pgbus
+  # Thread-safe rate counter inspired by LavinMQ's rate_stats macro.
+  # Tracks absolute counts via AtomicFixnum and computes rates as
+  # deltas between periodic snapshots.
+  #
+  # Usage:
+  #   counter = RateCounter.new(:enqueued, :dequeued, :failed)
+  #   counter.increment(:enqueued)
+  #   counter.snapshot!          # call periodically (e.g. every 5s)
+  #   counter.rate(:enqueued)    # => msgs/s since last snapshot
+  #   counter.rates              # => { enqueued: 12.4, dequeued: 10.1, failed: 0.2 }
+  class RateCounter
+    attr_reader :names
+
+    def initialize(*names)
+      @names = names.map(&:to_sym)
+      @counters = {}
+      @last_values = {}
+      @rates = {}
+      @last_snapshot_at = monotonic_now
+
+      @names.each do |name|
+        @counters[name] = Concurrent::AtomicFixnum.new(0)
+        @last_values[name] = 0
+        @rates[name] = 0.0
+      end
+    end
+
+    def increment(name, delta = 1)
+      @counters.fetch(name).increment(delta)
+    end
+
+    def count(name)
+      @counters.fetch(name).value
+    end
+
+    def rate(name)
+      @rates.fetch(name)
+    end
+
+    def rates
+      @names.to_h { |name| [name, @rates[name]] }
+    end
+
+    def counts
+      @names.to_h { |name| [name, @counters[name].value] }
+    end
+
+    def snapshot!
+      now = monotonic_now
+      elapsed = now - @last_snapshot_at
+      return if elapsed <= 0
+
+      @names.each do |name|
+        current = @counters[name].value
+        delta = current - @last_values[name]
+        @rates[name] = (delta / elapsed).round(1)
+        @last_values[name] = current
+      end
+
+      @last_snapshot_at = now
+    end
+
+    def to_h
+      {
+        counts: counts,
+        rates: rates
+      }
+    end
+
+    private
+
+    def monotonic_now
+      ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+    end
+  end
+end

data/lib/pgbus/recurring/schedule.rb

@@ -87,7 +87,7 @@ module Pgbus
       end
 
       def resolve_queue(task)
-        @config.queue_name(task.queue_name || @config.default_queue)
+        task.queue_name || @config.default_queue
       end
 
       def build_headers(task, run_at)

data/lib/pgbus/uniqueness.rb

@@ -0,0 +1,169 @@
+# frozen_string_literal: true
+
+require "active_support/concern"
+require "socket"
+
+module Pgbus
+  # Job uniqueness guarantees: prevent duplicate jobs from running concurrently.
+  #
+  # Unlike concurrency limits (which allow N concurrent jobs for the same key),
+  # uniqueness ensures AT MOST ONE job with a given key exists in the system
+  # at any time — from enqueue through completion.
+  #
+  # Lock lifecycle:
+  #   1. Enqueue: lock acquired (state: queued), no owner yet
+  #   2. Execution start: lock transitions to (state: executing, owner_pid: PID)
+  #   3. Completion/DLQ: lock released (row deleted)
+  #   4. Crash recovery: reaper detects orphaned locks by cross-referencing
+  #      owner_pid against pgbus_processes heartbeats
+  #   5. Last resort: expires_at TTL (default 24h) handles the case where
+  #      even the reaper can't run (entire supervisor dead)
+  #
+  # Strategies:
+  #   :until_executed  — Lock acquired at enqueue, held through execution, released on
+  #                      completion or DLQ. Prevents duplicate enqueue AND duplicate execution.
+  #
+  #   :while_executing — Lock acquired at execution start, released on completion.
+  #                      Allows duplicate enqueue (multiple copies in queue) but only one
+  #                      executes at a time.
+  #
+  # Usage:
+  #   class ImportOrderJob < ApplicationJob
+  #     ensures_uniqueness strategy: :until_executed,
+  #                        key: ->(order_id) { "import-order-#{order_id}" },
+  #                        on_conflict: :reject
+  #
+  #     def perform(order_id)
+  #       # Only one instance of this job per order_id can exist at a time
+  #     end
+  #   end
+  module Uniqueness
+    extend ActiveSupport::Concern
+
+    METADATA_KEY = "pgbus_uniqueness_key"
+    STRATEGY_KEY = "pgbus_uniqueness_strategy"
+    TTL_KEY = "pgbus_uniqueness_lock_ttl"
+
+    # 24 hours — last-resort fallback only. The reaper handles normal crash recovery.
+    DEFAULT_LOCK_TTL = 24 * 60 * 60
+
+    VALID_STRATEGIES = %i[until_executed while_executing].freeze
+    VALID_CONFLICTS = %i[reject discard log].freeze
+
+    class_methods do
+      def ensures_uniqueness(strategy: :until_executed, key: nil, lock_ttl: DEFAULT_LOCK_TTL, on_conflict: :reject)
+        raise ArgumentError, "strategy must be one of: #{VALID_STRATEGIES.join(", ")}" unless VALID_STRATEGIES.include?(strategy)
+        raise ArgumentError, "on_conflict must be one of: #{VALID_CONFLICTS.join(", ")}" unless VALID_CONFLICTS.include?(on_conflict)
+
+        valid_ttl = lock_ttl.is_a?(Numeric) || (defined?(ActiveSupport::Duration) && lock_ttl.is_a?(ActiveSupport::Duration))
+        raise ArgumentError, "lock_ttl must be a positive number or Duration" unless valid_ttl && lock_ttl.positive?
+        raise ArgumentError, "key must be callable (Proc or lambda)" if key && !key.respond_to?(:call)
+
+        @pgbus_uniqueness = {
+          strategy: strategy,
+          key: key || ->(*) { name },
+          lock_ttl: lock_ttl,
+          on_conflict: on_conflict
+        }.freeze
+      end
+
+      def pgbus_uniqueness
+        @pgbus_uniqueness
+      end
+    end
+
+    class << self
+      def resolve_key(active_job)
+        config = uniqueness_config(active_job)
+        return nil unless config
+
+        args = active_job.arguments
+        last = args.last
+        if last.is_a?(Hash) && last.each_key.all?(Symbol)
+          config[:key].call(*args[...-1], **last)
+        else
+          config[:key].call(*args)
+        end
+      end
+
+      def inject_metadata(active_job, payload_hash)
+        config = uniqueness_config(active_job)
+        return payload_hash unless config
+
+        key = resolve_key(active_job)
+        return payload_hash unless key
+
+        payload_hash.merge(
+          METADATA_KEY => key,
+          STRATEGY_KEY => config[:strategy].to_s,
+          TTL_KEY => config[:lock_ttl]
+        )
+      end
+
+      def extract_key(payload)
+        payload&.dig(METADATA_KEY)
+      end
+
+      def extract_strategy(payload)
+        payload&.dig(STRATEGY_KEY)&.to_sym
+      end
+
+      def uniqueness_config(active_job)
+        return nil unless active_job.class.respond_to?(:pgbus_uniqueness)
+
+        active_job.class.pgbus_uniqueness
+      end
+
+      # Acquire the uniqueness lock at enqueue time (:until_executed only).
+      # Lock state: queued, no owner_pid yet.
+      # Returns :acquired or :locked.
+      def acquire_enqueue_lock(key, active_job)
+        config = uniqueness_config(active_job)
+        return :no_lock unless config
+        return :no_lock unless config[:strategy] == :until_executed
+
+        acquired = JobLock.acquire!(
+          key,
+          job_class: active_job.class.name,
+          job_id: active_job.job_id,
+          state: "queued",
+          ttl: config[:lock_ttl]
+        )
+        acquired ? :acquired : :locked
+      end
+
+      # Transition a queued lock to executing state when the worker picks it up.
+      # Called for :until_executed jobs at execution start.
+      def claim_for_execution!(key, ttl:)
+        JobLock.claim_for_execution!(key, owner_pid: ::Process.pid, owner_hostname: Socket.gethostname, ttl: ttl)
+      end
+
+      # Acquire the uniqueness lock at execution time (:while_executing only).
+      # Lock state: executing with owner_pid.
+      # Returns true if acquired, false if another instance is running.
+      def acquire_execution_lock(key, payload)
+        strategy = extract_strategy(payload)
+        return true unless strategy == :while_executing
+
+        ttl = payload[TTL_KEY] || DEFAULT_LOCK_TTL
+
+        JobLock.acquire!(
+          key,
+          job_class: payload["job_class"],
+          job_id: payload["job_id"],
+          state: "executing",
+          owner_pid: ::Process.pid,
+          owner_hostname: Socket.gethostname,
+          ttl: ttl
+        )
+      end
+
+      # Release the uniqueness lock after execution completes.
+      def release_lock(key)
+        return unless key
+
+        JobLock.release!(key)
+      end
+    end
+  end
+end

data/lib/pgbus/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Pgbus
-  VERSION = "0.1.4"
+  VERSION = "0.1.6"
 end