cosmonats 0.6.0 → 0.8.0

data/lib/cosmo/job/processor.rb

@@ -5,27 +5,63 @@ require "timeout"
 module Cosmo
   module Job
     class Processor < ::Cosmo::Processor
+      # Headers the scheduler consumes rather than forwards. The +X-+ ones address the dispatch
+      # itself, and NATS sets +Nats-Scheduler+/+Nats-Schedule-Next+ on a message it fires and
+      # rejects a publish that carries them back.
+      DISPATCH_HEADERS = %w[X-Stream X-Subject X-Execute-At Nats-Expected-Stream
+                            Nats-Scheduler Nats-Schedule-Next].freeze
+
+      # @raise [UnknownJobStreamError] when --streams/--stream or COSMO_JOBS_STREAMS names an unconfigured stream
+      def self.validate_options!(options)
+        StreamFilter.from(options[:streams]).validate!
+      end
+
+      # @return [Array<String>] job streams pulled from, plus +scheduled+ when dispatching scheduled jobs
+      def subscriptions
+        names = @consumers.map { |(_, config, _)| config[:stream].to_s }
+        names << StreamFilter::SCHEDULED.to_s if scheduler? && scheduled_config
+        names
+      end
+
       private
 
       def setup
+        filter = StreamFilter.from(@options[:streams]).validate!
+
         # Initialize singletons before starting to process messages
-        API::Busy.instance
-        API::Counter.instance
+        API::Stats::Busy.instance
+        API::Stats::Totals.instance
         Limit.instance
+        Config.server_middleware
 
         jobs_config = Config.dig(:consumers, :jobs)
         jobs_config&.each do |stream_name, config|
-          next if stream_name == :scheduled # scheduled jobs are handled in schedule_loop
+          next if stream_name == StreamFilter::SCHEDULED # scheduled jobs are handled in schedule_loop
+          next unless filter.include?(stream_name)
 
           @consumers << subscribe(stream_name, config)
         end
+
+        log_subscriptions
+      end
+
+      def log_subscriptions
+        return if @consumers.empty?
+
+        names = @consumers.map { |(_, config, _)| config[:consumer] }
+        names << consumer_name(StreamFilter::SCHEDULED) if scheduler? && scheduled_config
+        Logger.info "subscribed: #{names.join(", ")}#{" (scheduler off)" unless scheduler?}"
+      end
+
+      def scheduled_config
+        Services.scheduled_consumer
       end
 
       def schedule_loop # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/MethodLength, Metrics/AbcSize
-        config = Config.dig(:consumers, :jobs, :scheduled)
+        config = scheduled_config
         return unless config
 
-        subscription, = subscribe(:scheduled, config)
+        subscription, = subscribe(StreamFilter::SCHEDULED, config)
         while running?
           break unless running?
 
@@ -33,13 +69,18 @@ module Cosmo
           timeout = ENV.fetch("COSMO_JOBS_SCHEDULER_FETCH_TIMEOUT", 5).to_f
           messages = fetch(subscription, batch_size: 100, timeout:)
           messages&.each do |message|
-            headers = message.header.except("X-Stream", "X-Subject", "X-Execute-At", "Nats-Expected-Stream")
+            headers = message.header.except(*DISPATCH_HEADERS)
             stream, subject, execute_at = message.header.values_at("X-Stream", "X-Subject", "X-Execute-At")
             headers["Nats-Expected-Stream"] = stream
+            scheduler = message.header["Nats-Scheduler"]
+            if scheduler
+              headers["X-Scheduled-By"] = scheduler
+              headers["X-Enqueued-At"] = message.metadata.timestamp.to_f.to_s
+            end
             execute_at = execute_at.to_i
 
             if now >= execute_at
-              client.publish(subject, message.data, headers: headers)
+              client.publish(subject, message.data, header: headers)
               message.ack
             else
               message.nak(delay: Config.to_ns(execute_at - now))
