pgbus 0.4.0 → 0.5.0

data/lib/pgbus/cli.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "optparse"
+
 module Pgbus
   module CLI
     module_function
@@ -9,7 +11,7 @@ module Pgbus
 
       case command
       when "start"
-        start_supervisor
+        start_supervisor(args[1..] || [])
       when "status"
         show_status
       when "queues"
@@ -25,13 +27,89 @@ module Pgbus
       end
     end
 
-    def start_supervisor
+    def start_supervisor(args = [])
+      apply_start_options(args)
       Pgbus.logger.info { "[Pgbus] Starting Pgbus #{Pgbus::VERSION}..." }
 
       supervisor = Process::Supervisor.new
       supervisor.run
     end
 
+    # Parses CLI flags for `pgbus start` and applies them to the global
+    # configuration before the supervisor boots. Designed to override the
+    # initializer config without requiring a redeploy.
+    def apply_start_options(args)
+      options = parse_start_options(args)
+
+      Pgbus.configuration.workers = options[:queues] if options[:queues]
+      apply_capsule_filter(options[:capsule]) if options[:capsule]
+      apply_role_filter(options)
+    end
+
+    def parse_start_options(args)
+      options = {}
+      OptionParser.new do |opts|
+        opts.banner = "Usage: pgbus start [options]"
+
+        opts.on("--queues STRING", "Override worker capsules (e.g. \"critical: 5; default: 10\")") do |v|
+          options[:queues] = v
+        end
+
+        opts.on("--capsule NAME", "Run only the named capsule from the configured workers") do |v|
+          options[:capsule] = v
+        end
+
+        opts.on("--workers-only", "Run only the worker processes (no scheduler/dispatcher/consumers/outbox)") do
+          options[:workers_only] = true
+        end
+
+        opts.on("--scheduler-only", "Run only the recurring scheduler (the cron pod pattern)") do
+          options[:scheduler_only] = true
+        end
+
+        opts.on("--dispatcher-only", "Run only the dispatcher (the maintenance pod pattern)") do
+          options[:dispatcher_only] = true
+        end
+      end.parse!(args.dup)
+      options
+    end
+
+    ROLE_FLAG_TO_ROLE = {
+      workers_only: :workers,
+      scheduler_only: :scheduler,
+      dispatcher_only: :dispatcher
+    }.freeze
+    private_constant :ROLE_FLAG_TO_ROLE
+
+    # Translates --workers-only / --scheduler-only / --dispatcher-only into
+    # the corresponding +Pgbus.configuration.roles+ array. Mutually exclusive —
+    # passing more than one of the three flags raises ArgumentError.
+    def apply_role_filter(options)
+      role_flags = options.slice(*ROLE_FLAG_TO_ROLE.keys).compact
+      return if role_flags.empty?
+
+      if role_flags.size > 1
+        raise ArgumentError,
+              "--workers-only, --scheduler-only, and --dispatcher-only are mutually exclusive " \
+              "(got: #{role_flags.keys.map { |k| "--#{k.to_s.tr("_", "-")}" }.join(", ")})"
+      end
+
+      Pgbus.configuration.roles = [ROLE_FLAG_TO_ROLE.fetch(role_flags.keys.first)]
+    end
+
+    def apply_capsule_filter(name)
+      capsule = Pgbus.configuration.capsule_named(name)
+      unless capsule
+        available = (Pgbus.configuration.workers || []).filter_map { |c| c[:name] || c["name"] }
+        raise ArgumentError,
+              "no capsule named #{name.inspect} (available: #{available.join(", ")})"
+      end
+
+      # Go through the public setter so any future normalization/validation
+      # in workers= is applied consistently to the CLI override path too.
+      Pgbus.configuration.workers = [capsule]
+    end
+
     def show_status
       processes = ProcessEntry.order(:kind, :created_at)
                               .select(:kind, :hostname, :pid, :metadata, :last_heartbeat_at)
@@ -65,7 +143,7 @@ module Pgbus
 
     def print_help
       puts <<~HELP
-        Usage: pgbus <command>
+        Usage: pgbus <command> [options]
 
         Commands:
           start    Start the Pgbus supervisor (workers + dispatcher)
@@ -73,6 +151,20 @@ module Pgbus
           queues   List queues with metrics
           version  Show version
           help     Show this help
