pgbus 0.0.1 → 0.1.2

data/lib/generators/pgbus/templates/upgrade_pgmq.rb.erb

@@ -0,0 +1,30 @@
+class UpgradePgmqToV<%= target_version_slug.camelize %> < ActiveRecord::Migration<%= migration_version %>
+  def up
+    # Drop existing PGMQ functions and types so they can be re-created
+    # at the new version. Tables (pgmq.meta, queue tables) are preserved.
+    execute Pgbus::PgmqSchema.drop_pgmq_functions_sql
+
+    # Re-create all PGMQ functions and types at the new version.
+    # This uses the vendored SQL which doesn't require the pgmq extension.
+    execute Pgbus::PgmqSchema.install_sql("<%= target_version %>")
+
+    # Record the upgrade
+    execute <<~SQL
+      CREATE TABLE IF NOT EXISTS pgbus_pgmq_schema_versions (
+        id SERIAL PRIMARY KEY,
+        version VARCHAR NOT NULL,
+        installed_at TIMESTAMP WITH TIME ZONE DEFAULT now() NOT NULL,
+        install_method VARCHAR NOT NULL DEFAULT 'embedded'
+      );
+
+      INSERT INTO pgbus_pgmq_schema_versions (version, install_method)
+      VALUES ('<%= target_version %>', 'upgrade');
+    SQL
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration,
+      "PGMQ schema downgrade is not supported. " \
+      "Restore from a database backup if needed."
+  end
+end

data/lib/generators/pgbus/upgrade_pgmq_generator.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module Pgbus
+  module Generators
+    class UpgradePgmqGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Upgrade PGMQ schema to the latest vendored version"
+
+      class_option :database,
+                   type: :string,
+                   default: nil,
+                   desc: "Use a separate database for pgbus tables (e.g. --database=pgbus)"
+
+      def create_migration
+        if separate_database?
+          migration_template "upgrade_pgmq.rb.erb",
+                             "db/pgbus_migrate/upgrade_pgmq_to_v#{target_version_slug}.rb"
+        else
+          migration_template "upgrade_pgmq.rb.erb",
+                             "db/migrate/upgrade_pgmq_to_v#{target_version_slug}.rb"
+        end
+      end
+
+      def display_post_upgrade
+        say ""
+        say "PGMQ upgrade migration created!", :green
+        say "  Target version: #{target_version}", :yellow
+        say ""
+        say "Next steps:"
+        say "  1. Review the migration in db/#{separate_database? ? "pgbus_migrate" : "migrate"}/"
+        say "  2. Run: rails db:migrate#{":#{options[:database]}" if separate_database?}"
+        say ""
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::Migration.current_version}]"
+      end
+
+      def target_version
+        Pgbus::PgmqSchema.latest_version
+      end
+
+      def target_version_slug
+        target_version.tr(".", "_")
+      end
+
+      def separate_database?
+        options[:database].present?
+      end
+    end
+  end
+end

data/lib/pgbus/active_job/adapter.rb

@@ -7,7 +7,7 @@ module Pgbus
     class Adapter
       def enqueue(active_job)
         queue = active_job.queue_name || Pgbus.configuration.default_queue
-        payload_hash = JSON.parse(Serializer.serialize_job(active_job))
+        payload_hash = Serializer.serialize_job_hash(active_job)
         payload_hash = Concurrency.inject_metadata(active_job, payload_hash)
         payload_hash = inject_batch_metadata(payload_hash)
 
@@ -16,7 +16,7 @@ module Pgbus
 
       def enqueue_at(active_job, timestamp)
         queue = active_job.queue_name || Pgbus.configuration.default_queue
-        payload_hash = JSON.parse(Serializer.serialize_job(active_job))
+        payload_hash = Serializer.serialize_job_hash(active_job)
         payload_hash = Concurrency.inject_metadata(active_job, payload_hash)
         payload_hash = inject_batch_metadata(payload_hash)
         delay = [(timestamp - Time.now.to_f).ceil, 0].max