@@ -49,6 +90,7 @@ module Cosmo
             # to escape #each and kill this thread — schedule_loop only runs once per processor,
             # so an unhandled exception would silently stop all future scheduled-job dispatch.
             Logger.error e
+            Cosmo.handle_error(e, { source: :scheduler, subject: message.subject })
             message.nak rescue nil
           end
 
@@ -56,58 +98,49 @@ module Cosmo
         end
       end
 
-      def process(messages, _) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+      def process(messages, _) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
         message = messages.first
         Logger.debug "received messages #{messages.inspect}"
         data = Utils::Json.parse(message.data)
-        unless data
-          Logger.error ArgumentError.new("malformed payload")
-          move_message(message)
-          return
-        end
+        return reject_message(message, ArgumentError.new("malformed payload")) unless data
 
         worker_class = Utils::String.safe_constantize(data[:class])
         unless worker_class
-          Logger.error ArgumentError.new("#{data[:class]} class not found")
-          move_message(message, data)
+          reject_message(message, ArgumentError.new("#{data[:class]} class not found"), data)
           notify_batch(data, success: false)
           return
         end
 
-        if worker_class.limits_concurrency?
-          slot = acquire_concurrency_slot(worker_class, message, data)
-          return if slot == false
-        end
-
-        duration = worker_class.default_options[:limit]&.dig(:duration)&.to_i
-
-        with_stats(message) do
+        begin
           sw = stopwatch
           Logger.with(jid: data[:jid])
           Logger.info "start"
 
           instance = build_worker(worker_class, data, message)
-          perform_job(instance, data: data, message: message, duration: duration)
+          Config.server_middleware.invoke(instance, data, message) do
+            perform_job(instance, data: data, message: message)
+          end
 
           message.ack
           notify_batch(data, success: true)
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "done" }
-          true
+        rescue Requeue => e
+          message.nak(delay: Config.to_ns(e.delay))
+          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "requeue[#{e.delay}s]" }
         rescue Timeout::Error => e
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail[timeout]" }
-          dropped = handle_failure(worker_class, message, data, e)
-          false if dropped
+          Cosmo.handle_error(e, error_context(message, data))
+          handle_failure(worker_class, message, data, e)
         rescue StandardError => e
           Logger.debug e
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail[error]" }
-          dropped = handle_failure(worker_class, message, data, e)
-          false if dropped
+          Cosmo.handle_error(e, error_context(message, data))
+          handle_failure(worker_class, message, data, e)
         rescue Exception # rubocop:disable Lint/RescueException
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail[exception]" }
           raise
         end
       ensure
-        Limit.instance.release(slot) if slot
         Logger.without(:jid)
         Logger.debug "processed message #{message.inspect}"
       end
@@ -115,34 +148,13 @@ module Cosmo
       def build_worker(worker_class, data, message)
         worker_class.new.tap do |worker|
           worker.jid = data[:jid]
-          worker.enqueued_at = message.metadata.timestamp
+          worker.enqueued_at = enqueued_at(message)
           worker.attempt = message.metadata.num_delivered
-          worker.scheduled_by = message.header&.dig("Nats-Scheduler")
+          worker.scheduled_by = scheduled_by(message)
           worker.batch_id = data[:batch_id]
         end
       end
 
-      # Tries to acquire a concurrency slot for the job.
-      # Returns the slot key (String) on success, or false if all slots are
-      # taken (a message is NAK'd with a delay of +retry_in+ before returning
-      def acquire_concurrency_slot(worker_class, message, data)
-        options = worker_class.concurrency_options
-        key = worker_class.concurrency_key(data[:args])
-
-        slot = Limit.instance.acquire(key, jid: data[:jid], limit: options[:limit], duration: options[:duration])
-        return slot if slot
-
-        message.nak(delay: Config.to_ns(options[:retry_in]))
-        Logger.debug "concurrency limit reached for #{data[:class]}, re-queueing back #{data[:jid]}"
-        false
-      rescue NATS::Error => e
-        # Unexpected KV failure (e.g. transient NATS error). NAK immediately so
-        # the message is retried rather than stuck in-flight until ack_wait expires.
-        Logger.error e
-        message.nak
-        false
-      end
-
       def handle_failure(worker_class, message, data, exception) # rubocop:disable Naming/PredicateMethod
         current_attempt = message.metadata.num_delivered
         desired_retries = data[:retry].to_i + 1
