cosmonats 0.7.0 → 0.8.0

data/lib/cosmo/job/limit.rb

@@ -13,14 +13,14 @@ module Cosmo
     # expires it automatically via the per-message Nats-TTL header. Both
     # paths leave the slot equally empty -- no tombstone, no delete marker.
     class Limit
-      BUCKET = "cosmo_jobs_limits"
+      BUCKET = "_cosmolimits"
 
       def self.instance
         @instance ||= new
       end
 
       def initialize
-        @kv = API::KV.new(BUCKET, allow_msg_ttl: true)
+        @kv = API::KV.new(BUCKET, allow_msg_ttl: true, replicas: Config.replicas)
       end
 
       # Try to acquire one of the numbered slots for +key+.

data/lib/cosmo/job/processor.rb

@@ -32,6 +32,7 @@ module Cosmo
         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|
@@ -53,7 +54,7 @@ module Cosmo
       end
 
       def scheduled_config
-        Config.dig(:consumers, :jobs, StreamFilter::SCHEDULED)
+        Services.scheduled_consumer
       end
 
       def schedule_loop # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/MethodLength, Metrics/AbcSize
@@ -89,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
 
@@ -96,7 +98,7 @@ 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)
@@ -109,40 +111,36 @@ module Cosmo
           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
@@ -157,27 +155,6 @@ module Cosmo
         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
@@ -231,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
 
@@ -249,16 +227,24 @@ module Cosmo
 
       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, 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
@@ -305,24 +291,13 @@ module Cosmo
         ENV.fetch("COSMO_JOBS_FETCH_TIMEOUT", 0.1).to_f
       end
 
-      def with_stats(message, &block)
-        API::Stats::Busy.instance.with(message) do
-          API::Stats::Totals.instance.with(&block)
-        end
-      end
-
       # @param job_instance [Cosmo::Job]
       # @param data [Hash]
       # @param message [NATS::Msg]
-      # @param duration [Float, nil]
       #
       # rubocop:disable-next Lint/UnusedMethodArgument
-      def perform_job(job_instance, data:, message:, duration: nil)
-        if duration
-          Timeout.timeout(duration, Timeout::Error, "execution expired after the #{duration}s duration limit") { job_instance.perform(*data[:args]) }
-        else
-          job_instance.perform(*data[:args])
-        end
+      def perform_job(job_instance, data:, message:)
+        job_instance.perform(*data[:args])
       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.rb

@@ -1,6 +1,8 @@
 # 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"
@@ -96,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

data/lib/cosmo/middleware/metrics.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Middleware
+    # Records every execution for the web UI's Metrics tab while +config.metrics+ is enabled: its run time when it
+    # succeeds, a failure otherwise, and the wait from enqueue to the job's first delivery. ActiveJob jobs count under
+    # their own class.
+    class Metrics
+      # @param job [Cosmo::Job]
+      # @param data [Hash]
+      # @param _message [NATS::Msg]
+      def call(job, data, _message) # rubocop:disable Metrics/AbcSize
+        return yield unless Config.metrics.enabled
+
+        name = job_name(job, data)
+        wait_ms = (Time.now - job.enqueued_at) * 1000 if job.attempt == 1 && job.enqueued_at
+        started = clock
+        result = yield
+        API::Stats::Metrics.instance.record(name, exec_ms: clock - started, wait_ms:)
+        result
+      rescue Exception # rubocop:disable Lint/RescueException
+        API::Stats::Metrics.instance.record(name, wait_ms:, failed: true) if name
+        raise
+      end
+
+      private
+
+      def job_name(job, data)
+        return job.class.name unless defined?(ActiveJobAdapter::Executor) && job.is_a?(ActiveJobAdapter::Executor)
+
+        data.dig(:args, 0, :job_class).to_s
+      end
+
+      def clock
+        Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_millisecond)
+      end
+    end
+  end
+end

data/lib/cosmo/{sentry/job_processor_middleware.rb → middleware/sentry.rb}

