pgbus 0.9.10 → 0.10.0

data/lib/pgbus/client.rb

@@ -7,6 +7,7 @@ require_relative "client/read_after"
 require_relative "client/ensure_stream_queue"
 require_relative "client/notify_stream"
 require_relative "client/connection_health"
+require_relative "client/resizable_pool"
 
 module Pgbus
   class Client
@@ -59,6 +60,12 @@ module Pgbus
         # operations through a mutex.
         @pgmq = PGMQ::Client.new(conn_opts, pool_size: 1, pool_timeout: config.pool_timeout)
         @pgmq_mutex = Mutex.new
+        # No dedicated streams pool on this path: a second PGMQ::Client would
+        # still funnel through the same non-thread-safe AR raw_connection.
+        # Stream publish + replay share the single serialized connection —
+        # @streams_pgmq aliases @pgmq so the code paths are uniform, and
+        # #with_streams_connection falls back to with_raw_connection.
+        @streams_pgmq = @pgmq
       else
         # With a String URL or Hash params, pgmq-ruby creates its own dedicated
         # PG::Connection per pool slot — no shared state with ActiveRecord.
@@ -72,13 +79,49 @@ module Pgbus
         # Both raise clean PG errors — no Ruby Timeout, no Thread#raise. Only
         # safe on this dedicated-connection branch — never on the shared-AR Proc
         # path, where statement_timeout would leak into application queries.
-        conn_opts = apply_connection_bounds(conn_opts)
+        conn_opts = wrap_session_gucs(apply_connection_bounds(conn_opts))
         @pgmq = PGMQ::Client.new(conn_opts, pool_size: config.resolved_pool_size, pool_timeout: config.pool_timeout)
         @pgmq_mutex = nil
+        # Dedicated streams pool (issue #315): isolates the durable-stream
+        # publish INSERT (#send_stream_message) and the dispatcher's per-wake
+        # replay reads (#read_after) from the job pool, so a saturated worker
+        # pool can't delay a broadcast on pool checkout, and each wake reuses a
+        # persistent connection instead of a fresh PG.connect per call. Its own
+        # PGMQ::Client → its own connection_pool, sized independently of worker
+        # thread counts.
+        # Build the streams pool from streams_connection_options (which defaults
+        # to connection_options but honors streams_database_url/host/port for a
+        # separate/direct streams DB — issue #315), bounds-applied, and tagged
+        # with a per-process application_name so the autoscaler can count peer
+        # processes from pg_stat_activity (issue #323 P1/P2). Snapshot it so a
+        # hot-swap rebuilds a byte-identical pool at a new size.
+        @streams_conn_opts = wrap_session_gucs(
+          tag_application_name(
+            apply_connection_bounds(config.streams_connection_options)
+          )
+        )
+        @streams_pgmq = PGMQ::Client.new(@streams_conn_opts, pool_size: config.streams_pool_size,
+                                                             pool_timeout: config.streams_pool_timeout)
       end
 
+      # Wrap the streams pool so its live reference can be atomically hot-swapped
+      # to a new size under load without losing broadcasts or leaking connections
+      # (issue #323 spike; #resize_streams_pool). All streams-pool access goes
+      # through this — see #streams_pool. Default behavior with no swap is
+      # byte-identical (one AtomicReference read + a counter bump per op).
+      @streams_pool = ResizablePool.new(
+        @streams_pgmq,
+        shared: @shared_connection,
+        drain_timeout: config.streams_pool_timeout + 1.0,
+        logger: Pgbus.logger
+      )
+
       @queues_created = Concurrent::Map.new
       @stream_indexes_created = Concurrent::Map.new