@@ -154,7 +166,7 @@ module Cosmo
         end
 
         warn_capped(message, data, capped_at) if capped_at
-        data[:dead] ? move_message(message, data) : drop_message(message, data)
+        data[:dead] ? move_message(message, data, exception) : drop_message(message, data)
         notify_batch(data, success: false)
         true
       end
@@ -196,6 +208,7 @@ module Cosmo
         delay.is_a?(Numeric) && delay.positive? ? delay : default_retry_delay(current_attempt)
       rescue StandardError => e
         Logger.error e
+        Cosmo.handle_error(e, { source: :retry_in, class: data[:class], jid: data[:jid] })
         default_retry_delay(current_attempt)
       end
 
@@ -207,26 +220,62 @@ module Cosmo
         config = config.dup
         config[:batch_size] = 1
         config[:stream] = stream_name
-        config[:consumer] = "consumer-#{stream_name}"
+        config[:consumer] = consumer_name(stream_name)
         subscription = client.subscribe(config[:subject], config[:consumer], config.except(:subject, :priority, :stream, :batch_size, :consumer))
         [subscription, config, nil]
       end
 
       def drop_message(message, data)
         message.term
-        Logger.debug "job dropped #{data[:jid]}"
+        Logger.debug "job dropped #{data&.dig(:jid)}"
+      end
+
+      # Logs why a message can't be processed at all and parks it in the DLQ.
+      def reject_message(message, error, data = nil)
+        Logger.error error
+        Cosmo.handle_error(error, error_context(message, data, source: :reject))
+        move_message(message, data, error)
+      end
+
+      def error_context(message, data, source: :job)
+        meta = message.metadata
+        { source:, stream: meta.stream, subject: message.subject, attempt: meta.num_delivered }.merge(Hash(data))
       end
 
-      def move_message(message, data = nil)
+      def move_message(message, data = nil, exception = nil)
+        return drop_message(message, data) unless Config.dead.enabled
+
         klass = data ? Utils::String.underscore(data[:class]) : "default"
         headers = { "X-Stream" => message.metadata.stream, "X-Subject" => message.subject }
+        headers.merge!(Failure.headers(exception)) if exception
         Client.instance.publish("jobs.dead.#{klass}", message.data, header: headers)
         message.ack
         Logger.debug "job moved #{data&.dig(:jid)} to DLQ"
       end
 
       def scheduler?
-        true
+        @options.fetch(:scheduler, true)
+      end
+
+      # +Nats-Scheduler+ on a message NATS fired directly into a job stream, and +X-Scheduled-By+
+      # on one the scheduler dispatched onward, which may not carry the reserved header.
+      def scheduled_by(message)
+        header = message.header or return
+
+        header["X-Scheduled-By"] || header["Nats-Scheduler"]
+      end
+
+      # A cron message is enqueued when NATS fires it, not when the scheduler dispatches it on, so
+      # the firing forwards its own timestamp in +X-Enqueued-At+ to survive the re-publish.
+      def enqueued_at(message)
+        forwarded = message.header&.dig("X-Enqueued-At")
+
+        forwarded ? Time.at(forwarded.to_f) : message.metadata.timestamp
+      end
+
+      # Durable and per stream, so every process pulls from the same consumer and shares the work.
+      def consumer_name(stream_name)
+        "consumer-#{stream_name}"
       end
 
       def consumers
@@ -242,26 +291,14 @@ module Cosmo
         ENV.fetch("COSMO_JOBS_FETCH_TIMEOUT", 0.1).to_f
       end
 
-      def with_stats(message, &block)
-        API::Busy.instance.with(message) do
-          API::Counter.instance.with(&block)
-        end
-      end
-
       # @param job_instance [Cosmo::Job]
       # @param data [Hash]
       # @param message [NATS::Msg]
