cosmonats 0.4.3 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5704f0d1aa99b76acfecee7b986cab37a9bc8863832c868edc28a3ffeeb7f418
-  data.tar.gz: 9128efdd2e934778ea6d5e6412cbcd0d8424743a881b0a56edb665eefe8463b8
+  metadata.gz: d13ca716127cde6f5ee26e9a859a91f68d8f0e91a4e43c48b4ebceaae1a2dfef
+  data.tar.gz: 2a421516f48ae13cadec5470d3a3113f28f403a979319e36667866f2367b201c
 SHA512:
-  metadata.gz: 8a64d21839c7a7f18da7a47eb10a40209126d350b5bfc3778d75f110b4bc05034e558f87566f1140fd71e0a0b44f1019f49832f1e6f66460da06fdcc1ba8e7c1
-  data.tar.gz: 21c1231e93af54a7ca68fc021415657dd9bb332893779ec936dc140e947055fff4e2c2a8fa8962d59a34ab3b6599ef9aa67b81b6dcf19d9b190df756f260b6f2
+  metadata.gz: c0aee10ad2470135778ef9d16732f041ac67426b3d5aeaeeaf899c473b4ec5af5018c0d1d7f3aeb862dea21d41741e85e9f84070f3c802a95014f48c65c95eac
+  data.tar.gz: fe1a8cb37bbd553b1bc69eaf35ac0134bf155c222d9f49b539286509884a194c8fc830c05bdb35052e51ebf83055fdf7f030a2d697b6c0912ba1884eacd46628

data/README.md

