pgbus 0.14.0 → 0.14.2

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

@@ -0,0 +1,13 @@
+class AddPgbusBatchCallbackJobs < ActiveRecord::Migration<%= migration_version %>
+  def up
+    add_column :pgbus_batches, :on_finish_job, :jsonb unless column_exists?(:pgbus_batches, :on_finish_job)
+    add_column :pgbus_batches, :on_success_job, :jsonb unless column_exists?(:pgbus_batches, :on_success_job)
+    add_column :pgbus_batches, :on_failure_job, :jsonb unless column_exists?(:pgbus_batches, :on_failure_job)
+  end
+
+  def down
+    remove_column :pgbus_batches, :on_finish_job if column_exists?(:pgbus_batches, :on_finish_job)
+    remove_column :pgbus_batches, :on_success_job if column_exists?(:pgbus_batches, :on_success_job)
+    remove_column :pgbus_batches, :on_failure_job if column_exists?(:pgbus_batches, :on_failure_job)
+  end
+end

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

@@ -0,0 +1,50 @@
+class AddPgbusBatchExecutions < ActiveRecord::Migration<%= migration_version %>
+  def up
+    create_table :pgbus_batch_executions do |t|
+      t.string :batch_id, null: false
+      t.string :job_id, null: false
+      t.bigint :msg_id
+      t.string :queue_name
+      t.datetime :created_at, null: false, default: -> { "CURRENT_TIMESTAMP" }
+    end
+
+    add_index :pgbus_batch_executions, :job_id,
+      unique: true, name: "idx_pgbus_batch_executions_job_id"
+    add_index :pgbus_batch_executions, :batch_id,
+      name: "idx_pgbus_batch_executions_batch_id"
+    add_index :pgbus_batch_executions, :created_at,
+      where: "msg_id IS NULL",
+      name: "idx_pgbus_batch_executions_orphans"
+    add_foreign_key :pgbus_batch_executions, :pgbus_batches,
+      column: :batch_id, primary_key: :batch_id, on_delete: :cascade
+
+    if column_exists?(:pgbus_batches, :on_discard_class) && !column_exists?(:pgbus_batches, :on_failure_class)
+      rename_column :pgbus_batches, :on_discard_class, :on_failure_class
+    end
+
+    # Drain/restart pgbus workers after this migration: a process that still
+    # writes discarded_jobs will error once the column is gone.
+    if column_exists?(:pgbus_batches, :discarded_jobs)
+      if column_exists?(:pgbus_batches, :failed_jobs)
+        execute "UPDATE pgbus_batches SET failed_jobs = failed_jobs + discarded_jobs"
+        remove_column :pgbus_batches, :discarded_jobs
+      else
+        rename_column :pgbus_batches, :discarded_jobs, :failed_jobs
+      end
+    end
+  end
+
+  def down
+    unless column_exists?(:pgbus_batches, :discarded_jobs)
+      add_column :pgbus_batches, :discarded_jobs, :integer, null: false, default: 0
+      if column_exists?(:pgbus_batches, :failed_jobs)
+        execute "UPDATE pgbus_batches SET discarded_jobs = failed_jobs"
+        execute "UPDATE pgbus_batches SET failed_jobs = 0"
+      end
+    end
+    if column_exists?(:pgbus_batches, :on_failure_class) && !column_exists?(:pgbus_batches, :on_discard_class)
+      rename_column :pgbus_batches, :on_failure_class, :on_discard_class
+    end
+    drop_table :pgbus_batch_executions if table_exists?(:pgbus_batch_executions)
+  end
+end

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

@@ -44,6 +44,8 @@ Pgbus.configure do |c|
 
   # --- Dispatcher (maintenance tasks) ------------------------------------
   # c.dispatch_interval = 1.0
+  # c.batch_retention = 7.days          # finished-batch cleanup; nil disables
+  # c.batch_sweep_interval = 5.minutes  # stalled-batch repair sweep
 
   # --- Realtime broadcast isolation (turbo-rails) ------------------------
   # Route turbo-rails' async render+broadcast jobs to a dedicated queue so a

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

@@ -104,12 +104,14 @@ class CreatePgbusTables < ActiveRecord::Migration<%= migration_version %>
       t.string :description
       t.string :on_finish_class
       t.string :on_success_class
-      t.string :on_discard_class
+      t.string :on_failure_class
+      t.jsonb :on_finish_job
+      t.jsonb :on_success_job
+      t.jsonb :on_failure_job
       t.jsonb :properties, default: {}
       t.integer :total_jobs, null: false, default: 0
       t.integer :completed_jobs, null: false, default: 0
       t.integer :failed_jobs, null: false, default: 0