+
+        Options for `start`:
+          --queues STRING    Override worker capsules from the CLI
+                             (e.g. "critical: 5; default: 10")
+          --capsule NAME     Run only the named capsule from the configured
+                             workers (useful for one-capsule-per-process
+                             deployments)
+          --workers-only     Run only worker processes (no scheduler/
+                             dispatcher/consumers/outbox — for worker-only
+                             containers)
+          --scheduler-only   Run only the recurring scheduler (the cron pod
+                             pattern — exactly one of these per deployment)
+          --dispatcher-only  Run only the dispatcher (the maintenance pod
+                             pattern)
       HELP
     end
   end

data/lib/pgbus/client.rb

@@ -9,6 +9,14 @@ module Pgbus
     PGMQ_REQUIRE_MUTEX = Mutex.new
     private_constant :PGMQ_REQUIRE_MUTEX
 
+    # Throttle window for PGMQ's enable_notify_insert trigger. Postgres
+    # NOTIFYs are coalesced into one wake-up per window, so a value of 250ms
+    # means: at most 4 broadcasts/sec per queue, regardless of insert rate.
+    # The trigger is a Postgres-level concern; exposing it as a setting
+    # never came up in practice and changing it on the fly would require
+    # re-running the trigger DDL on every queue.
+    NOTIFY_THROTTLE_MS = 250
+
     def initialize(config = Pgbus.configuration)
       # Define the PGMQ module before requiring the gem so that Zeitwerk's
       # eager_load (called inside pgmq.rb) can resolve the constant.
@@ -33,9 +41,10 @@ module Pgbus
       else
         # With a String URL or Hash params, pgmq-ruby creates its own dedicated
         # PG::Connection per pool slot — no shared state with ActiveRecord.
-        # Use the configured pool_size and let pgmq-ruby's connection_pool handle
+        # Use the resolved pool size (auto-tuned from worker thread counts
+        # unless explicitly set) and let pgmq-ruby's connection_pool handle
         # concurrency internally (no mutex needed).
-        @pgmq = PGMQ::Client.new(conn_opts, pool_size: config.pool_size, pool_timeout: config.pool_timeout)
+        @pgmq = PGMQ::Client.new(conn_opts, pool_size: config.resolved_pool_size, pool_timeout: config.pool_timeout)
         @pgmq_mutex = nil
       end
 
@@ -221,6 +230,49 @@ module Pgbus
       result
     end
 
+    # Check whether a message exists in the given queue.
+    #
+    # Pass either +msg_id+ for a fast primary-key lookup, or +uniqueness_key+
+    # to scan the queue for any message whose payload carries that key in the
+    # +pgbus_uniqueness_key+ JSONB field. The latter is used by the dispatcher
+    # reaper to determine if a uniqueness lock with msg_id=0 (placeholder)
+    # still has a corresponding queue message.
+    #
+    # +queue_name+ may be either a logical name (e.g. "default") or an already
+    # prefixed physical name (e.g. "pgbus_default"). The client normalizes both.
+    #
+    # Returns:
+    #   true  — the message definitely exists in the queue
+    #   false — the message definitely does not exist
+    #   nil   — could not determine (e.g. queue table missing or unknown error).
+    #           Callers MUST treat nil as "exists" for safety.
+    def message_exists?(queue_name, msg_id: nil, uniqueness_key: nil)
+      has_msg_id = !msg_id.nil?
+      has_uniqueness_key = !uniqueness_key.nil?
+      raise ArgumentError, "pass exactly one of msg_id or uniqueness_key" unless has_msg_id ^ has_uniqueness_key
+
+      full_name = resolve_full_queue_name(queue_name)
+      sanitized = QueueNameValidator.sanitize!(full_name)
+
+      synchronized do
+        with_raw_connection do |conn|
+          if has_msg_id
+            msg_id_present?(conn, sanitized, msg_id.to_i)
+          else
+            uniqueness_key_present?(conn, sanitized, uniqueness_key)
+          end
+        end
+      end
+    rescue ActiveRecord::StatementInvalid => e
+      raise unless undefined_table_error?(e)
+
+      nil
+    rescue StandardError => e
+      raise unless defined?(PG::UndefinedTable) && e.is_a?(PG::UndefinedTable)
+
+      nil
+    end
+
     def purge_archive(queue_name, older_than:, batch_size: 1000)
       full_name = config.queue_name(queue_name)
       sanitized = QueueNameValidator.sanitize!(full_name)
@@ -266,6 +318,42 @@ module Pgbus
 
     private
 