@@ -88,7 +88,7 @@ module Pgbus
       def enqueue_immediate(queue, jobs)
         return if jobs.empty?
 
-        payloads = jobs.map { |j| JSON.parse(Serializer.serialize_job(j)) }
+        payloads = jobs.map { |j| Serializer.serialize_job_hash(j) }
         msg_ids = Pgbus.client.send_batch(queue, payloads)
 
         unless msg_ids.is_a?(Array) && msg_ids.size == jobs.size
@@ -104,6 +104,3 @@ module Pgbus
     end
   end
 end
-
-# Register the adapter with ActiveJob (register method added in Rails 7.2+)
-ActiveJob::QueueAdapters.register(:pgbus, Pgbus::ActiveJob::Adapter) if ActiveJob::QueueAdapters.respond_to?(:register)

data/lib/pgbus/active_job/executor.rb

@@ -21,12 +21,18 @@ module Pgbus
           return :dead_lettered
         end
 
-        job = ::ActiveJob::Base.deserialize(payload)
-        execute_job(job)
-        client.archive_message(queue_name, message.msg_id.to_i)
-        signal_concurrency(payload)
-        signal_batch_completed(payload)
-        instrument("pgbus.job_completed", queue: queue_name, job_class: payload["job_class"])
+        job_class = payload["job_class"]
+
+        job_succeeded = false
+
+        Instrumentation.instrument("pgbus.executor.execute", queue: queue_name, job_class: job_class) do
+          job = ::ActiveJob::Base.deserialize(payload)
+          execute_job(job)
+          client.archive_message(queue_name, message.msg_id.to_i)
+          job_succeeded = true
+        end
+
+        instrument("pgbus.job_completed", queue: queue_name, job_class: job_class)
         :success
       rescue StandardError => e
         handle_failure(message, queue_name, e)
@@ -34,6 +40,14 @@ module Pgbus
         # Don't signal concurrency on transient failure — the job will be retried.
         # Semaphore is released only on success or dead-lettering.
         :failed
+      ensure
+        # Signal concurrency and batch only when the job was archived successfully.
+        # job_succeeded is set AFTER archive_message, so if archive fails the
+        # semaphore slot stays held until VT expires and the job is retried.
+        if job_succeeded
+          signal_concurrency(payload)
+          signal_batch_completed(payload)
+        end
       end
 
       private
@@ -67,12 +81,12 @@ module Pgbus
         key = Concurrency.extract_key(payload)
         return unless key
 
-        Concurrency::Semaphore.release(key)
-
-        released = Concurrency::BlockedExecution.release_next(key)
-        return unless released
-
-        client.send_message(released[:queue_name], released[:payload])
+        # Atomic permit handoff: try to promote a blocked job first.
+        # promote_next wraps delete + enqueue in a transaction so neither is lost.
+        # If promoted, the slot stays occupied (no release needed).
+        # Only release the semaphore if there's nothing to promote.
+        promoted = Concurrency::BlockedExecution.promote_next(key, client: client)
+        Concurrency::Semaphore.release(key) unless promoted
       rescue StandardError => e
         Pgbus.logger.warn { "[Pgbus] Concurrency signal failed: #{e.message}" }
       end

data/lib/pgbus/batch.rb

@@ -29,124 +29,117 @@ module Pgbus
       self
     end
 
-    # Record a completed job. Returns the batch status after update.
+    # Record a completed job. Returns the batch row after update.
     def self.job_completed(batch_id)
       update_counter(batch_id, "completed_jobs")
     end
 
-    # Record a discarded (dead-lettered) job. Returns the batch status after update.
+    # Record a discarded (dead-lettered) job. Returns the batch row after update.
     def self.job_discarded(batch_id)
       update_counter(batch_id, "discarded_jobs")
     end
 
     # Find a batch record by ID. Returns a hash or nil.
     def self.find(batch_id)