-      t.integer :discarded_jobs, null: false, default: 0
       t.string :status, null: false, default: "pending"
       t.datetime :created_at, null: false, default: -> { "CURRENT_TIMESTAMP" }
       t.datetime :finished_at
@@ -120,6 +122,24 @@ class CreatePgbusTables < ActiveRecord::Migration<%= migration_version %>
     add_index :pgbus_batches, :status,
       name: "idx_pgbus_batches_status"
 
+    create_table :pgbus_batch_executions do |t|
+      t.string :batch_id, null: false
+      t.string :job_id, null: false
+      t.bigint :msg_id
+      t.string :queue_name
+      t.datetime :created_at, null: false, default: -> { "CURRENT_TIMESTAMP" }
+    end
+
+    add_index :pgbus_batch_executions, :job_id,
+      unique: true, name: "idx_pgbus_batch_executions_job_id"
+    add_index :pgbus_batch_executions, :batch_id,
+      name: "idx_pgbus_batch_executions_batch_id"
+    add_index :pgbus_batch_executions, :created_at,
+      where: "msg_id IS NULL",
+      name: "idx_pgbus_batch_executions_orphans"
+    add_foreign_key :pgbus_batch_executions, :pgbus_batches,
+      column: :batch_id, primary_key: :batch_id, on_delete: :cascade
+
     # Recurring task definitions (synced from config/recurring.yml)
     create_table :pgbus_recurring_tasks do |t|
       t.string :key, null: false
@@ -186,6 +206,7 @@ class CreatePgbusTables < ActiveRecord::Migration<%= migration_version %>
     drop_table :pgbus_stream_queues
     drop_table :pgbus_recurring_executions
     drop_table :pgbus_recurring_tasks
+    drop_table :pgbus_batch_executions
     drop_table :pgbus_batches
     drop_table :pgbus_blocked_executions
     drop_table :pgbus_semaphores

data/lib/pgbus/active_job/adapter.rb

@@ -10,9 +10,12 @@ module Pgbus
         payload_hash = Serializer.serialize_job_hash(active_job)
         payload_hash = Concurrency.inject_metadata(active_job, payload_hash)
         payload_hash = Uniqueness.inject_metadata(active_job, payload_hash)
-        payload_hash = inject_batch_metadata(payload_hash)
+        payload_hash = inject_batch_metadata(payload_hash, active_job: active_job)
 
-        return active_job if uniqueness_rejected?(active_job, payload_hash)
+        if uniqueness_rejected?(active_job, payload_hash, queue: queue)
+          uncount_batch_job(payload_hash)
+          return active_job
+        end
 
         enqueue_with_concurrency(active_job, queue, payload_hash)
       end
@@ -22,18 +25,22 @@ module Pgbus
         payload_hash = Serializer.serialize_job_hash(active_job)
         payload_hash = Concurrency.inject_metadata(active_job, payload_hash)
         payload_hash = Uniqueness.inject_metadata(active_job, payload_hash)
-        payload_hash = inject_batch_metadata(payload_hash)
+        payload_hash = inject_batch_metadata(payload_hash, active_job: active_job)
         delay = [(timestamp - Time.current.to_f).ceil, 0].max
 
-        return active_job if uniqueness_rejected?(active_job, payload_hash)
+        if uniqueness_rejected?(active_job, payload_hash, queue: queue)
+          uncount_batch_job(payload_hash)
+          return active_job
+        end
 
         enqueue_with_concurrency(active_job, queue, payload_hash, delay: delay)
       end
 
       def enqueue_all(active_jobs)
-        # Jobs with uniqueness must go through individual enqueue to acquire locks
-        unique, bulk = active_jobs.partition { |j| Uniqueness.uniqueness_config(j) }
-        unique.each do |j|
+        # Jobs with uniqueness or concurrency must go through individual enqueue
+        # to acquire locks/semaphores — the bulk path cannot (issue #413)
+        individual, bulk = active_jobs.partition { |j| Uniqueness.uniqueness_config(j) || concurrency_config(j) }
+        individual.each do |j|
           if scheduled_in_future?(j)
             enqueue_at(j, j.scheduled_at.to_f)
           else
@@ -41,9 +48,12 @@ module Pgbus
           end
         end
 
