pgbus 0.21.3 → 0.22.0

data/lib/pgbus/pgmq_schema.rb

@@ -16,6 +16,23 @@ module Pgbus
     SCHEMA_DIR = File.expand_path("pgmq_schema", __dir__).freeze
     FIXUPS_DIR = File.join(SCHEMA_DIR, "fixups").freeze
 
+    # Wrapped around the vendored SQL by install_sql; see there.
+    SAVE_SEARCH_PATH_SQL = <<~SQL
+      DO $pgbus_save_path$
+      BEGIN
+        PERFORM pg_catalog.set_config('pgbus.install_saved_search_path', pg_catalog.current_setting('search_path'), true);
+      END
+      $pgbus_save_path$;
+    SQL
+
+    RESTORE_SEARCH_PATH_SQL = <<~SQL
+      DO $pgbus_restore_path$
+      BEGIN
+        PERFORM pg_catalog.set_config('search_path', pg_catalog.current_setting('pgbus.install_saved_search_path'), true);
+      END
+      $pgbus_restore_path$;
+    SQL
+
     class << self
       # Returns the latest vendored PGMQ version string.
       def latest_version
@@ -78,11 +95,21 @@ module Pgbus
       # Strips the extension-only pg_dump config blocks since they're
       # irrelevant when not installed as an extension.
       #
+      # From 1.13.1 upstream pins `search_path = pg_catalog, pg_temp` for the
+      # install with a transaction-local set_config, which outlives the install
+      # inside a migration (or a caller's transaction, in the client's lazy
+      # install): the next unqualified CREATE TABLE would target pg_catalog.
+      # The caller's path is saved before upstream's text and restored after,
+      # both transaction-local, so the install's pinning still applies to its
+      # own statements and nothing after them. Run the result as one string:
+      # a statement-by-statement load would lose the saved path between
+      # statements.
+      #
       # @param version [String] defaults to latest
       # @return [String] SQL
       def install_sql(version = latest_version)
-        sql = sql_for_version(version)
-        strip_extension_only_blocks(sql)
+        sql = strip_extension_only_blocks(sql_for_version(version))
+        "#{SAVE_SEARCH_PATH_SQL}\n#{sql}\n#{RESTORE_SEARCH_PATH_SQL}"
       end
 
       # Returns SQL to create the version tracking table and record an installation.
@@ -231,9 +258,12 @@ module Pgbus
           ""
         )
 
-        # Remove _belongs_to_pgmq function (checks pg_depend on extension)
+        # Remove _belongs_to_pgmq function (checks pg_depend on extension).
+        # From 1.13.1 its terminator carries `SET search_path = ...`; ending
+        # the match at a bare `LANGUAGE plpgsql;` would run on to the next
+        # function without that clause and strip everything in between.
         sql.gsub(
-          /CREATE FUNCTION pgmq\._belongs_to_pgmq.*?LANGUAGE plpgsql;/m,
+          /CREATE FUNCTION pgmq\._belongs_to_pgmq.*?\$\$\s*LANGUAGE plpgsql[^;]*;/m,
           ""
         )
       end

data/lib/pgbus/process/completer.rb