@@ -2,28 +2,39 @@
 
 # rubocop:disable-next Metrics/MethodLength, Metrics/AbcSize
 module Cosmo
-  module Sentry
-    module JobProcessorMiddleware
+  module Middleware
+    # Wraps every job execution in a Sentry transaction. Errors reach Sentry through {ERROR_HANDLER}, which also
+    # captures those outside jobs (stream batches, fetches, the scheduler). Needs +sentry-ruby+, so it is not loaded
+    # by default:
+    #
+    #   require "cosmo/middleware/sentry"
+    #
+    #   Cosmo.configure do |config|
+    #     config.server_middleware { |chain| chain.add Cosmo::Middleware::Sentry }
+    #     config.error_handlers << Cosmo::Middleware::Sentry::ERROR_HANDLER
+    #   end
+    class Sentry
       NAME_PREFIX = "Cosmonats"
       OP_NAME = "queue.cosmonats"
       SPAN_ORIGIN = "auto.queue.cosmonats"
       STATUS_OK = 200
       STATUS_FAIL = 500
 
-      # @param job_instance [Cosmo::Job]
+      ERROR_HANDLER = lambda do |error, context|
+        next unless ::Sentry.initialized?
+
+        ::Sentry.capture_exception(error, contexts: { cosmonats: context }, hint: { background: true, integration: "cosmonats" })
+      end
+
+      # @param job [Cosmo::Job]
       # @param data [Hash]
       # @param message [NATS::Msg]
-      # @param duration [Float, nil]
-      def perform_job(job_instance, data:, message:, duration: nil)
-        unless ::Sentry.initialized?
-          super
-          return
-        end
+      def call(job, data, message)
+        return yield unless ::Sentry.initialized?
 
         scope = ::Sentry.get_current_scope
-        transaction_name = "#{NAME_PREFIX}/#{job_instance.class.name}"
+        transaction_name = "#{NAME_PREFIX}/#{job.class.name}"
         scope.set_transaction_name(transaction_name, source: :task)
-
         transaction = ::Sentry.start_transaction(
           name: scope.transaction_name,
           source: scope.transaction_source,
@@ -35,29 +46,14 @@ module Cosmo
         transaction&.set_data("messaging.message.retry.count", data[:retry] || 0)
 
         begin
-          super
-
+          result = yield
           transaction&.set_http_status(STATUS_OK)
           transaction&.finish
-        rescue StandardError => e
-          ::Sentry.capture_exception(
-            e,
-            contexts: {
-              cosmonats: data.merge(
-                nats_stream: message.metadata.stream,
-                nats_subject: message.subject,
-                timeout_duration: duration
-              )
-            },
-            hint: {
-              background: true,
-              integration: "cosmonats"
-            }
-          )
+          result
+        rescue StandardError
           transaction&.set_http_status(STATUS_FAIL)
           transaction&.finish
-
-          raise e
+          raise
         end
       end
     end

data/lib/cosmo/middleware/totals.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Middleware
+    # Counts every execution as processed or failed, so a job retried twice and then succeeding adds two failures
+    # and one processed. A counter that cannot be written is logged rather than raised, so stats never fail a job.
+    class Totals
+      # @param _job [Cosmo::Job]
+      # @param _data [Hash]
+      # @param _message [NATS::Msg]
+      def call(_job, _data, _message)
+        result = yield
+        increment(:processed)
+        result
+      rescue Exception # rubocop:disable Lint/RescueException
+        increment(:failed)
+        raise
+      end
+
+      private
+
+      def increment(key)
+        API::Stats::Totals.instance.increment(key)
+      rescue StandardError => e
+        Logger.debug "Totals #{key} counter error: #{e.class} #{e.message}"
+      end
+    end
+  end
+end

data/lib/cosmo/middleware.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+require "cosmo/middleware/chain"
+require "cosmo/middleware/limit"
+require "cosmo/middleware/busy"
+require "cosmo/middleware/totals"
+require "cosmo/middleware/metrics"