-      # @param duration [Float, nil]
       #
-      # rubocop:disable Lint/UnusedMethodArgument
-      def perform_job(job_instance, data:, message:, duration: nil)
-        if duration
-          Timeout.timeout(duration) { job_instance.perform(*data[:args]) }
-        else
-          job_instance.perform(*data[:args])
-        end
+      # rubocop:disable-next Lint/UnusedMethodArgument
+      def perform_job(job_instance, data:, message:)
+        job_instance.perform(*data[:args])
       end
-      # rubocop:enable Lint/UnusedMethodArgument
     end
   end
 end

data/lib/cosmo/job/requeue.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Job
+    # Raised by a middleware to put the job back on its stream for later: the processor naks it with +delay+ instead
+    # of failing it. The redelivery still counts as a delivery attempt toward the job's +retry+ and +max_deliver+.
+    # Raise it from a middleware that runs before the built-in stats middleware, or the attempt is counted as failed.
+    #
+    #   raise Cosmo::Job::Requeue.new(30) if maintenance?
+    class Requeue < StandardError
+      attr_reader :delay
+
+      # @param delay [Numeric] seconds before the job is delivered again
+      def initialize(delay = 0)
+        @delay = delay
+        super("requeued for #{delay}s")
+      end
+    end
+  end
+end

data/lib/cosmo/job/stream_filter.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Job
+    # Which job streams a worker subscribes to: every configured stream, or the subset named by
+    # +--streams+/+--stream+, falling back to the +COSMO_JOBS_STREAMS+ environment variable.
+    #
+    # The scheduled stream is a service stream. Naming it selects nothing, it runs unless the worker
+    # was started with +--no-scheduler+.
+    class StreamFilter
+      SCHEDULED = :scheduled
+      ENV_NAME = "COSMO_JOBS_STREAMS"
+
+      # @param names [Array<String>, nil] stream names from the command line
+      # @return [StreamFilter]
+      def self.from(names)
+        new(names || ENV[ENV_NAME]&.split(","))
+      end
+
+      # @param names [Array<String>, nil]
+      def initialize(names)
+        @names = Array(names).map(&:strip).reject(&:empty?)
+      end
+
+      # @return [Boolean] whether every configured stream is selected
+      def all?
+        @names.empty?
+      end
+
+      # @param stream_name [String, Symbol]
+      # @return [Boolean] whether the stream is selected
+      def include?(stream_name)
+        all? || @names.include?(stream_name.to_s)
+      end
+
+      # @return [StreamFilter] self, so it can be chained onto {.from}
+      # @raise [UnknownJobStreamError] when a name is not configured
+      def validate!
+        return self if all?
+
+        unknown = @names - configured - [SCHEDULED.to_s]
+        raise UnknownJobStreamError.new(unknown, configured) if unknown.any?
+
+        self
+      end
+
+      private
+
+      def configured
+        @configured ||= Config.dig(:consumers, :jobs)&.keys.to_a.map(&:to_s) - [SCHEDULED.to_s]
+      end
+    end
+  end
+end

data/lib/cosmo/job.rb

@@ -1,6 +1,10 @@
 # frozen_string_literal: true
 
 require "cosmo/job/data"
+require "cosmo/job/enqueuer"
+require "cosmo/job/requeue"
+require "cosmo/job/failure"
+require "cosmo/job/stream_filter"
 require "cosmo/job/limit"
 require "cosmo/job/processor"
 
@@ -94,30 +98,14 @@ module Cosmo
       end
 
       def perform(*args, async: true, **options)
-        batch = Batch.current if async
-        options[:batch_id] = batch.bid if batch
-        data = Data.new(name, args, default_options.merge(options))
-        unless async
-          payload = Utils::Json.parse(data.to_args[1])
-          raise ArgumentError, "Cannot parse payload" unless payload
-
-          new.perform(*payload[:args])
-          return
-        end
+        options = default_options.merge(options)
+        return Enqueuer.enqueue(name, args, options, batch: Batch.current) if async
 