@@ -0,0 +1,187 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Process
+    # Batched completion (issue #606, `config.completion_batching`): one per
+    # worker, it archives the messages of jobs that finish close together with
+    # one `archive_batch` per physical queue instead of one `pgmq.archive` each.
+    #
+    # There is no timer. An idle completer flushes a lone id at once; ids that
+    # arrive while a flush is in flight wait for the next one, so under load a
+    # batch is "what arrived during the previous round trip".
+    # `batch_size` only caps the ids per statement.
+    #
+    # The job thread waits for its own result (`complete` blocks), so the
+    # executor's post-archive sequence runs on it as before. Threads only: a
+    # blocking wait would stall an Async reactor (Configuration#validate!).
+    #
+    # completed_jobs: :delete (the config setting, issue #607) flushes with
+    # `delete_batch` instead; the results keep their names, so :archived
+    # means "completed" and :already_archived "already gone".
+    class Completer
+      # Raised on a waiting job thread whose flush can no longer happen: the
+      # flusher was stopped past its join budget, or died.
+      class Stopped < Pgbus::Error; end
+
+      # One job thread's wait for its id's result.
+      class Waiter
+        def initialize
+          @mutex = Mutex.new
+          @done = ConditionVariable.new
+          @state = nil
+          @value = nil
+        end
+
+        def resolve(value) = settle(:value, value)
+
+        def fail(error) = settle(:error, error)
+
+        def value
+          @mutex.synchronize do
+            @done.wait(@mutex) until @state
+            raise @value if @state == :error
+
+            @value
+          end
+        end
+
+        private
+
+        # First settle wins: a waiter already served is not failed again by
+        # the flusher's last-resort cleanup.
+        def settle(state, value)
+          @mutex.synchronize do
+            next if @state
+
+            @state = state
+            @value = value
+            @done.broadcast
+          end
+        end
+      end
+
+      BATCH_OPS = { archive: :archive_batch, delete: :delete_batch }.freeze
+
+      def initialize(client:, batch_size:, completed_jobs: :archive)
+        @batch_op = BATCH_OPS.fetch(completed_jobs) do
+          raise ArgumentError, "completed_jobs must be :archive or :delete (got #{completed_jobs.inspect})"
+        end
+        @client = client
+        @batch_size = batch_size
+        @mutex = Mutex.new
+        @wake = ConditionVariable.new
+        @pending = {}
+        @taken = nil
+        @stopping = false
+        # Owned by this completer: one per worker, built with it. #stop asks
+        # it to drain and exit, then joins it (Worker#shutdown, #immediate_shutdown).
+        @thread = Thread.new { run } # rubocop:disable ThreadSafety/NewThread
+        @thread.name = "pgbus-completer"
+      end
+
+      # Archives msg_id from the physical (already prefixed) queue and returns
+      # :archived, or :already_archived when the message was gone (another
+      # worker archived it). Raises the flush's exception on this thread.
+      def complete(queue, msg_id)
+        waiter = @mutex.synchronize do
+          next if @stopping
+
+          entry = Waiter.new
+          (@pending[queue] ||= []) << [msg_id, entry]
+          @wake.signal
+          entry
+        end
+        return waiter.value if waiter
+
+        # Stopped: complete on the caller's thread, with the same results.
+        @client.public_send(@batch_op, queue, [msg_id], prefixed: false).empty? ? :already_archived : :archived
+      end
+
+      # Flushes whatever is pending, then joins the flusher. Past timeout the
+      # flusher is killed and every waiter it had not served raises Stopped.
+      # Later `complete` calls archive on their own thread. Idempotent.
+      def stop(timeout: nil)
+        @mutex.synchronize do
+          @stopping = true
+          @wake.signal
+        end
+        return if @thread.join(timeout)
+
+        @thread.kill
+        @thread.join
+      end
+
+      def alive? = @thread.alive?
+
+      # Ids waiting for the next flush (test and introspection seam).
+      def pending_count
+        @mutex.synchronize { @pending.sum { |_queue, entries| entries.size } }
+      end
+
+      # True while a flush is running (test and introspection seam).
+      def flushing?
+        @mutex.synchronize { !@taken.nil? }
+      end
+
+      private
+
+      def run
+        while (batch = take)
+          flush(batch)
+          @mutex.synchronize { @taken = nil }
+        end
+      ensure
+        fail_unserved
+      end
+
+      # Waits for pending ids and takes all of them; nil once stopping with
+      # nothing left.
+      def take
+        @mutex.synchronize do
+          @wake.wait(@mutex) while @pending.empty? && !@stopping
+          next if @pending.empty?
+
+          @taken = @pending
+          @pending = {}
+          @taken
+        end
+      end
+
+      def flush(batch)
+        batch.each do |queue, entries|
+          entries.each_slice(@batch_size) { |slice| flush_slice(queue, slice) }
+        end
+      end
+
+      # pgmq returns one row per archived id (as a String); an id already gone
+      # is absent. A duplicated id (the same message held by two threads after
+      # a redelivery) is archived once, so only its first waiter gets :archived.
+      def flush_slice(queue, slice)
+        archived = @client.public_send(@batch_op, queue, slice.map(&:first).uniq, prefixed: false).to_set(&:to_i)
+        slice.each { |msg_id, waiter| waiter.resolve(archived.delete?(msg_id) ? :archived : :already_archived) }
+      rescue Exception => e # rubocop:disable Lint/RescueException -- delivered to the waiters, never swallowed
+        log_flush_failure(queue, slice, e) unless e.is_a?(StandardError)
+        slice.map(&:last).each { |waiter| waiter.fail(e) }
+      end
+
+      def log_flush_failure(queue, slice, error)
+        Pgbus.logger.error do
+          "[Pgbus::Completer] #{@batch_op} on #{queue} (#{slice.size} ids) raised #{error.class}: #{error.message}"
+        end
+      end
+
+      # The flusher is exiting: nothing may be left waiting on it.
+      def fail_unserved
+        leftovers = @mutex.synchronize do
+          entries = [@taken, @pending].compact.flat_map(&:values).flatten(1)
+          @taken = nil
+          @pending = {}
+          @stopping = true
+          entries
+        end
+        error = Stopped.new("the completer stopped before archiving this message")
+        leftovers.map(&:last).each { |waiter| waiter.fail(error) }
+      end
+    end
+  end
+end