-        bulk.group_by { |j| j.queue_name || Pgbus.configuration.default_queue }.each do |queue, jobs|
+        # Group by priority too: send_batch routes through the queue strategy,
+        # so a mixed-priority bulk send needs one produce_batch per level.
+        bulk.group_by { |j| [j.queue_name || Pgbus.configuration.default_queue, j.try(:priority)] }
+            .each do |(queue, priority), jobs|
           immediate, scheduled = jobs.partition { |j| !scheduled_in_future?(j) }
-          enqueue_immediate(queue, immediate)
+          enqueue_immediate(queue, immediate, priority: priority)
           scheduled.each { |j| enqueue_at(j, j.scheduled_at.to_f) }
         end
 
@@ -56,6 +66,8 @@ module Pgbus
         key = Concurrency.extract_key(payload_hash)
         concurrency = concurrency_config(active_job)
         priority = active_job.try(:priority)
+        msg_id = nil
+        blocked = false
 
         if key && concurrency
           result = Concurrency::Semaphore.acquire(key, concurrency[:limit], concurrency[:duration])
@@ -64,34 +76,48 @@ module Pgbus
             msg_id = Pgbus.client.send_message(queue, payload_hash, delay: delay, priority: priority)
             active_job.provider_job_id = msg_id
           else
-            handle_conflict(concurrency, active_job, key, queue, payload_hash, priority: priority)
+            blocked = handle_conflict(concurrency, active_job, key, queue, payload_hash, priority: priority)
           end
         else
           msg_id = Pgbus.client.send_message(queue, payload_hash, delay: delay, priority: priority)
           active_job.provider_job_id = msg_id
         end
 
+        # Bind before backfill so a live message is never left with an unbound
+        # uniqueness row if execution-row bookkeeping raises.
+        bind_acquired_uniqueness_lock(queue, msg_id) if msg_id
+        Batch.backfill_execution(payload_hash, msg_id, physical_queue(queue, priority)) if msg_id
+        # A retry re-enqueue that is now live (sent, or parked as a blocked
+        # execution) must stop the original attempt from signalling completion.
+        Batch.note_retry_reenqueued(payload_hash["job_id"]) if (msg_id || blocked) && retry_retagged?(payload_hash)
+        uniqueness_key = Thread.current[:pgbus_acquired_uniqueness_key]
+        UniquenessKey.clear_bind_stamp!(uniqueness_key) if uniqueness_key
         Thread.current[:pgbus_acquired_uniqueness_key] = nil
         active_job
       rescue StandardError => e
-        # Roll back the uniqueness lock if enqueue failed
-        rollback_key = Thread.current[:pgbus_acquired_uniqueness_key]
-        if rollback_key
-          begin
-            Uniqueness.release_lock(rollback_key)
-          rescue StandardError => rollback_error
-            Pgbus.logger.warn { "[Pgbus] Lock rollback failed: #{rollback_error.message}" }
-          end
+        if msg_id.nil?
+          rollback_acquired_uniqueness_lock
+          uncount_batch_job(payload_hash)
+        else
+          # Message is live: drop the thread-local so a later discard on this
+          # thread cannot release that job's uniqueness lock, but do not
+          # DELETE the pgbus_uniqueness_keys row.
           Thread.current[:pgbus_acquired_uniqueness_key] = nil
         end
         raise e
       end
 
+      def physical_queue(queue, priority)
+        Pgbus.client.target_queue(queue, priority)
+      end
+
       def concurrency_config(active_job)
         active_job.class.respond_to?(:pgbus_concurrency) && active_job.class.pgbus_concurrency
       end
 
-      def handle_conflict(concurrency, active_job, key, queue, payload_hash, priority: nil)
+      # Returns true when the job was parked as a blocked execution (it will
+      # run later), false when it was dropped.
+      def handle_conflict(concurrency, active_job, key, queue, payload_hash, priority: nil) # rubocop:disable Naming/PredicateMethod
         case concurrency[:on_conflict]
         when :block
           Concurrency::BlockedExecution.insert(
@@ -101,14 +127,37 @@ module Pgbus
             priority: priority || Pgbus.configuration.default_priority,
             duration: concurrency[:duration]
           )
+          return true
         when :discard
           Pgbus.logger.info { "[Pgbus] Discarding job #{active_job.class.name}: concurrency limit for #{key}" }