+      # Guards the one-time build of the publisher autoscale trigger (issue #323).
+      # NOT @pgmq_mutex — that is nil on the dedicated path (the only path the
+      # trigger exists on), so it wouldn't serialize concurrent first-publishers.
+      @streams_trigger_mutex = Mutex.new
       @queue_strategy = QueueFactory.for(config)
       @schema_ensured = schema_ensured
       @connection_health = ConnectionHealth.new(
@@ -122,9 +165,22 @@ module Pgbus
     # carries the underlying error plus which config source was in use.
     def verify_connection!
       synchronized do
-        @pgmq.with_connection { |conn| conn.exec("SELECT 1") }
+        @pgmq.with_connection do |conn|
+          conn.exec("SELECT 1")
+          # When require_primary is set, reject a connection that landed on a
+          # read-only replica at boot rather than letting a read/write-splitting
+          # pooler silently route pgmq's VOLATILE read/archive to a standby,
+          # where workers read nothing and jobs stop with a healthy heartbeat
+          # (issue #332). Off by default, so a single-primary deployment is
+          # unaffected.
+          Process::PrimaryValidator.validate_primary!(conn) if config.require_primary
+        end
       end
       true
+    rescue Process::ReplicaConnectionError => e
+      raise ConfigurationError,
+            "Database connection via #{connection_source} landed on a read-only replica " \
+            "(require_primary is set): #{e.message}"
     rescue PGMQ::Errors::ConnectionError, PG::Error => e
       raise ConfigurationError, "Database connection failed via #{connection_source}: #{e.message}"
     end
@@ -141,6 +197,17 @@ module Pgbus
       true
     end
 
+    # Whether the job connection currently lands on a read-only replica
+    # (pg_is_in_recovery() => t). Used by the doctor to warn about a
+    # read/write-splitting pooler that could route pgmq's VOLATILE read/archive
+    # to a standby, silently stalling job processing (issue #332). Raw PG error
+    # propagates so the caller can render the reason.
+    def in_recovery?
+      with_raw_connection do |conn|
+        conn.exec(Process::PrimaryValidator::RECOVERY_QUERY).getvalue(0, 0) == "t"
+      end
+    end
+
     # The logical queue names pgbus expects to exist based on the configuration
     # (default queue + worker capsules + recurring tasks). Public wrapper around
     # collect_configured_queues so the doctor can diff configured-vs-existing
@@ -240,6 +307,38 @@ module Pgbus
       end
     end
 
+    # Durable stream broadcast. Unlike #send_message, this ALWAYS targets the
+    # bare queue (config.queue_name) and never the priority strategy's
+    # _p0.._pN sub-queues: streams are delivered by a non-consuming peek
+    # (read_after) on the bare queue, and the streamer LISTENs on the bare
+    # channel, so a broadcast routed to _p1 would never reach the browser
+    # (issue #310). ensure_stream_queue creates the bare queue + NOTIFY
+    # trigger + archive index, mirroring this bare-name write path.
+    def send_stream_message(stream_name, payload, headers: nil, delay: 0)
+      target = config.queue_name(stream_name)
+      # Capture the produced msg_id — it is this method's return value (callers
+      # like Stream#broadcast rely on it), so the autoscale trigger below must NOT
+      # become the last expression.
+      msg_id = Instrumentation.instrument("pgbus.client.send_message", queue: target) do
+        with_stale_connection_retry do
+          ensure_stream_queue(stream_name)
+          # Publish through the dedicated streams pool (issue #315) so a
+          # saturated job pool can't block a broadcast on pool checkout. On the
+          # shared-AR path @streams_pgmq aliases @pgmq and synchronized still
+          # serializes on the mutex.
+          synchronized do
+            streams_pool.produce(target, serialize(payload), headers: headers && serialize(headers), delay: delay)
+          end
+        end
+      end
+      # Opportunistically autoscale the streams pool from the publish path so a
+      # pure-publisher process (no streamer) still grows under a broadcast storm
+      # (issue #323 follow-up). Throttled + fail-soft — never delays or breaks the
+      # broadcast; nil (a no-op) unless autoscale is on and the pool is dedicated.
+      streams_pool_trigger&.maybe_check
+      msg_id
+    end
+
     def send_batch(queue_name, payloads, headers: nil, delay: 0)
       full_name = config.queue_name(queue_name)
       serialized, serialized_headers = serialize_batch(payloads, headers)
@@ -446,6 +545,16 @@ module Pgbus
       {}
     end
 
+    # Same shape as #pool_stats but for the dedicated streams pool (issue #315).
+    # On the shared-AR path @streams_pgmq aliases @pgmq, so this reports the job
+    # pool's counters — accurate, since streams share that connection there.
+    def streams_pool_stats
+      streams_pool.stats.merge(pool_timeout: config.streams_pool_timeout)
+    rescue StandardError => e
+      Pgbus.logger.debug { "[Pgbus::Client] streams_pool_stats unavailable: #{e.class}: #{e.message}" }
+      {}
+    end
+
     def list_queues
       with_stale_connection_retry do
         synchronized { @pgmq.list_queues }
@@ -642,7 +751,47 @@ module Pgbus
     end
 
     def close
-      synchronized { @pgmq.close }
+      # Stop the publisher autoscale executor (if one was ever built) so its
+      # background thread doesn't leak (issue #323). Outside `synchronized` — it
+      # takes no pool lock and shutdown waits on a possibly-running check.
+      @streams_pool_trigger.shutdown if defined?(@streams_pool_trigger) && @streams_pool_trigger
+      synchronized do
+        @pgmq.close
+        # Close the CURRENT streams pool too (issue #315) so its connections
+        # don't leak. close_current reads the live (possibly hot-swapped, #323)
+        # pool once under the swap mutex and skips it when it aliases @pgmq (the
+        # shared-AR path) so we don't double-close the same pool.
+        @streams_pool.close_current(job_pool: @pgmq)
+      end
+    end
+
+    # Opt-in hot-swap of the dedicated streams pool to a new size (issue #323
+    # spike). Builds a fresh PGMQ::Client at new_size with the SAME bounds-applied
+    # connection options, atomically swaps the live reference, then drains +
+    # closes the old pool (bounded, never Thread#kill). NOT called automatically —
+    # there is no control loop here; a caller triggers it explicitly.
+    #
+    # No-op on the shared-AR (Proc) path (the streams pool aliases the job pool,
+    # which is non-thread-safe and forced to pool_size 1 — swapping it would
+    # corrupt the job pool), and no-op when the size is unchanged.
+    #
+    # @return [ResizablePool::SwapStats] on a swap, or {swapped: false, reason:}
+    def resize_streams_pool(new_size)
+      raise ArgumentError, "new_size must be a positive integer" unless new_size.is_a?(Integer) && new_size.positive?
+      return { swapped: false, reason: :shared_connection } if @shared_connection
+      return { swapped: false, reason: :unchanged } if streams_pool.stats[:size] == new_size
+
+      from_size = streams_pool.stats[:size]
+      new_pgmq = PGMQ::Client.new(
+        @streams_conn_opts, pool_size: new_size, pool_timeout: config.streams_pool_timeout
+      )
+      @streams_pool.swap(new_pgmq, from_size: from_size, to_size: new_size)
+    end
+
+    # Accumulated streams-pool swap telemetry (issue #323) — for the bench and a
+    # future control loop. Zero-valued before any swap.
+    def streams_swap_stats
+      @streams_pool.stats_snapshot
     end
 
     private
@@ -767,7 +916,14 @@ module Pgbus
                PG.connect(opts)
              when Hash
                owned = true
-               PG.connect(**opts)
+               # :variables is a database.yml convention, not a libpq keyword —
+               # strip it before PG.connect and apply the GUCs via SET so this
+               # raw bootstrap/DDL connection matches the pooled connections
+               # (issue #332). Empty/absent variables is a plain connect.
+               variables = opts[:variables]
+               conn = PG.connect(**opts.except(:variables))
+               variables&.each { |name, value| conn.exec("SET #{name} = '#{value}'") }
+               conn
              else
                raise ConfigurationError, "Cannot resolve raw PG connection from #{opts.class}"
              end
@@ -776,6 +932,22 @@ module Pgbus
       conn&.close if owned
     end
 
+    # Yields a PG connection from the dedicated streams pool for the streamer's
+    # replay reads (read_after / stream_current_msg_id / stream_oldest_msg_id).
+    # On the dedicated (String/Hash) path this checks out a persistent pooled
+    # connection — no fresh PG.connect per call (issue #315). On the shared-AR
+    # (Proc) path there is no separate pool (@streams_pgmq aliases @pgmq and
+    # points at the non-thread-safe AR raw_connection), so we fall back to
+    # with_raw_connection, which reuses that same shared connection. Callers
+    # already wrap this in `synchronized`, so the shared path stays serialized.
+    def with_streams_connection(&)
+      if @shared_connection
+        with_raw_connection(&)
+      else
+        streams_pool.with_connection(&)
+      end
+    end
+
     def ensure_single_queue(full_name)
       return if @queues_created[full_name]
 
@@ -871,6 +1043,40 @@ module Pgbus
       end
     end
 
+    # The ResizablePool wrapping the streams pool. All streams-pool reads
+    # (produce / with_connection / stats) go through it so the underlying
+    # PGMQ::Client can be atomically hot-swapped to a new size (issue #323).
+    attr_reader :streams_pool
+
+    # Lazily-built publisher-side autoscale trigger (issue #323). nil (a no-op)
+    # unless streams_pool_autoscale is on AND this is the dedicated-connection
+    # path (resize is a no-op on the shared-AR path). Built once — on the first
+    # publish, then reused. Runs its headroom query through the job pool (@pgmq),
+    # so it never competes with the streams pool it resizes.
+    #
+    # Double-checked under @streams_trigger_mutex so concurrent first-publishers
+    # (the integration spec fires 8) can't each build their own trigger — that
+    # would give each thread a trigger with its own throttle, defeating the
+    # once-per-interval throttle on the first window. Steady-state publishes hit
+    # the fast `defined?` path with no lock.
+    def streams_pool_trigger
+      return @streams_pool_trigger if defined?(@streams_pool_trigger)
+
+      @streams_trigger_mutex.synchronize do
+        return @streams_pool_trigger if defined?(@streams_pool_trigger)
+
+        @streams_pool_trigger =
+          if config.streams_pool_autoscale && !@shared_connection
+            autoscaler = Streams::PoolAutoscaler.new(client: self, config: config)
+            Streams::PoolTrigger.new(
+              autoscaler: autoscaler, job_pool: @pgmq,
+              interval: config.streams_pool_autoscale_interval,
+              application_name_prefix: config.streams_application_name
+            )
+          end
+      end
+    end
+
     # Substrings that indicate the pooled PG::Connection was already dead
     # *before* pgmq-ruby tried to use it — typically killed by a connection
     # pooler (PgBouncer server_idle_timeout / client_idle_timeout), an admin
@@ -1100,6 +1306,57 @@ module Pgbus
     # read_timeout is already failing — and keeps a single server-side mechanism.
     #
     # Returns conn_opts unchanged when read_timeout is nil (bounding disabled).
+    # Stamp a per-process application_name on the streams-pool connection options
+    # so the autoscaler can count peer processes via
+    # `pg_stat_activity.application_name LIKE '<prefix>_%'` (issue #323 P1). The
+    # suffix is the pid so DISTINCT application_name is an exact process count.
+    # application_name is a cosmetic session GUC — appending it can't break the
+    # connection. The Proc (shared-AR) path never reaches here.
+    def tag_application_name(conn_opts)
+      name = "#{config.streams_application_name}_#{::Process.pid}"
+      case conn_opts
+      when Hash
+        conn_opts.merge(application_name: name)
+      when String
+        # Two libpq string forms (mirrors #append_connection_bounds): URI form
+        # carries params as `?key=value&…` query pairs; key=value conninfo form
+        # is space-separated. Appending wins over any earlier application_name.
+        if conn_opts.start_with?("postgres://", "postgresql://")
+          separator = conn_opts.include?("?") ? "&" : "?"
+          "#{conn_opts}#{separator}application_name=#{name}"
+        else
+          "#{conn_opts} application_name=#{name}"
+        end
+      else
+        conn_opts
+      end
+    end
+
+    # In :session GUC mode, a Hash conn_opts carrying database.yml `:variables`
+    # must apply those GUCs via post-connect `SET` rather than the libpq
+    # `options` STARTUP param (which a transaction-mode PgBouncer rejects). We
+    # can't pass `:variables` to PG.connect (not a libpq keyword), so wrap the
+    # opts in a fresh-connect factory Proc: pgmq-ruby natively accepts a callable
+    # per pool slot (pgmq connection.rb), and it must return a UNIQUE
+    # PG::Connection each call (pgmq guards against a shared object). Applies to
+    # a Hash with `:variables` only — a String URL or no variables passes through
+    # unchanged, as does :options mode (where forward_connection_variables
+    # already baked the GUCs into `options`). See issue #332.
+    def wrap_session_gucs(conn_opts)
+      return conn_opts unless config.connection_guc_mode == :session
+      return conn_opts unless conn_opts.is_a?(Hash)
+
+      variables = conn_opts[:variables]
+      return conn_opts if variables.nil? || variables.empty?
+
+      pg_opts = conn_opts.except(:variables)
+      lambda do
+        conn = PG.connect(pg_opts)
+        variables.each { |name, value| conn.exec("SET #{name} = '#{value}'") }
+        conn
+      end
+    end
+
     def apply_connection_bounds(conn_opts)
       timeout = config.read_timeout
       return conn_opts unless timeout&.positive?