data/lib/pgbus/process/consumer.rb

@@ -264,7 +264,7 @@ module Pgbus
           dispatch(handlers, message, queue_name)
         end
 
-        Pgbus.client.archive_message(queue_name, message.msg_id.to_i)
+        archive_handled(queue_name, message)
         # A redelivery that succeeds leaves no stale failure behind. First
         # deliveries never had one: no query on the hot path (#484's gate).
         if message.read_ct.to_i > 1
@@ -367,6 +367,20 @@ module Pgbus
       # as retrying with its error. No ErrorReporter call: Handler#process!
       # already instruments pgbus.event_failed, which feeds the reporters.
       # A handler queue has no priority sub-queues: its one table is the source.
+      # The handlers ran. An archive whose reply was lost after it may have
+      # committed, and whose re-run found the message gone (issue #607
+      # review), leaves nothing to retry: the message is handled.
+      def archive_handled(queue_name, message)
+        Pgbus.client.archive_message(queue_name, message.msg_id.to_i)
+      rescue StandardError => e # pgmq loads lazily: completion_lost? checks the class
+        raise unless Client::NativeStatements.completion_lost?(e)
+
+        Pgbus.logger.warn do
+          "[Pgbus] Consumer archive of #{message.msg_id} on #{queue_name} lost its reply and the message is gone: " \
+            "counted as handled (#{e.message})"
+        end
+      end
+
       def record_failed_event(message, queue_name, error)
         FailedEventRecorder.record!(
           queue_name: queue_name, msg_id: message.msg_id.to_i, source_queue: config.queue_name(queue_name),

data/lib/pgbus/process/worker.rb

@@ -62,6 +62,12 @@ module Pgbus
                 "fair_share and group_mode are mutually exclusive — this worker has group_mode " \
                 "#{(@group_mode || config.group_mode).inspect} while config.fair_share is set"
         end
+        # The completer blocks a job thread on its flush, which would stall an
+        # Async reactor (issue #606); the mode can arrive here without validate!.
+        if config.completion_batching && @execution_mode == :async
+          raise Pgbus::ConfigurationError,
+                "completion_batching supports execution_mode :threads only; this worker runs :async"
+        end
         @single_active_consumer = single_active_consumer
         @consumer_priority = consumer_priority
         @lifecycle = Lifecycle.new
@@ -94,7 +100,16 @@ module Pgbus
           else
             stat_buffer
           end
-        @executor = Pgbus::ActiveJob::Executor.new(stat_buffer: @stat_buffer)
+        # One per worker when config.completion_batching (issue #606): archives
+        # the jobs that finish close together in one archive_batch, or deletes
+        # them with delete_batch under config.completed_jobs = :delete (#607).
+        @completer = if config.completion_batching
+                       Completer.new(client: Pgbus.client, batch_size: config.completion_batch_size,
+                                     completed_jobs: config.completed_jobs)
+                     end
+        # The worker's config, so the executor and the completer read one
+        # completed_jobs (issue #607).
+        @executor = Pgbus::ActiveJob::Executor.new(config: config, stat_buffer: @stat_buffer, completer: @completer)
         @wake_signal = wake_signal || WakeSignal.new
         @pool = ExecutionPools.build(
           mode: @execution_mode,
@@ -239,6 +254,13 @@ module Pgbus
         Pgbus.stopping = true
         @lifecycle.transition_to!(:stopped)
         @wake_signal.notify!
+        # Before the pool: a job thread killed while it waits on a flush would
+        # skip its signals while the flush still commits (archived, never
+        # signalled). Stopping first resolves every waiter; a job that finishes
+        # inside the grace then archives on its own thread. Bounded by the same
+        # grace: a flush wedged on a hung Postgres is abandoned, and its jobs
+        # fail with Completer::Stopped and re-run after visibility_timeout.
+        @completer&.stop(timeout: config.immediate_shutdown_grace)
         ExecutionPools.terminate(@pool, grace: config.immediate_shutdown_grace)
       end
 
@@ -927,7 +949,13 @@ module Pgbus
         # config.drain_timeout for in-flight jobs. A job still running here has
         # proven it won't finish; waiting another full window would push the
         # worker past the supervisor's shutdown_timeout deadline (issue #386).
+        residual_deadline = monotonic_now + POOL_TERMINATION_WAIT
         @pool.wait_for_termination(POOL_TERMINATION_WAIT)
+        # After the pool: every waiter is a job thread, so the drain already
+        # flushed theirs; this flushes any left and joins the flusher. It shares
+        # the pool's residual budget rather than adding its own, so shutdown
+        # stays inside the supervisor's drain_timeout + 5 default deadline.
+        @completer&.stop(timeout: [residual_deadline - monotonic_now, 0].max)
         VisibilityHeartbeat.stop
         @stat_buffer&.stop
         @queue_lock&.unlock_all

data/lib/pgbus/serializer.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "time"
 
 module Pgbus
   module Serializer

data/lib/pgbus/version.rb

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

data/lib/pgbus/visibility_heartbeat.rb

@@ -33,8 +33,13 @@ module Pgbus
     # the whole table is bounded by the execution pool's capacity plus the
     # read-ahead buffer (threads + read_ahead per process, issue #486) — a
     # handful of entries per process, not one per enqueued job.
+    #
+    # `key` is the entry's @entries key, [queue_name, msg_id], built once with
+    # the entry (#track, #hold): #register, #unregister and #registered? each
+    # used to build their own (issue #604). A tenth member stays in the
+    # 160-byte slot.
     Entry = Struct.new(:client, :queue_name, :prefixed, :msg_id, :job_class, :extended_at, :extensions,
-                       :concurrency, :on_beat, keyword_init: true)
+                       :concurrency, :on_beat, :key, keyword_init: true)
 
     # Per-job opt-out, included on ActiveJob::Base by the engine:
     #
@@ -81,9 +86,10 @@ module Pgbus
                 concurrency: nil, on_beat: nil)
         return yield unless config.visibility_heartbeat
 