+          # The job will never run: roll back an :until_executed uniqueness lock
+          # acquired earlier in this enqueue (no executor will release it), and
+          # uncount it from its batch so completion is not waited on forever.
+          rollback_acquired_uniqueness_lock
+          uncount_batch_job(payload_hash)
         when :raise
           raise ConcurrencyLimitExceeded, "Concurrency limit reached for key: #{key}"
         end
+        false
       end
 
-      def uniqueness_rejected?(active_job, payload_hash)
+      # Releases an :until_executed lock this enqueue acquired, if any.
+      # Used when the job is dropped before a message is sent (concurrency
+      # :discard, or send_message raising) — otherwise the lock is orphaned
+      # because no executor will ever release it.
+      def rollback_acquired_uniqueness_lock
+        rollback_key = Thread.current[:pgbus_acquired_uniqueness_key]
+        return unless rollback_key
+
+        begin
+          Uniqueness.release_lock(rollback_key)
+        rescue StandardError => e
+          Pgbus.logger.warn { "[Pgbus] Lock rollback failed: #{e.message}" }
+        end
+        Thread.current[:pgbus_acquired_uniqueness_key] = nil
+      end
+
+      def uniqueness_rejected?(active_job, payload_hash, queue:)
         uniqueness_key = Uniqueness.extract_key(payload_hash)
         return false unless uniqueness_key
 
@@ -123,7 +172,7 @@ module Pgbus
         # See issue #333.
         return false if active_job.executions.to_i.positive?
 
-        result = Uniqueness.acquire_enqueue_lock(uniqueness_key, active_job)
+        result = Uniqueness.acquire_enqueue_lock(uniqueness_key, active_job, queue_name: queue)
 
         # :no_lock means no enqueue-time lock needed (e.g. :while_executing strategy)
         return false if result == :no_lock
@@ -147,25 +196,90 @@ module Pgbus
         end
       end
 
-      def inject_batch_metadata(payload_hash)
+      def bind_acquired_uniqueness_lock(queue, msg_id)
+        key = Thread.current[:pgbus_acquired_uniqueness_key]
+        return unless key
+
+        Uniqueness.bind_lock(key, queue_name: queue, msg_id: msg_id)
+      rescue StandardError => e
+        Pgbus.logger.warn { "[Pgbus] Uniqueness bind failed: #{e.message}" }
+      end
+
+      # Reverses inject_batch_metadata for a job discarded at enqueue time
+      # (uniqueness duplicate or concurrency :discard conflict): the message is
+      # never sent, so it can never signal completion. Only applies while the
+      # tagging batch's block is still active on this thread.
+      def uncount_batch_job(payload_hash)
+        batch_id = payload_hash[Batch::METADATA_KEY]
+        return unless batch_id
+
+        if batch_id == Thread.current[:pgbus_batch_id]
+          Batch.untrack_enqueue(payload_hash)
+        elsif retry_retagged?(payload_hash)
+          # The retry never became live; the original attempt's row and its
+          # normal completion signal stand.
+          Batch.forget_retry_reenqueued(payload_hash["job_id"])
+        end
+      end
+
+      # Tagged for a batch while no Batch#enqueue block is active on this
+      # thread — only a retry_on re-enqueue gets there (issue #424).
+      def retry_retagged?(payload_hash)
+        payload_hash[Batch::METADATA_KEY] && Thread.current[:pgbus_batch_id].nil?
+      end
+
+      # A retry_on re-enqueue of a job that is already a batch member: same
+      # job_id (executions > 0), batch_id carried by the BatchId mixin. It
+      # rejoins its batch without being counted again. A first-attempt job
+      # that merely has a batch_id outside a block is NOT tagged — membership
+      # stays explicit; only the callback_batch_id never re-tags.
+      def retry_batch_id_for(active_job)
+        return nil unless active_job.respond_to?(:batch_id)
+        return nil unless active_job.executions.to_i.positive?
+
+        active_job.batch_id
+      end
+
+      # Tag the payload with the active batch and count it in (guarded
+      # increment + execution row, before the send — Batch.track_enqueue). Pass
+      # track: false to only tag, when the caller counts a bulk once.
+      def inject_batch_metadata(payload_hash, active_job: nil, track: true)
         batch_id = Thread.current[:pgbus_batch_id]
-        return payload_hash unless batch_id
+        if batch_id
+          tagged = payload_hash.merge(Batch::METADATA_KEY => batch_id)
+          Batch.track_enqueue(tagged) if track
+          return tagged
+        end
 