-      return nil unless defined?(ActiveRecord::Base)
-
-      result = ActiveRecord::Base.connection.exec_query(
-        "SELECT * FROM pgbus_batches WHERE batch_id = $1",
-        "Pgbus Batch Find",
-        [batch_id]
-      )
-      result.first
+      BatchEntry.find_by(batch_id: batch_id)&.attributes
     end
 
     # Delete finished batches older than the given threshold.
     def self.cleanup(older_than:)
-      return 0 unless defined?(ActiveRecord::Base)
-
-      result = ActiveRecord::Base.connection.exec_query(
-        "DELETE FROM pgbus_batches WHERE status = 'finished' AND finished_at < $1 RETURNING id",
-        "Pgbus Batch Cleanup",
-        [older_than]
-      )
-      result.to_a.size
+      BatchEntry.stale(before: older_than).delete_all
     end
 
     private
 
     def create_record
-      return unless defined?(ActiveRecord::Base)
-
-      ActiveRecord::Base.connection.exec_query(
-        <<~SQL,
-          INSERT INTO pgbus_batches
-            (batch_id, description, on_finish_class, on_success_class, on_discard_class, properties, status)
-          VALUES ($1, $2, $3, $4, $5, $6, 'pending')
-        SQL
-        "Pgbus Batch Create",
-        [batch_id, description, on_finish&.name, on_success&.name, on_discard&.name, JSON.generate(properties)]
+      BatchEntry.create!(
+        batch_id: batch_id,
+        description: description,
+        on_finish_class: on_finish&.name,
+        on_success_class: on_success&.name,
+        on_discard_class: on_discard&.name,
+        properties: JSON.generate(properties),
+        status: "pending"
       )
     end
 
     def count_jobs(&)
+      previous_batch_id = Thread.current[:pgbus_batch_id]
+      previous_count = Thread.current[:pgbus_batch_job_count]
+
       Thread.current[:pgbus_batch_id] = batch_id
-      @job_count = 0
+      Thread.current[:pgbus_batch_job_count] = 0
 
       yield
 
       @job_count = Thread.current[:pgbus_batch_job_count] || 0
     ensure
-      Thread.current[:pgbus_batch_id] = nil
-      Thread.current[:pgbus_batch_job_count] = nil
+      Thread.current[:pgbus_batch_id] = previous_batch_id
+      Thread.current[:pgbus_batch_job_count] = previous_count
     end
 
     def update_total
-      return unless defined?(ActiveRecord::Base)
+      if @job_count.zero?
+        # Finish empty batches immediately — no jobs to signal completion
+        BatchEntry.where(batch_id: batch_id).update_all(
+          total_jobs: 0,
+          status: "finished",
+          finished_at: Time.current
+        )
+        fire_empty_batch_callbacks
+      else
+        BatchEntry.where(batch_id: batch_id).update_all(total_jobs: @job_count, status: "processing")
+      end
+    end
 
-      ActiveRecord::Base.connection.exec_query(
-        "UPDATE pgbus_batches SET total_jobs = $1, status = 'processing' WHERE batch_id = $2",
-        "Pgbus Batch Update Total",
-        [@job_count, batch_id]
-      )
+    def fire_empty_batch_callbacks
+      record = BatchEntry.find_by(batch_id: batch_id)
+      return unless record
+
+      properties = parse_properties(record.properties)
+      self.class.send(:enqueue_callback, record.on_finish_class, properties) if record.on_finish_class
+      self.class.send(:enqueue_callback, record.on_success_class, properties) if record.on_success_class
+    end
+
+    def parse_properties(props)
+      JSON.parse(props.presence || "{}")
+    rescue JSON::ParserError => e
+      Pgbus.logger.error { "[Pgbus] Invalid batch properties JSON: #{e.message}" }
+      {}
     end
 
     class << self
       private
 
       def update_counter(batch_id, column)
-        return nil unless defined?(ActiveRecord::Base)
-
-        result = ActiveRecord::Base.connection.exec_query(
-          <<~SQL,
-            UPDATE pgbus_batches
-            SET #{column} = #{column} + 1,
-                status = CASE
-                  WHEN completed_jobs + discarded_jobs + 1 = total_jobs THEN 'finished'
-                  ELSE status
-                END,
-                finished_at = CASE
-                  WHEN completed_jobs + discarded_jobs + 1 = total_jobs THEN NOW()
-                  ELSE finished_at
-                END
-            WHERE batch_id = $1
-            RETURNING status, total_jobs, completed_jobs, discarded_jobs, on_finish_class, on_success_class, on_discard_class, properties
-          SQL
-          "Pgbus Batch Counter",
-          [batch_id]
-        )
+        result = BatchEntry.increment_counter!(batch_id, column)
+        return nil unless result
 
-        row = result.first
-        return nil unless row
-
-        fire_callbacks(row) if row["status"] == "finished"
-        row
+        fire_callbacks(result[:record]) if result[:just_finished]
+        result
       end
 
-      def fire_callbacks(row)
-        properties = JSON.parse(row["properties"] || "{}")
-        all_succeeded = row["discarded_jobs"].to_i.zero?
-
-        enqueue_callback(row["on_finish_class"], properties) if row["on_finish_class"]
-        enqueue_callback(row["on_success_class"], properties) if row["on_success_class"] && all_succeeded
-        enqueue_callback(row["on_discard_class"], properties) if row["on_discard_class"] && !all_succeeded
+      def fire_callbacks(record)
+        properties = begin
+          JSON.parse(record.properties.presence || "{}")
+        rescue JSON::ParserError => e
+          Pgbus.logger.error { "[Pgbus] Invalid batch properties JSON: #{e.message}" }
+          {}
+        end
+        all_succeeded = record.discarded_jobs.zero?
+
+        enqueue_callback(record.on_finish_class, properties) if record.on_finish_class
+        enqueue_callback(record.on_success_class, properties) if record.on_success_class && all_succeeded
+        enqueue_callback(record.on_discard_class, properties) if record.on_discard_class && !all_succeeded
       end
 
       def enqueue_callback(class_name, properties)
-        job_class = class_name.constantize
+        job_class = class_name.safe_constantize
+        unless job_class && job_class < ::ActiveJob::Base
+          Pgbus.logger.error { "[Pgbus] Batch callback class invalid or not an ActiveJob: #{class_name}" }
+          return
+        end
         job_class.perform_later(properties)
-      rescue NameError => e
-        Pgbus.logger.error { "[Pgbus] Batch callback class not found: #{class_name}: #{e.message}" }
       end
     end
   end

data/lib/pgbus/cli.rb

@@ -33,24 +33,19 @@ module Pgbus
     end
 
     def show_status
-      if defined?(ActiveRecord::Base)
-        processes = ActiveRecord::Base.connection.execute(
-          "SELECT kind, hostname, pid, metadata, last_heartbeat_at FROM pgbus_processes ORDER BY kind, created_at"
-        )
+      processes = ProcessEntry.order(:kind, :created_at)
+                              .select(:kind, :hostname, :pid, :metadata, :last_heartbeat_at)
 
-        if processes.none?
-          puts "No Pgbus processes running."
-          return
-        end
+      if processes.none?
+        puts "No Pgbus processes running."
+        return
+      end
 
-        puts "KIND         HOST                 PID      HEARTBEAT                      METADATA"
-        puts "-" * 100
-        processes.each do |p|
-          puts format("%-12s %-20s %-8s %-30s %s",
-                      p["kind"], p["hostname"], p["pid"], p["last_heartbeat_at"], p["metadata"])
-        end
-      else
-        puts "ActiveRecord not available. Run from a Rails context."
+      puts "KIND         HOST                 PID      HEARTBEAT                      METADATA"
+      puts "-" * 100
+      processes.each do |p|
+        puts format("%-12s %-20s %-8s %-30s %s",
+                    p.kind, p.hostname, p.pid, p.last_heartbeat_at, p.metadata)
       end
     end
 

data/lib/pgbus/client.rb

@@ -6,61 +6,78 @@ module Pgbus
   class Client
     attr_reader :pgmq, :config
 
+    PGMQ_REQUIRE_MUTEX = Mutex.new
+    private_constant :PGMQ_REQUIRE_MUTEX
+
     def initialize(config = Pgbus.configuration)
-      require "pgmq-ruby"
+      # Define the PGMQ module before requiring the gem so that Zeitwerk's
+      # eager_load (called inside pgmq.rb) can resolve the constant.
+      # Without this, Ruby 4.0 + Zeitwerk 2.7.5 raises NameError because
+      # eager_load runs const_get(:Client) on PGMQ before the module is defined.
+      PGMQ_REQUIRE_MUTEX.synchronize do
+        Object.const_set(:PGMQ, Module.new) unless defined?(::PGMQ)
+        require "pgmq"
+      end
       @config = config
       @pgmq = PGMQ::Client.new(
         config.connection_options,
         pool_size: config.pool_size,
         pool_timeout: config.pool_timeout
       )
-      @queues_created = {}
-      @mutex = Mutex.new
+      @queues_created = Concurrent::Map.new
     end
 
     def ensure_queue(name)
       full_name = config.queue_name(name)
-      @mutex.synchronize do
-        return if @queues_created[full_name]
+      return if @queues_created[full_name]
 
+      @queues_created.compute_if_absent(full_name) do
         @pgmq.create(full_name)
         @pgmq.enable_notify_insert(full_name, throttle_interval_ms: config.notify_throttle_ms) if config.listen_notify
-        @queues_created[full_name] = true
+        true
       end
     end
 
     def ensure_dead_letter_queue(name)
       dlq_name = config.dead_letter_queue_name(name)
-      @mutex.synchronize do
-        return if @queues_created[dlq_name]
+      return if @queues_created[dlq_name]
 
+      @queues_created.compute_if_absent(dlq_name) do
         @pgmq.create(dlq_name)
-        @queues_created[dlq_name] = true
+        true
       end
     end
 
     def send_message(queue_name, payload, headers: nil, delay: 0)
       full_name = config.queue_name(queue_name)
       ensure_queue(queue_name)
-      @pgmq.produce(full_name, serialize(payload), headers: headers && serialize(headers), delay: delay)
+      Instrumentation.instrument("pgbus.client.send_message", queue: full_name) do
+        @pgmq.produce(full_name, serialize(payload), headers: headers && serialize(headers), delay: delay)
+      end
     end
 
     def send_batch(queue_name, payloads, headers: nil, delay: 0)
       full_name = config.queue_name(queue_name)
       ensure_queue(queue_name)
-      serialized = payloads.map { |p| serialize(p) }
-      serialized_headers = headers&.map { |h| serialize(h) }
-      @pgmq.produce_batch(full_name, serialized, headers: serialized_headers, delay: delay)
+      Instrumentation.instrument("pgbus.client.send_batch", queue: full_name, size: payloads.size) do
+        serialized = payloads.map { |p| serialize(p) }
+        serialized_headers = headers&.map { |h| serialize(h) }
+        @pgmq.produce_batch(full_name, serialized, headers: serialized_headers, delay: delay)
+      end
     end
 
     def read_message(queue_name, vt: nil)
       full_name = config.queue_name(queue_name)
-      @pgmq.read(full_name, vt: vt || config.visibility_timeout)
+      Instrumentation.instrument("pgbus.client.read_message", queue: full_name) do
+        @pgmq.read(full_name, vt: vt || config.visibility_timeout)
+      end
     end
 
     def read_batch(queue_name, qty:, vt: nil)
       full_name = config.queue_name(queue_name)
-      @pgmq.read_batch(full_name, vt: vt || config.visibility_timeout, qty: qty)
+      Instrumentation.instrument("pgbus.client.read_batch", queue: full_name, qty: qty) do
+        @pgmq.read_batch(full_name, vt: vt || config.visibility_timeout, qty: qty)
+      end
     end
 
     def read_with_poll(queue_name, qty:, vt: nil, max_poll_seconds: 5, poll_interval_ms: 100)

data/lib/pgbus/concurrency/blocked_execution.rb

@@ -8,65 +8,60 @@ module Pgbus
       class << self
         # Insert a blocked execution for a job that hit the concurrency limit.
         def insert(concurrency_key:, queue_name:, payload:, duration:, priority: 0)
-          expires_at = Time.now.utc + duration
-
-          execute(<<~SQL, "Pgbus Blocked Insert", [concurrency_key, queue_name, JSON.generate(payload), priority, expires_at])
-            INSERT INTO pgbus_blocked_executions
-              (concurrency_key, queue_name, payload, priority, expires_at)
-            VALUES ($1, $2, $3, $4, $5)
-          SQL
+          Pgbus::BlockedExecution.create!(
+            concurrency_key: concurrency_key,
+            queue_name: queue_name,
+            payload: JSON.generate(payload),
+            priority: priority,
+            expires_at: Time.now.utc + duration
+          )
         end
 
         # Release the next blocked execution for a given concurrency key.
         # Returns the released row (queue_name, payload) or nil if none.
         def release_next(concurrency_key)
-          return nil unless defined?(ActiveRecord::Base)
+          Pgbus::BlockedExecution.release_next!(concurrency_key)
+        end
 
-          result = execute(<<~SQL, "Pgbus Blocked Release", [concurrency_key])
-            DELETE FROM pgbus_blocked_executions
-            WHERE id = (
-              SELECT id FROM pgbus_blocked_executions
-              WHERE concurrency_key = $1
-              ORDER BY priority ASC, created_at ASC
-              LIMIT 1
-              FOR UPDATE SKIP LOCKED
-            )
-            RETURNING queue_name, payload
-          SQL
+        # Atomically promote the next blocked execution: delete the row and enqueue
+        # the job in a single transaction. Returns true if a job was promoted, false
+        # otherwise. This avoids losing a blocked row if enqueue fails.
+        def promote_next(concurrency_key, client:, delay: 0)
+          released = nil
+          Pgbus::BlockedExecution.transaction do
+            released = release_next(concurrency_key)
+            raise ActiveRecord::Rollback unless released
 
-          row = result.first
-          return nil unless row
+            actual_delay = resolve_delay(released[:payload], delay)
+            client.send_message(released[:queue_name], released[:payload], delay: actual_delay)
+          end
 
-          { queue_name: row["queue_name"], payload: JSON.parse(row["payload"]) }
+          !!released
+        rescue StandardError => e
+          Pgbus.logger.warn { "[Pgbus] Promote blocked execution failed for #{concurrency_key}: #{e.message}" }
+          false
         end
 
         # Delete blocked executions that have expired.
         # Returns the count of deleted rows.
         def expire_stale
-          result = execute(<<~SQL, "Pgbus Blocked Expire", [Time.now.utc])
-            DELETE FROM pgbus_blocked_executions
-            WHERE expires_at < $1
-            RETURNING id
-          SQL
-
-          result.to_a.size
+          Pgbus::BlockedExecution.expired(Time.now.utc).delete_all
         end
 
         # Count blocked executions for a given key. Useful for testing/monitoring.
         def count_for(concurrency_key)
-          result = execute(<<~SQL, "Pgbus Blocked Count", [concurrency_key])
-            SELECT COUNT(*) AS cnt FROM pgbus_blocked_executions WHERE concurrency_key = $1
-          SQL
-
-          result.first&.fetch("cnt", 0).to_i
+          Pgbus::BlockedExecution.where(concurrency_key: concurrency_key).count
         end
 
         private
 
-        def execute(sql, name, binds)
-          return [] unless defined?(ActiveRecord::Base)
+        def resolve_delay(payload, default_delay)
+          scheduled_at = payload["scheduled_at"]
+          return default_delay unless scheduled_at
 
-          ActiveRecord::Base.connection.exec_query(sql, name, binds)
+          [Time.parse(scheduled_at).to_f - Time.now.to_f, 0].max.ceil
+        rescue StandardError
+          default_delay
         end
       end
     end