-        publish(data, batch)
-      end
+        payload = Utils::Json.parse(Data.new(name, args, options).to_json)
+        raise ArgumentError, "Cannot parse payload" unless payload
 
-      # The batch is reserved a pending slot before we know the publish will
-      # succeed (must happen in that order -- see Batch#jobs). Roll it back
-      # if it never actually made it onto the stream, so the batch doesn't
-      # hang waiting for a completion that will never arrive.
-      def publish(data, batch)
-        batch&.register_job!
-        Publisher.publish_job(data)
-      rescue StandardError
-        batch&.rollback_job!
-        raise
+        new.perform(*payload[:args])
+        nil
       end
 
       def perform_async(*args)

data/lib/cosmo/logger.rb

@@ -61,7 +61,12 @@ module Cosmo
     class << self
       extend Forwardable
 
-      delegate %i[info error debug warn fatal trace] => :instance
+      delegate %i[info error debug warn fatal] => :instance
+    end
+
+    # Dropped when the logger has no TRACE level (a plain +::Logger+), rather than flooding its DEBUG.
+    def self.trace(...)
+      instance.trace(...) if instance.respond_to?(:trace)
     end
 
     def self.with(...)
@@ -82,10 +87,18 @@ module Cosmo
     def self.coerce_level(level)
       level.to_s.downcase == "trace" ? TRACE : ::Logger::Severity.coerce(level)
     end
-    private_class_method :coerce_level
 
+    # @param logger [::Logger] any stdlib-compatible logger. Its formatter is left alone, so context such as
+    #   +jid+ and +elapsed+ is only printed when it uses {SimpleFormatter}.
     def self.instance=(logger)
       @instance = logger
     end
+
+    # +COSMO_LOG_LEVEL+, when set, wins over +level+.
+    #
+    # @param level [Symbol, String, Integer] +:trace+, +:debug+, +:info+, +:warn+, +:error+, or +:fatal+
+    def self.level=(level)
+      instance.level = coerce_level(ENV.fetch("COSMO_LOG_LEVEL", level))
+    end
   end
 end

data/lib/cosmo/middleware/busy.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Middleware
+    # Lists the job on the Web UI's Busy page while it runs.
+    class Busy
+      # @param _job [Cosmo::Job]
+      # @param _data [Hash]
+      # @param message [NATS::Msg]
+      def call(_job, _data, message, &)
+        API::Stats::Busy.instance.with(message, &)
+      end
+    end
+  end
+end