+    # Accept either a logical name ("default") or an already-prefixed
+    # physical name ("pgbus_default") and return the physical name.
+    # Coerces symbols to strings so callers can pass either form.
+    def resolve_full_queue_name(queue_name)
+      name = queue_name.to_s
+      prefix = "#{config.queue_prefix}_"
+      name.start_with?(prefix) ? name : config.queue_name(name)
+    end
+
+    def msg_id_present?(conn, sanitized, msg_id)
+      result = conn.exec_params(
+        "SELECT 1 FROM pgmq.q_#{sanitized} WHERE msg_id = $1 LIMIT 1",
+        [msg_id]
+      )
+      result.ntuples.positive?
+    end
+
+    def uniqueness_key_present?(conn, sanitized, uniqueness_key)
+      result = conn.exec_params(
+        "SELECT 1 FROM pgmq.q_#{sanitized} " \
+        "WHERE message::jsonb ->> 'pgbus_uniqueness_key' = $1 LIMIT 1",
+        [uniqueness_key]
+      )
+      result.ntuples.positive?
+    end
+
+    # Detect "relation does not exist" via the underlying PG error type.
+    # Falls back to message matching only if PG::UndefinedTable is undefined
+    # (very old pg gem) — never relies on locale-sensitive text.
+    def undefined_table_error?(error)
+      cause = error.respond_to?(:cause) ? error.cause : nil
+      return true if defined?(PG::UndefinedTable) && cause.is_a?(PG::UndefinedTable)
+
+      false
+    end
+
     def collect_configured_queues
       queues = Set.new
       queues << config.default_queue
@@ -352,7 +440,7 @@ module Pgbus
       @queues_created.compute_if_absent(full_name) do
         synchronized do
           @pgmq.create(full_name)
-          @pgmq.enable_notify_insert(full_name, throttle_interval_ms: config.notify_throttle_ms) if config.listen_notify
+          @pgmq.enable_notify_insert(full_name, throttle_interval_ms: NOTIFY_THROTTLE_MS) if config.listen_notify
         end
         true
       end

data/lib/pgbus/configuration/capsule_dsl.rb