-        entry = Entry.new(client: client, queue_name: queue_name, prefixed: prefixed, msg_id: msg_id.to_i,
+        msg_id = msg_id.to_i
+        entry = Entry.new(client: client, queue_name: queue_name, prefixed: prefixed, msg_id: msg_id,
                           job_class: job_class, extended_at: monotonic_now, extensions: 0,
-                          concurrency: concurrency, on_beat: on_beat)
+                          concurrency: concurrency, on_beat: on_beat, key: [queue_name, msg_id].freeze)
         register(entry, config)
         begin
           yield
@@ -100,8 +106,10 @@ module Pgbus
       def hold(client:, queue_name:, msg_id:, prefixed: true, job_class: nil, config: Pgbus.configuration)
         return unless config.visibility_heartbeat
 
-        entry = Entry.new(client: client, queue_name: queue_name, prefixed: prefixed, msg_id: msg_id.to_i,
-                          job_class: job_class, extended_at: monotonic_now, extensions: 0)
+        msg_id = msg_id.to_i
+        entry = Entry.new(client: client, queue_name: queue_name, prefixed: prefixed, msg_id: msg_id,
+                          job_class: job_class, extended_at: monotonic_now, extensions: 0,
+                          key: [queue_name, msg_id].freeze)
         register(entry, config)
         entry
       end
@@ -159,7 +167,7 @@ module Pgbus
             @running = false
             @thread = nil
           end
-          @entries[key_for(entry)] = entry
+          @entries[entry.key] = entry
           # Start the ticker once per process.
           unless @running && @thread&.alive?
             @running = true
@@ -173,7 +181,7 @@ module Pgbus
       # job's entry (issue #486).
       def unregister(entry)
         synchronize do
-          key = key_for(entry)
+          key = entry.key
           @entries.delete(key) if @entries[key].equal?(entry)
         end
       end
@@ -265,11 +273,7 @@ module Pgbus
       end
 
       def registered?(entry)
-        synchronize { @entries[key_for(entry)].equal?(entry) }
-      end
-
-      def key_for(entry)
-        [entry.queue_name, entry.msg_id]
+        synchronize { @entries[entry.key].equal?(entry) }
       end
 
       # A literal block, not `(&)`: ThreadSafety/ClassInstanceVariable accepts

data/sig/pgbus/client.rbs

@@ -33,6 +33,7 @@ module Pgbus
     def message_with_job_id?: (untyped queue_name, job_id: untyped, ?prefixed: untyped) -> bool?
     def metrics: (?untyped queue_name) -> untyped
     def move_to_dead_letter: (untyped queue_name, untyped message, ?headers: untyped) -> untyped