data/lib/cosmo/middleware/chain.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Middleware
+    # An ordered list of middleware classes wrapped around job execution, outermost first.
+    #
+    # A middleware is any class whose instances respond to +call(job, data, message)+ and yield to run
+    # the rest of the chain. Every invocation builds fresh instances, so a middleware may keep per-job
+    # state in instance variables. Adding a class that is already in the chain moves it instead.
+    #
+    #   class Timing
+    #     def initialize(threshold:) = @threshold = threshold
+    #
+    #     def call(job, data, message)
+    #       yield
+    #     end
+    #   end
+    #
+    #   Cosmo.configure do |config|
+    #     config.server_middleware { |с| с.add Timing, threshold: 5 }
+    #   end
+    class Chain
+      include Enumerable
+
+      Entry = Struct.new(:klass, :args, :kwargs) do
+        def build
+          klass.new(*args, **kwargs)
+        end
+      end
+
+      # @yieldparam chain [Chain]
+      def initialize
+        @entries = []
+        yield self if block_given?
+      end
+
+      # @yieldparam entry [Entry]
+      def each(&)
+        @entries.each(&)
+      end
+
+      # Appends +klass+, so it runs innermost, closest to the job.
+      #
+      # @param klass [Class]
+      # @return [Chain]
+      def add(klass, *args, **kwargs)
+        remove(klass)
+        @entries << Entry.new(klass, args, kwargs)
+        self
+      end
+
+      # Inserts +klass+ first, so it runs outermost.
+      #
+      # @param klass [Class]
+      # @return [Chain]
+      def prepend(klass, *args, **kwargs)
+        remove(klass)
+        @entries.unshift(Entry.new(klass, args, kwargs))
+        self
+      end
+
+      # Inserts +klass+ right before +existing+, or first when +existing+ is not in the chain.
+      #
+      # @param existing [Class]
+      # @param klass [Class]
+      # @return [Chain]
+      def insert_before(existing, klass, *args, **kwargs)
+        entry = take(klass, args, kwargs)
+        @entries.insert(index(existing) || 0, entry)
+        self
+      end
+
+      # Inserts +klass+ right after +existing+, or last when +existing+ is not in the chain.
+      #
+      # @param existing [Class]
+      # @param klass [Class]
+      # @return [Chain]
+      def insert_after(existing, klass, *args, **kwargs)
+        entry = take(klass, args, kwargs)
+        position = index(existing)
+        @entries.insert(position ? position + 1 : @entries.size, entry)
+        self
+      end
+
+      # @param klass [Class]
+      # @return [Chain]
+      def remove(klass)
+        @entries.delete_if { _1.klass == klass }
+        self
+      end
+
+      # @param klass [Class]
+      # @return [Boolean]
+      def exists?(klass)
+        @entries.any? { _1.klass == klass }
+      end
+      alias include? exists?
+
+      # @return [Boolean]
+      def empty?
+        @entries.empty?
+      end
+
+      # @return [Chain]
+      def clear
+        @entries.clear
+        self
+      end
+
+      # Runs the block wrapped in every middleware, passing +args+ to each one's +call+.
+      #
+      # @return [Object] the outermost middleware's return value, normally the block's
+      def invoke(*args, &)
+        return yield if empty?
+
+        traverse(map(&:build), 0, args, &)
+      end
+
+      private
+
+      def traverse(chain, position, args, &block)
+        return yield if position >= chain.size
+
+        chain[position].call(*args) { traverse(chain, position + 1, args, &block) }
+      end
+
+      def index(klass)
+        @entries.index { _1.klass == klass }
+      end
+
+      def take(klass, args, kwargs)
+        position = index(klass)
+        position ? @entries.delete_at(position) : Entry.new(klass, args, kwargs)
+      end
+    end
+  end
+end

data/lib/cosmo/middleware/limit.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require "timeout"
+
+module Cosmo
+  module Middleware
+    # Enforces a job class's +limit:+ option: a concurrency slot per key, requeueing the job while every slot is
+    # taken, and the +duration+ timeout. Jobs without +limit:+ pass straight through. It runs first in the chain,
+    # so a requeued job never reaches the stats or user middleware.
+    class Limit
+      # @param job [Cosmo::Job]
+      # @param data [Hash]
+      # @param _message [NATS::Msg]
+      # @raise [Job::Requeue] while every concurrency slot is taken
+      def call(job, data, _message, &)
+        limit = job.class.default_options[:limit]
+        return yield unless limit
+
+        slot = acquire(job.class, data)
+        timeout(limit[:duration], &)
+      ensure
+        Job::Limit.instance.release(slot) if slot
+      end
+
+      private
+
+      def acquire(job_class, data)
+        options = job_class.concurrency_options
+        return unless options
+
+        key = job_class.concurrency_key(data[:args])
+        slot = Job::Limit.instance.acquire(key, jid: data[:jid], limit: options[:limit], duration: options[:duration])
+        raise Job::Requeue, options[:retry_in] unless slot
+
+        slot
+      rescue NATS::Error => e
+        Logger.error e
+        Cosmo.handle_error(e, { source: :limit, class: job_class.name, jid: data[:jid] })
+        raise Job::Requeue
+      end
+
+      def timeout(duration, &)
+        return yield unless duration
+
+        seconds = duration.to_i
+        Timeout.timeout(seconds, Timeout::Error, "execution expired after the #{seconds}s duration limit", &)
+      end
+    end
+  end
+end