@@ -73,6 +73,7 @@ bundle exec cosmo -C config/cosmo.yml -c 20 streams # Streams only
 - [Advanced Usage](#-advanced-usage)
   - [Cron](#cron)
   - [Priority Queues](#priority-queues)
+  - [Batches](#batches)
   - [Concurrency Limiting](#concurrency-limiting)
   - [Custom Serializers](#custom-serializers)
   - [Error Handling](#error-handling)
@@ -147,6 +148,7 @@ nothing else to run.
 - **Job uniqueness** — prevent duplicate execution
 - **Concurrency limits** — cap simultaneous executions per class or per key
 - **Cron scheduling** — recurring jobs manageable live from the web UI
+- **Batches** — group jobs and fire a callback once the whole group finishes, including nested batches
 
 ### 🌊 Stream Processing
 - **Real-time event streams** — process continuous data feeds
@@ -204,8 +206,8 @@ concurrency: 5                     # Number of worker threads
 consumers:                         # Declare consumer groups for streams, things that pull messages and process them
   jobs:                            # Consumer configs for jobs (or streams)
     default:                       # Stream name
-      ack_policy: explicit         # Acknowledgment required for each message, can be explicit, none, or all
-      max_deliver: 10              # Max retry attempts before sending to a dead stream
+      ack_policy: explicit         # Acknowledgment required for each message can be explicit, none, or all
+      max_deliver: 30              # Max retry attempts before sending to a dead stream. Safety ceiling only, keep it above every job class's own retry
       max_ack_pending: 10          # Max messages waiting for ack, if exceeded, the server will stop delivering new messages until some are acked
       ack_wait: 15                 # Seconds to wait for ack before redelivering
       subject: jobs.%{name}.>      # Subject pattern for this consumer, %{name} replaced with stream name, becomes `jobs.default.>`
@@ -328,6 +330,7 @@ The `>` wildcard matches everything after that prefix. Think of subjects as topi
 timeout: 25                 # Shutdown timeout in seconds
 concurrency: &concurrency 1 # Number of worker threads
 max_retries: &max_retries 3 # Default max retries
+batch_expiry: 259200         # Seconds before a Batch's tracking data expires (default: 3 days)
 
 stream_config: &stream_config
   storage: file         # storage type (file or memory)
@@ -340,7 +343,7 @@ stream_config: &stream_config
 
 consumer_config: &consumer_config
   ack_policy: explicit    # ack policy (explicit, none, all), each individual message must be acknowledged
-  max_deliver: 10         # maximum number of times a message will be delivered before it's considered failed
+  max_deliver: 30         # maximum number of times a message will be delivered before it's considered failed. keep it above every job class's own retry; a job that still exceeds it is capped and dead-lettered early with a warning
   max_ack_pending: 20     # maximum number of messages with pending ack for this consumer
   ack_wait: 60            # time in seconds to wait for an ack before redelivering the message
   subject: jobs.%{name}.> # subject pattern for consumer, %{name} will be replaced with stream name
@@ -466,6 +469,58 @@ class UrgentJob
 end
 ```
 
+### Batches
+
+Group jobs together and fire a callback once every one of them has finished. The registered class
+is plain Ruby — it implements `on_complete(status, opts)` / `on_success(status, opts)`, not
+`perform` — and runs on the job worker pool, not inline on whichever thread finalized the batch.
+`status` is `{ bid:, total:, succeeded:, failed: }` and `opts` is whatever you passed to `#on`.
+
+```ruby
+batch = Cosmo::Batch.new
+batch.jobs do
+  ImportJob.perform_async(1)
+  ImportJob.perform_async(2)
+end
+batch.on(:complete, NotifyUser, user_id: 1)  # fires once every job has finished, pass or fail
+batch.on(:success, NotifyUser, user_id: 1)   # fires only if none of them failed
+
+class NotifyUser
+  def on_complete(status, opts)
+    UserMailer.batch_done(opts[:user_id], status[:succeeded], status[:total]).deliver_later
+  end
+
+  def on_success(status, opts)
+    UserMailer.batch_succeeded(opts[:user_id]).deliver_later
+  end
+end
+```
+
+- `#jobs` is the only place membership is tracked — jobs enqueued outside the block aren't part of
+  the batch. Call it at least once (an empty block is fine) to close the batch.
+- `:complete` always fires once every job is done. `:success` fires only if none were dead-lettered
+  or dropped after exhausting retries. Both can be registered before or after `#jobs` — a callback
+  registered after the batch has already finished still fires.
+
+**Nested batches** — a running job can spawn its own sub-batch, which counts as one pending unit
+  of its parent and propagates any failure upward as a single failure unit:
+
+  ```ruby
+  class ImportJob
+    include Cosmo::Job
+
+    def perform(account_id)
+      sub_batch = Cosmo::Batch.new(parent: batch_id)  # batch_id is this job's own batch, if any
+      sub_batch.jobs { SyncRecordJob.perform_async(account_id) }
+    end
+  end
+  ```
+- Batch tracking data (pending counts, callbacks, results) expires automatically after
+  `Config[:batch_expiry]` seconds (default: 3 days).
+- Open and finished batches — with pending/succeeded/failed counts — are listed live in the web UI's **Batches** tab.
+- Only `Cosmo::Job`-based jobs are tracked; the ActiveJob adapter doesn't currently participate in
+  batches (see [`docs/active_job.md`](docs/active_job.md)).
+
 ### Concurrency Limiting
 
 ```ruby
@@ -521,6 +576,22 @@ class ResilientJob
 end
 ```
 
+By default, a failed job is redelivered after `attempt**4 + 15` seconds. Override that per job class
+with `retry_in`, given the 1-based attempt count and the exception that was raised:
+```ruby
+class ThrottledApiJob
+  include Cosmo::Job
+  options retry: 5, retry_in: ->(count, exception) { exception.is_a?(RateLimitedError) ? 60 : count * 10 }
+
+  def perform(...)
+    # ...
+  end
+end
+```
+If the proc returns something non-numeric/non-positive, or raises, the default backoff is used
+instead. Note: if this job class also sets `limit: { concurrency: ... }` (see above), `count`
+includes deliveries that were turned away for lack of a free slot, not just failed attempts.
+
 ### Testing
 
 ```ruby
@@ -541,10 +612,10 @@ assert_kind_of String, jid
 config.active_job.queue_adapter = :cosmonats
 ```
 The ActiveJob queue name maps directly to a Cosmo stream. Use `cosmo_options` for anything
-Cosmo-specific — retries, DLQ behavior, or overriding the target stream:
+Cosmo-specific — retries, DLQ behavior, overriding the target stream, or a custom retry delay:
 ```ruby
 class ReportJob < ApplicationJob
-  cosmo_options retry: 5, dead: false, stream: :critical
+  cosmo_options retry: 5, dead: false, stream: :critical, retry_in: ->(count, exception) { count * 10 }
 
   def perform(report_id)
     Report.find(report_id).generate!
@@ -656,6 +727,7 @@ sudo systemctl enable cosmo && sudo systemctl start cosmo
 - **Jobs** — enqueued, scheduled, busy, and dead views, with per-job retry and delete
 - **Streams** — per-stream state (messages, bytes, consumers) with pause/resume
 - **Crons** — every schedule deployed in NATS, with run-now and delete
+- **Batches** — open and finished batches with pending/succeeded/failed counts
 - Summary counters (processed / failed / busy / enqueued / retries / scheduled / dead) backed by a NATS KV counter, no separate metrics store needed
 
 **Structured logs:**

data/lib/cosmo/active_job/executor.rb

@@ -8,6 +8,15 @@ module Cosmo
 
       options stream: :default
 
+      # Resolves +retry_in+ from the underlying ActiveJob class (declared via +cosmo_options+),
+      # since the processor only ever sees +Executor+ as the worker class for ActiveJob-dispatched jobs.
+      def self.retry_in(data)
+        job_class = Utils::String.safe_constantize(data.dig(:args, 0, :job_class))
+        return super unless job_class.respond_to?(:get_cosmo_options)
+
+        job_class.get_cosmo_options[:retry_in] || super
+      end
+
       def perform(job_data)
         ::ActiveJob::Base.execute(Utils::Hash.stringify_keys(job_data))
       end

data/lib/cosmo/active_job/options.rb

@@ -5,7 +5,7 @@ module Cosmo
     # Adds +cosmo_options+ to ActiveJob classes.
     #
     #   class MyJob < ApplicationJob
-    #     cosmo_options retry: 5, dead: false
+    #     cosmo_options retry: 5, dead: false, retry_in: ->(count, exception) { count * 30 }
     #
     #     def perform(user_id)
     #       # ...
@@ -13,11 +13,14 @@ module Cosmo
     #   end
     #
     # Options mirror those accepted by +Cosmo::Job+:
-    #   retry:  [Integer]  Number of retries before giving up (default: 3)
-    #   dead:   [Boolean]  Move to DLQ when retries exhausted? (default: true)
-    #   stream: [Symbol]   Override the NATS stream (default: derived from queue_name)
+    #   retry:    [Integer, Boolean]  Number of retries before giving up (default: 3). +false+ means 0 (no retries)
+    #   dead:     [Boolean]  Move to DLQ when retries exhausted? (default: true)
+    #   stream:   [Symbol]   Override the NATS stream (default: derived from queue_name)
+    #   retry_in: [Proc]     <tt>->(count, exception) { }</tt> computing seconds to wait before redelivering a
+    #             failed job (+count+ is the 1-based delivery attempt that just failed). Falls back to the
+    #             default backoff if unset, non-numeric/non-positive, or raising.
     module Options
-      VALID_OPTIONS = %i[retry dead stream].freeze
+      VALID_OPTIONS = %i[retry dead stream retry_in].freeze
 
       def self.included(base)
         base.extend(ClassMethods)

data/lib/cosmo/api/batch.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module API
+    # Read-model over Cosmo::Batch's KV/counter state, for the status API and Web UI.
+    class Batch
+      LIMIT = 25
+
+      def self.all(limit: LIMIT)
+        kv.keys("*.meta", limit: limit).map { |key| new(key.delete_suffix(".meta")) }
+      end
+
+      def self.kv
+        Cosmo::Batch.kv
+      end
+
+      def self.counter
+        Cosmo::Batch.counter
+      end
+
+      attr_reader :bid
+
+      def initialize(bid)
+        @bid = bid
+      end
+
+      def meta
+        Utils::Json.parse(kv.get("#{bid}.meta")&.value) || {}
+      end
+
+      def parent_id
+        meta[:parent_id]
+      end
+
+      def created_at
+        meta[:created_at]
+      end
+
+      def ready?
+        !!kv.get("#{bid}.ready")
+      end
+
+      # { total:, pending:, succeeded:, failed: } -- read from the +ready+
+      # snapshot once finalized (counters are reset by then), or live off the
+      # counters (with +pending+ still moving) while the batch is still open.
+      def stats
+        entry = kv.get("#{bid}.ready")
+        return Utils::Json.parse(entry.value).merge(pending: 0) if entry
+
+        live_stats
+      end
+
+      def callback(event)
+        entry = kv.get("#{bid}.callback.#{event}")
+        Utils::Json.parse(entry&.value)
+      end
+
+      def callbacks
+        { success: callback(:success), complete: callback(:complete) }
+      end
+
+      def fired?(event)
+        !!kv.get("#{bid}.fired.#{event}")
+      end
+
+      private
+
+      def live_stats
+        total = counter.get("#{bid}.total")
+        failed = counter.get("#{bid}.failed")
+        pending = counter.get("#{bid}.pending")
+        { total: total, pending: pending, succeeded: total - failed - pending, failed: failed }
+      end
+
+      def kv
+        self.class.kv
+      end
+
+      def counter
+        self.class.counter
+      end
+    end
+  end
+end

data/lib/cosmo/api/counter.rb

@@ -21,19 +21,20 @@ module Cosmo
         increment(:failed)
       end
 
-      def increment(key, by: 1)
-        publish(key, "+#{by}")
+      def increment(key, by: 1, msg_id: nil)
+        publish(key, "+#{by}", msg_id: msg_id)
       end
       alias incr increment
 
-      def decrement(key, by: 1)
-        publish(key, "-#{by}")
+      def decrement(key, by: 1, msg_id: nil)
+        publish(key, "-#{by}", msg_id: msg_id)
       end
       alias decr decrement
 
       def reset(key)
         client.purge(STREAM_NAME, subject(key))
       end
+      alias purge reset
 
       def get(key)
         raw = client.get_message(STREAM_NAME, direct: true, subject: subject(key))
@@ -44,11 +45,21 @@ module Cosmo
 
       private
 
-      def publish(key, value)
+      # @param msg_id [String, nil] when given, sent as Nats-Msg-Id so a redelivered
+      #   caller (e.g. a job whose ack was lost and retried) collapses onto the
+      #   stream's duplicate_window instead of double-applying the +/- delta.
+      # @return [Integer, nil] the resulting counter value, or +nil+ when this exact
+      #   +msg_id+ was already applied (a deduped publish's PubAck carries no +val+
+      #   at all -- coercing that to 0 would look identical to "counter is now 0",
+      #   so callers must treat +nil+ as "no new information" rather than a real value).
+      def publish(key, value, msg_id: nil)
         rescued = nil
+        headers = { "Nats-Incr" => value }
+        headers["Nats-Msg-Id"] = msg_id if msg_id
 
         begin
-          client.publish(subject(key), "", header: { "Nats-Incr" => value }).val.to_i
+          ack = client.publish(subject(key), "", header: headers)
+          ack.val.to_i unless ack.duplicate
         rescue NATS::JetStream::Error::NoStreamResponse
           raise if rescued
 

data/lib/cosmo/api/cron/entry.rb

@@ -84,7 +84,7 @@ module Cosmo
                              jid: SecureRandom.hex(12),
                              class: @class_name,
                              args: @args,
-                             retry: ::Cosmo::Job::Data::DEFAULTS[:retry],
+                             retry: ::Cosmo::Job::Data.default_retry,
                              dead: ::Cosmo::Job::Data::DEFAULTS[:dead]
                            })
         end

data/lib/cosmo/api/cron.rb

@@ -65,7 +65,7 @@ module Cosmo
                                      jid: SecureRandom.hex(12),
                                      class: body[:class],
                                      args: body[:args] || [],
-                                     retry: body[:retry] || Job::Data::DEFAULTS[:retry],
+                                     retry: body[:retry] || Job::Data.default_retry,
                                      dead: body[:dead].nil? ? Job::Data::DEFAULTS[:dead] : body[:dead]
                                    })
         client.publish(target, payload, stream: stream_name)

data/lib/cosmo/api/kv.rb

@@ -25,6 +25,13 @@ module Cosmo
         # nop
       end
 
+      # Creates +key+ if it does not already exist -- a plain CAS-if-absent with
+      # no per-message TTL involved. Raises NATS::KeyValue::KeyWrongLastSequenceError
+      # when the key is already live, so callers can use it as a fire-once gate.
+      def create(key, value)
+        kv.create(key, value.to_s)
+      end
+
       # Writes a KV-Operation tombstone. On a ttl-bearing bucket this leaves the
       # subject occupied, so a subsequent #set(ttl:) CAS with last_seq: 0 will
       # keep failing -- use #erase on those buckets instead.

data/lib/cosmo/api.rb

@@ -5,6 +5,7 @@ require "cosmo/api/counter"
 require "cosmo/api/kv"
 require "cosmo/api/stats"
 require "cosmo/api/cron"
+require "cosmo/api/batch"
 
 module Cosmo
   module API

data/lib/cosmo/batch/callback.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Cosmo
+  class Batch
+    # The actual job that gets enqueued when a batch finishes. It looks up
+    # the plain Ruby class you passed to Batch#on and calls on_success or
+    # on_complete on it.
+    class Callback
+      include Cosmo::Job
+
+      def perform(class_name, event, status, opts)
+        klass = Utils::String.safe_constantize(class_name)
+        klass&.new&.public_send("on_#{event}", status, opts)
+      end
+    end
+  end
+end

data/lib/cosmo/batch/dispatcher.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+module Cosmo
+  class Batch
+    # The engine behind Batch.notify: decides when a batch is done and
+    # fires its callbacks.
+    module Dispatcher
+      private
+
+      # Decrements pending. If this decrement just brought it to zero,
+      # finalize the batch. This never touches total or failed - that's what
+      # lets #jobs release its placeholder slot without messing up the stats
+      # we report back.
+      def release_pending(bid, msg_id: nil)
+        pending = counter.decrement("#{bid}.pending", msg_id: msg_id)
+        finalize(bid) if pending.to_i.zero?
+      end
+
+      # Runs once per batch. Has its own "already done" guard, separate from
+      # try_fire's, because calling #jobs again on an already-finished batch
+      # would otherwise re-run this and double-report to the parent batch.
+      def finalize(bid)
+        kv.create("#{bid}.finalized", "1")
+
+        total = counter.get("#{bid}.total")
+        failed = counter.get("#{bid}.failed")
+        stats = { total: total, succeeded: total - failed, failed: failed }
+        kv.set("#{bid}.ready", Utils::Json.dump(stats))
+        purge_counters(bid)
+
+        EVENTS.each { |event| try_fire(bid, event, stats: stats) }
+        parent_propagate(bid, failed: failed)
+      rescue NATS::KeyValue::KeyWrongLastSequenceError
+        # already finalized - this is a stray completion from a reused Batch, ignore it
+      end
+
+      def purge_counters(bid)
+        %w[total pending failed].each { |key| counter.purge("#{bid}.#{key}") }
+      end
+
+      # Tells the parent batch "one of your jobs (me) is done", counting any
+      # failure in this batch as one failure for the parent. Same shape as
+      # #notify - the child's bid just stands in for a jid here.
+      def parent_propagate(bid, failed:)
+        meta = Utils::Json.parse(kv.get("#{bid}.meta")&.value)
+        parent_id = meta && meta[:parent_id]
+        return unless parent_id
+
+        notify(parent_id, bid, success: failed.zero?)
+      end
+
+      # Adds this batch as one job of the parent batch.
+      def link(parent_id)
+        counter.increment("#{parent_id}.total")
+        counter.increment("#{parent_id}.pending")
+      end
+
+      # Fires the callback for +event+, if one is registered and the batch is
+      # ready. Safe to call any number of times from either #on or a job
+      # finishing - it only actually fires once.
+      def try_fire(bid, event, stats: nil)
+        callback_entry = kv.get("#{bid}.callback.#{event}")
+        stats ||= Utils::Json.parse(kv.get("#{bid}.ready")&.value)
+        return unless fireable?(callback_entry, event, stats)
+
+        kv.create("#{bid}.fired.#{event}", "1")
+        dispatch_callback(callback_entry, bid, event, stats)
+      rescue NATS::KeyValue::KeyWrongLastSequenceError
+        # someone else already fired this event, do nothing
+      end
+
+      def fireable?(callback_entry, event, stats)
+        return false unless callback_entry && stats
+        return false if event == :success && stats[:failed].to_i.positive?
+
+        true
+      end
+
+      def dispatch_callback(callback_entry, bid, event, stats)
+        callback = Utils::Json.parse(callback_entry.value)
+        status = stats.merge(bid: bid)
+        Callback.perform_async(callback[:class], event, status, callback[:opts] || {})
+      end
+    end
+  end
+end

data/lib/cosmo/batch.rb

@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "cosmo/batch/callback"
+require "cosmo/batch/dispatcher"
+
+module Cosmo
+  # Groups jobs together and runs a callback once they've all finished.
+  #
+  #   batch = Cosmo::Batch.new
+  #   batch.jobs { MyJob.perform_async(1); MyJob.perform_async(2) }
+  #   batch.on(:complete, MyCallback, user_id: 1)
+  #
+  # :complete fires once every job is done, pass or fail.
+  # :success only fires if none of them ended up dead-lettered or dropped.
+  #
+  # The callback is a plain Ruby class, not a job. It just needs an
+  # on_complete/on_success(status, options) method:
+  #
+  #   class MyCallback
+  #     def on_complete(status, options)
+  #       Notifier.notify(options[:user_id], "#{status[:succeeded]}/#{status[:total]} done")
+  #     end
+  #   end
+  #
+  # It still runs on the worker pool, not on the thread that finished the
+  # batch - it's wrapped in an internal Batch::Callback job.
+  #
+  #   status - { bid:, total:, succeeded:, failed: }
+  #   options - arbitrary custom arguments passed to #on
+  #
+  # Only jobs enqueued inside a #jobs block count. Nested batches must be created explicitly,
+  # by passing the parent's batch_id:
+  #
+  #   Cosmo::Batch.new(parent: batch_id)
+  #
+  # A nested batch counts as one job of its parent. If anything fails in it,
+  # the failure is propagated to the parent counter either. The parent succeeds
+  # when everything succeeds. Always call #jobs at least once, even with nothing
+  # in it - a batch that never gets a #jobs call never closes and never fires.
+  #
+  # == How this is actually stored in NATS
+  #
+  # Every batch has an id (bid, a random hex string). State lives
+  # in two places, both keyed by that bid:
+  #
+  # 1. Three atomic counters (via API::Counter, namespace "batch", backed by
+  #    the shared _cosmostats stream):
+  #
+  #      <bid>.total    - how many jobs have ever joined this batch
+  #      <bid>.pending  - how many of them are still running
+  #      <bid>.failed   - how many ended up dead-lettered or dropped
+  #
+  #    Each increment/decrement is one atomic NATS publish, and the reply
+  #    tells you the resulting value right away - no separate read needed.
+  #    Whichever decrement of +pending+ happens to bring it down to exactly
+  #    0 is the one that finalizes the batch (see #release_pending). Since
+  #    it's a single atomic counter, exactly one decrement can ever be "the
+  #    one that hits zero", so the batch can't finish twice or too early.
+  #
+  # 2. A KV bucket (BUCKET, ttl'd so old batches clean themselves up):
+  #
+  #      <bid>.meta               - { parent_id, created_at }, written once at creation
+  #      <bid>.callback.success   - { class, opts } from #on(:success, ...)
+  #      <bid>.callback.complete  - { class, opts } from #on(:complete, ...)
+  #      <bid>.ready              - { total, succeeded, failed }, written once
+  #                                  by #finalize, right after reading the counters
+  #      <bid>.fired.success      - exists once the :success callback has run
+  #      <bid>.fired.complete     - exists once the :complete callback has run
+  #      <bid>.finalized          - exists once #finalize has run for this bid
+  #
+  #    The fired.* and finalized keys are write-once: kv.create raises if the
+  #    key already exists. That's what stops the same callback (or the same
+  #    finalize) from running twice when two things race to trigger it - a
+  #    job finishing at the same moment someone calls #on, for example.
+  class Batch
+    extend Dispatcher
+
+    BUCKET = "cosmo_jobs_batches"
+    DEFAULT_EXPIRY = "3d"
+    EVENTS = %i[success complete].freeze
+
+    def self.current
+      Thread.current[:cosmo_batch]
+    end
+
+    # Called by Job::Processor when a job is done for good - acked, or
+    # dead-lettered/dropped. Never call this for a retry. JID is used to dedupe: if the same
+    # job's ack gets lost, and it's redelivered, we don't want to count it twice.
+    def self.notify(bid, jid, success:)
+      counter.increment("#{bid}.failed", msg_id: "batch.#{bid}.#{jid}.failed") unless success
+      release_pending(bid, msg_id: "batch.#{bid}.#{jid}.pending")
+    end
+
+    def self.counter
+      @counter ||= API::Counter.new("batch")
+    end
+
+    def self.kv
+      @kv ||= API::KV.new(BUCKET, ttl: Utils::Duration.parse(Config[:batch_expiry] || DEFAULT_EXPIRY))
+    end
+
+    attr_reader :bid, :parent_id
+
+    def initialize(parent: nil)
+      @bid = SecureRandom.hex(8)
+      @parent_id = parent
+      kv.set("#{@bid}.meta", Utils::Json.dump({ parent_id: @parent_id, created_at: Time.now.to_i }))
+      link(@parent_id) if @parent_id
+    end
+
+    # Every perform_async/perform_at/perform_in called inside this block
+    # joins the batch. Holds a placeholder "pending" slot for the whole
+    # block, so the batch can't finish while you're still adding jobs, even
+    # if the first one completes instantly. The placeholder never counts
+    # toward total, so it doesn't skew the stats. Safe to call more than
+    # once. If the block raises, we still release the placeholder (so the
+    # batch isn't stuck) and let the error propagate as normal.
+    def jobs
+      previous = Thread.current[:cosmo_batch]
+      Thread.current[:cosmo_batch] = self
+      counter.increment("#{bid}.pending")
+      begin
+        yield
+      ensure
+        Thread.current[:cosmo_batch] = previous
+        release_pending
+      end
+    end
+
+    # Registers a callback for +event+ (:success or :complete). Can be
+    # called before or after #jobs - if the batch already finished, it
+    # fires right away.
+    def on(event, klass, **opts)
+      raise ArgumentError, "event must be :success or :complete" unless EVENTS.include?(event)
+
+      kv.set("#{bid}.callback.#{event}", Utils::Json.dump({ class: klass.name, opts: opts }))
+      try_fire(event)
+    end
+
+    # Called by Job#perform for each job added inside #jobs.
+    def register_job!
+      counter.increment("#{bid}.total")
+      counter.increment("#{bid}.pending")
+    end
+
+    # Undoes #register_job! when the publish itself failed, so the job
+    # never really joined the batch. Unlike a real completion, this also
+    # reverts total, not just pending.
+    def rollback_job!
+      counter.decrement("#{bid}.total")
+      release_pending
+    end
+
+    private
+
+    # Bridges instance methods to the private class methods Dispatcher adds to Batch
+    def counter = self.class.counter
+    def kv = self.class.kv
+    def link(parent_id) = self.class.send(:link, parent_id)
+    def release_pending = self.class.send(:release_pending, bid)
+    def try_fire(event) = self.class.send(:try_fire, bid, event)
+  end
+end