-        Thread.current[:pgbus_batch_job_count] = (Thread.current[:pgbus_batch_job_count] || 0) + 1
-        payload_hash.merge(Batch::METADATA_KEY => batch_id)
+        retry_batch_id = active_job && retry_batch_id_for(active_job)
+        return payload_hash unless retry_batch_id
+
+        tagged = payload_hash.merge(Batch::METADATA_KEY => retry_batch_id)
+        Batch.track_retry(tagged)
+        tagged
       end
 
-      def enqueue_immediate(queue, jobs)
+      def enqueue_immediate(queue, jobs, priority: nil)
         return if jobs.empty?
 
-        payloads = jobs.map { |j| Serializer.serialize_job_hash(j) }
-        msg_ids = Pgbus.client.send_batch(queue, payloads)
+        payloads = jobs.map { |j| inject_batch_metadata(Serializer.serialize_job_hash(j), track: false) }
+        # One guarded increment for the whole bulk, not one per job.
+        Batch.track_enqueue(payloads) if Thread.current[:pgbus_batch_id]
+        physical = physical_queue(queue, priority)
+        msg_ids = nil
+        msg_ids = Pgbus.client.send_batch(queue, payloads, priority: priority)
 
         unless msg_ids.is_a?(Array) && msg_ids.size == jobs.size
           raise Pgbus::EnqueueError, "Pgbus batch enqueue failed: expected #{jobs.size} ids, got #{msg_ids&.size || 0}"
         end
 
         jobs.zip(msg_ids).each { |job, id| job.provider_job_id = id }
+        payloads.zip(msg_ids).each { |payload, id| Batch.backfill_execution(payload, id, physical) }
+      rescue Pgbus::EnqueueError
+        Array(payloads).each_with_index do |payload, index|
+          Batch.untrack_enqueue(payload) if msg_ids.nil? || msg_ids[index].nil?
+        end
+        raise
       rescue Pgbus::SchemaNotReady => e
         Pgbus.logger.error { "[Pgbus] #{e.message}" }
         raise

data/lib/pgbus/active_job/batch_id.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module ActiveJob
+    # Gives every ActiveJob a handle on the batch it belongs to.
+    #
+    # Two distinct ids, mirroring solid_queue's ActiveJob::BatchId:
+    #
+    # * +batch_id+ — the batch this job is a *member* of. Assigned by the
+    #   executor from the payload's +pgbus_batch_id+ before +perform+, so a
+    #   running job can call +batch.enqueue+ to add siblings.
+    # * +callback_batch_id+ — the batch this job *reports on*. Set on
+    #   +on_finish+/+on_success+/+on_failure+ jobs at fire time. A callback is
+    #   never a member of the batch it reports on, so its +batch_id+ is nil.
+    #
+    # Both round-trip through +serialize+/+deserialize+, and both are omitted
+    # from the serialized hash when unset — an unbatched job's payload is
+    # byte-for-byte what it was before this mixin existed.
+    module BatchId
+      extend ActiveSupport::Concern
+
+      included do
+        attr_accessor :batch_id, :callback_batch_id
+      end
+
+      def serialize
+        data = super
+        data["batch_id"] = batch_id if batch_id
+        data["callback_batch_id"] = callback_batch_id if callback_batch_id
+        data
+      end
+
+      def deserialize(job_data)
+        super
+        self.batch_id = job_data["batch_id"]
+        self.callback_batch_id = job_data["callback_batch_id"]
+      end
+
+      # The batch this job reports on, or is a member of. nil when neither.
+      def batch
+        return @batch if defined?(@batch)
+
+        id = callback_batch_id || batch_id
+        @batch = id && Pgbus::Batch.find(id)
+      end
+    end
+  end
+end

data/lib/pgbus/active_job/executor.rb

@@ -56,9 +56,14 @@ module Pgbus
             # released on completion or DLQ.
             nil
           when :while_executing
-            # Acquire the lock now. If another worker is already executing
-            # this job, skip it — VT will expire and it'll be retried.
-            unless Uniqueness.acquire_execution_lock(uniqueness_key, payload)
+            # Acquire the lock now, bound to this message. If another worker is
+            # already executing this job, skip it — VT will expire and it'll be
+            # retried. A row left by a crashed attempt of THIS message is
+            # re-acquired, not treated as a duplicate.
+            acquired = Uniqueness.acquire_execution_lock(
+              uniqueness_key, payload, msg_id: message.msg_id.to_i, queue_name: queue_name
+            )
+            unless acquired
               Pgbus.logger.info { "[Pgbus] Skipping duplicate execution for #{job_class}" }
               return :skipped
             end