+    def native_ready?: (String full_name) -> bool
     def notify_enabled?: (untyped queue_name) -> bool
     def notify_stream: (untyped stream_name, untyped payload) -> untyped
     def oldest_claimable_ages: (?untyped queue_name) -> untyped

data/sig/pgbus/configuration.rbs

@@ -30,6 +30,12 @@ module Pgbus
     def capsule_named: (untyped name) -> untyped
     def circuit_breaker_enabled: () -> bool
     def circuit_breaker_enabled=: (untyped) -> untyped
+    def completed_jobs: () -> (:archive | :delete)
+    def completed_jobs=: (untyped value) -> (:archive | :delete)
+    def completion_batch_size: () -> Integer
+    def completion_batch_size=: (untyped) -> untyped
+    def completion_batching: () -> bool
+    def completion_batching=: (untyped) -> untyped
     def connection_guc_mode: () -> Symbol
     def connection_guc_mode=: (untyped mode) -> untyped
     def connection_options: () -> untyped
@@ -114,6 +120,8 @@ module Pgbus
     def metrics_backend=: (untyped) -> untyped
     def metrics_enabled: () -> bool
     def metrics_enabled=: (untyped) -> untyped
+    def native_statements: () -> bool
+    def native_statements=: (bool value) -> bool
     def notify_presence_change: (untyped stream_name, untyped change) -> untyped
     def outbox_batch_size: () -> Integer
     def outbox_batch_size=: (untyped) -> untyped

data/sig/pgbus/process/completer.rbs

@@ -0,0 +1,26 @@
+module Pgbus
+  module Process
+    # Batched completion (issue #606): one per worker when
+    # config.completion_batching is on.
+    class Completer
+      class Stopped < Pgbus::Error
+      end
+
+      class Waiter
+        def initialize: () -> void
+        def resolve: (Symbol value) -> void
+        def fail: (Exception error) -> void
+        def value: () -> Symbol
+      end
+
+      BATCH_OPS: Hash[Symbol, Symbol]
+
+      def initialize: (client: Pgbus::Client, batch_size: Integer, ?completed_jobs: :archive | :delete) -> void
+      def complete: (String queue, Integer msg_id) -> (:archived | :already_archived)
+      def stop: (?timeout: Numeric?) -> void
+      def alive?: () -> bool
+      def pending_count: () -> Integer
+      def flushing?: () -> bool
+    end
+  end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pgbus
 version: !ruby/object:Gem::Version
-  version: 0.21.3
+  version: 0.22.0
 platform: ruby
 authors:
 - Mikael Henriksson
@@ -356,6 +356,7 @@ files:
 - lib/pgbus/client/connection_health.rb
 - lib/pgbus/client/ensure_stream_queue.rb
 - lib/pgbus/client/fair_read.rb
+- lib/pgbus/client/native_statements.rb
 - lib/pgbus/client/notify_lock_retry.rb
 - lib/pgbus/client/notify_stream.rb
 - lib/pgbus/client/read_after.rb
@@ -456,11 +457,14 @@ files:
 - lib/pgbus/outbox/poller.rb
 - lib/pgbus/pgmq_schema.rb
 - lib/pgbus/pgmq_schema/fixups/pgmq_v1.13.0.sql
+- lib/pgbus/pgmq_schema/fixups/pgmq_v1.13.1.sql
 - lib/pgbus/pgmq_schema/pgmq_v1.11.0.sql
 - lib/pgbus/pgmq_schema/pgmq_v1.11.1.sql
 - lib/pgbus/pgmq_schema/pgmq_v1.12.0.sql
 - lib/pgbus/pgmq_schema/pgmq_v1.13.0.sql
+- lib/pgbus/pgmq_schema/pgmq_v1.13.1.sql
 - lib/pgbus/process/claim_buffer.rb
+- lib/pgbus/process/completer.rb
 - lib/pgbus/process/consumer.rb
 - lib/pgbus/process/consumer_priority.rb
 - lib/pgbus/process/dispatcher.rb
@@ -614,6 +618,7 @@ files:
 - sig/pgbus/client.rbs
 - sig/pgbus/configuration.rbs
 - sig/pgbus/configuration/snapshot.rbs
+- sig/pgbus/process/completer.rbs
 - sig/pgbus/streams.rbs
 - sig/pgbus/streams/stream.rbs
 homepage: https://github.com/zoolutions/pgbus