@@ -0,0 +1,190 @@
+# frozen_string_literal: true
+
+module Pgbus
+  class Configuration
+    # Parses the capsule string DSL into the internal worker config array.
+    #
+    # The DSL lets users describe worker thread pools (capsules) compactly:
+    #
+    #   "*: 5"                                # one capsule, all queues, 5 threads
+    #   "critical: 5; default: 10"            # two capsules
+    #   "critical, default: 5"                # one capsule, list = strict priority
+    #   "default, mailers: 10; *: 2"          # specialized + fallback wildcard
+    #   "staging_*: 3"                        # prefix wildcard
+    #
+    # Operators:
+    #   ,    queue separator within a capsule (list order = strict priority)
+    #   ;    capsule separator (each becomes its own thread pool)
+    #   :N   thread count for the capsule (default 5)
+    #   *    wildcard, matches all queues
+    #   *_   trailing wildcard, prefix match (e.g. "staging_*")
+    #
+    # Returns +Array<Hash>+ in the same shape as the legacy +workers:+ array,
+    # so the rest of the codebase can consume it without changes:
+    #
+    #   [{ queues: ["critical"], threads: 5 },
+    #    { queues: ["default", "mailers"], threads: 10 }]
+    #
+    # Validation runs at parse time. Errors include the offending input and
+    # a clear message naming the rule that was violated.
+    class CapsuleDSL
+      DEFAULT_THREADS = 5
+      WILDCARD = "*"
+
+      # Valid queue tokens accepted by the DSL:
+      #   "*"         bare wildcard (matches all queues at runtime)
+      #   "default"   bare queue name
+      #   "staging_*" prefix wildcard (matches queues by prefix at runtime)
+      #
+      # The character class for queue names — alphanumerics and underscores
+      # only — intentionally matches Pgbus::QueueNameValidator::VALID_QUEUE_NAME_PATTERN.
+      # Hyphens, dots, and other punctuation are NOT permitted because PGMQ
+      # interpolates queue names directly into SQL identifiers (q_<name>,
+      # a_<name>) and only those characters are safe there. If you change
+      # this pattern, also update QueueNameValidator to keep them in sync.
+      #
+      # The "*" tokens (bare and trailing) are DSL-only — they get expanded
+      # to concrete queue names at runtime before reaching QueueNameValidator.
+      QUEUE_NAME_PATTERN = /\A(?:\*|[a-zA-Z0-9_]+\*?)\z/
+      THREAD_COUNT_PATTERN = /\A\d+\z/
+
+      class ParseError < ArgumentError; end
+
+      def self.parse(input)
+        new(input).parse
+      end
+
+      def initialize(input)
+        @input = input
+      end
+
+      def parse
+        validate_input_type!
+        validate_input_not_empty!
+
+        capsules = split_capsules(@input).map { |segment| parse_capsule(segment) }
+        validate_no_duplicate_queues_across_capsules!(capsules)
+        capsules
+      end
+
+      private
+
+      def validate_input_type!
+        return if @input.is_a?(String)
+
+        raise ParseError,
+              "expected String, got #{@input.class} (#{@input.inspect})"
+      end
+
+      def validate_input_not_empty!
+        return if @input.strip != ""
+
+        raise ParseError, "empty capsule string — must declare at least one queue"
+      end
+
+      # Splits on `;` and trims whitespace, dropping a trailing empty segment
+      # so `"a; b;"` is treated the same as `"a; b"`. Leading-empty segments
+      # ("; a") are kept and rejected by parse_capsule with a clearer error.
+      def split_capsules(string)
+        segments = string.strip.split(";").map(&:strip)
+        segments.pop if segments.last == "" # trailing semicolon
+        segments
+      end
+
+      def parse_capsule(segment)
+        if segment == ""
+          raise ParseError,
+                "empty capsule in #{@input.inspect} — leading or doubled semicolons are not allowed"
+        end
+
+        queues_part, threads_part = split_on_threads(segment)
+        queues = parse_queue_list(queues_part, segment)
+        threads = parse_thread_count(threads_part, segment)
+        validate_no_duplicates_within_capsule!(queues, segment)
+
+        { queues: queues, threads: threads }
+      end
+
+      # Splits a capsule segment on the LAST `:` so queue names containing
+      # underscores or digits are unaffected. Returns [queues_part, threads_part]
+      # where threads_part is nil when no `:` is present.
+      def split_on_threads(segment)
+        idx = segment.rindex(":")
+        return [segment, nil] unless idx
+
+        [segment[0...idx].strip, segment[(idx + 1)..].strip]
+      end
+
+      def parse_queue_list(queues_part, original_segment)
+        if queues_part.empty?
+          raise ParseError,
+                "expected queue name before ':' in #{original_segment.inspect}"
+        end
+
+        queues = queues_part.split(",").map(&:strip)
+        queues.each { |q| validate_queue_name!(q, original_segment) }
+        queues
+      end
+
+      def validate_queue_name!(name, original_segment)
+        return if name.match?(QUEUE_NAME_PATTERN)
+
+        raise ParseError,
+              "invalid character in queue name #{name.inspect} (in #{original_segment.inspect}) — " \
+              "queue names must match /[a-zA-Z0-9_]+\\*?/"
+      end
+
+      def parse_thread_count(threads_part, original_segment)
+        return DEFAULT_THREADS if threads_part.nil?
+
+        if threads_part.empty?
+          raise ParseError,
+                "expected thread count after ':' in #{original_segment.inspect}"
+        end
+
+        unless threads_part.match?(THREAD_COUNT_PATTERN)
+          raise ParseError,
+                "invalid thread count #{threads_part.inspect} in #{original_segment.inspect} — " \
+                "must be a positive integer"
+        end
+
+        n = threads_part.to_i
+        if n.zero?
+          raise ParseError,
+                "thread count must be a positive integer, got 0 in #{original_segment.inspect}"
+        end
+
+        n
+      end
+
+      def validate_no_duplicates_within_capsule!(queues, original_segment)
+        seen = {}
+        queues.each do |q|
+          if seen[q]
+            raise ParseError,
+                  "queue #{q.inspect} listed twice in capsule #{original_segment.inspect} — " \
+                  "duplicate queues within a capsule are not allowed"
+          end
+          seen[q] = true
+        end
+      end
+
+      def validate_no_duplicate_queues_across_capsules!(capsules)
+        return if capsules.size < 2
+
+        seen = {}
+        capsules.each_with_index do |capsule, idx|
+          capsule[:queues].each do |q|
+            if seen[q]
+              label = (q == WILDCARD ? "wildcard '*'" : "queue #{q.inspect}")
+              raise ParseError,
+                    "#{label} appears in two capsules (positions #{seen[q]} and #{idx + 1}) " \
+                    "in #{@input.inspect} — each queue can only be assigned to one capsule"
+            end
+            seen[q] = idx + 1
+          end
+        end
+      end
+    end
+  end
+end