@@ -67,6 +72,7 @@ module Pgbus
 
         Pgbus.logger.debug { "[Pgbus::Executor] deserialized #{tag} job_class=#{job_class}" }
         job_succeeded = false
+        retried = false
 
         msg_id = message.msg_id.to_i
         instrument_payload = {
@@ -85,13 +91,22 @@ module Pgbus
           # (issue #368). Pass this executor's config so an injected allowlist
           # is not silently ignored in favour of Pgbus.configuration.
           job = Serializer.deserialize_job_data(payload, configuration: config)
+          # Batch membership rides the pgbus metadata key, not the serialized
+          # job data, so hand it to the job before perform — that is what makes
+          # `batch` (and `batch.enqueue` for open batches) work inside a job.
+          assign_batch_id(job, payload)
           Pgbus.logger.debug { "[Pgbus::Executor] running #{tag} job_class=#{job_class}" }
           execute_job(job)
+          # retry_on re-enqueues from inside perform_now and returns normally:
+          # this attempt is done (archive it) but the job is not — the retry
+          # message carries the batch tag and signals on its own outcome.
+          retried = Batch.retry_reenqueued?(payload["job_id"])
           Pgbus.logger.debug { "[Pgbus::Executor] perform_returned #{tag} job_class=#{job_class}" }
           archive_from(queue_name, msg_id, source_queue: source_queue)
           Pgbus.logger.debug { "[Pgbus::Executor] archived #{tag} job_class=#{job_class}" }
-          FailedEventRecorder.clear!(queue_name: queue_name, msg_id: msg_id)
           job_succeeded = true
+          release_uniqueness_lock(uniqueness_key)
+          FailedEventRecorder.clear!(queue_name: queue_name, msg_id: msg_id)
         end
 
         instrument("pgbus.job_completed", queue: queue_name, job_class: job_class)
@@ -108,6 +123,10 @@ module Pgbus
         # silently lost control flow — no failed event row, no job_failed
         # notification, uniqueness lock held until VT expired. See issue #126.
         handle_failure(message, queue_name, e, payload: payload)
+        # A failed :while_executing attempt is no longer executing — release
+        # so the retry (same message, after VT) can acquire. :until_executed
+        # keeps its lock until success/DLQ by design (#126, #333).
+        release_uniqueness_lock(uniqueness_key) if uniqueness_strategy == :while_executing
         instrument(
           "pgbus.job_failed",
           queue: queue_name,
@@ -129,15 +148,32 @@ module Pgbus
         # 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
+          # Uniqueness is released once, immediately after archive. A second
+          # key-only DELETE here can drop a successor that acquired the same
+          # key if the first DELETE committed but the client raised afterward.
           signal_concurrency(payload)
-          signal_batch_completed(payload)
-          # Release uniqueness lock on successful completion (both strategies)
-          Uniqueness.release_lock(uniqueness_key) if uniqueness_key
+          signal_batch_completed(payload) unless retried
         end
+        Batch.clear_retry_reenqueued
       end
 
       private
 
+      def assign_batch_id(job, payload)
+        batch_id = payload[Batch::METADATA_KEY]
+        return unless batch_id && job.respond_to?(:batch_id=)
+
+        job.batch_id = batch_id
+      end
+
+      def release_uniqueness_lock(key)
+        return unless key
+
+        Uniqueness.release_lock(key)
+      rescue StandardError => e
+        Pgbus.logger.warn { "[Pgbus] Uniqueness release failed: #{e.message}" }
+      end
+
       def execute_job(job)
         if defined?(Rails) && Rails.respond_to?(:application) && Rails.application
           wrapper = reloading? ? Rails.application.reloader : Rails.application.executor
@@ -286,7 +322,7 @@ module Pgbus
         batch_id = payload[Batch::METADATA_KEY]
         return unless batch_id
 
-        Batch.job_completed(batch_id)
+        Batch.job_completed(batch_id, job_id: payload["job_id"])
       rescue StandardError => e
         Pgbus.logger.warn { "[Pgbus] Batch completion signal failed: #{e.message}" }
       end
@@ -295,7 +331,7 @@ module Pgbus
         batch_id = payload[Batch::METADATA_KEY]
         return unless batch_id
 
-        Batch.job_discarded(batch_id)
+        Batch.job_discarded(batch_id, job_id: payload["job_id"])
       rescue StandardError => e
         Pgbus.logger.warn { "[Pgbus] Batch discard signal failed: #{e.message}" }
       end