cosmonats 0.1.3 → 0.2.0

data/lib/cosmo/job/processor.rb

@@ -2,22 +2,32 @@
 
 module Cosmo
   module Job
-    class Processor < ::Cosmo::Processor
+    class Processor < ::Cosmo::Processor # rubocop:disable Metrics/ClassLength
       def initialize(pool, running, options)
         super
         @weights = []
       end
 
+      def stop(timeout = Config[:timeout])
+        @running.make_false
+        @pool.shutdown
+        @consumers.each { |(s, _)| s.unsubscribe rescue nil }
+        @pool.wait_for_termination(timeout)
+        [@work_thread, @schedule_thread].compact.each { _1.join(timeout) || _1.kill }
+        @consumers.clear
+      end
+
       private
 
       def run_loop
-        Thread.new { work_loop }
-        Thread.new { schedule_loop }
+        @work_thread = Thread.new { work_loop }
+        @schedule_thread = Thread.new { schedule_loop }
       end
 
       def setup
         jobs_config = Config.dig(:consumers, :jobs)
         jobs_config&.each do |stream_name, config|
+          config = config.dup
           consumer_name = "consumer-#{stream_name}"
           subject = config.delete(:subject)
           priority = config.delete(:priority)
@@ -28,7 +38,11 @@ module Cosmo
       end
 
       def work_loop # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/MethodLength, Metrics/AbcSize
+        shutdown = false
+
         while running?
+          break if shutdown
+
           @weights.shuffle.each do |stream_name|
             break unless running?
 
@@ -36,10 +50,11 @@ module Cosmo
               timeout = ENV.fetch("COSMO_JOBS_FETCH_TIMEOUT", 0.1).to_f
               @pool.post do
                 subscription = @consumers.find { |(_, sn)| sn == stream_name }&.first
-                messages = fetch_messages(subscription, batch_size: 1, timeout:)
+                messages = lock(stream_name) { fetch(subscription, batch_size: 1, timeout:) }
                 process(messages) if messages&.any?
               end
             rescue Concurrent::RejectedExecutionError
+              shutdown = true
               break # pool doesn't accept new jobs, we are shutting down
             end
 
@@ -55,7 +70,7 @@ module Cosmo
           now = Time.now.to_i
           timeout = ENV.fetch("COSMO_JOBS_SCHEDULER_FETCH_TIMEOUT", 5).to_f
           subscription = @consumers.find { |(_, sn)| sn == :scheduled }&.first
-          messages = fetch_messages(subscription, batch_size: 100, timeout:)
+          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")
             stream, subject, execute_at = message.header.values_at("X-Stream", "X-Subject", "X-Execute-At")
@@ -80,17 +95,19 @@ module Cosmo
         Logger.debug "received messages #{messages.inspect}"
         data = Utils::Json.parse(message.data)
         unless data
-          Logger.debug ArgumentError.new("malformed payload")
+          Logger.error ArgumentError.new("malformed payload")
+          move_message(message)
           return
         end
 
         worker_class = Utils::String.safe_constantize(data[:class])
         unless worker_class
-          Logger.debug ArgumentError.new("#{data[:class]} class not found")
+          Logger.error ArgumentError.new("#{data[:class]} class not found")
+          move_message(message, data)
           return
         end
 
-        begin
+        with_stats(message) do
           sw = stopwatch
           Logger.with(jid: data[:jid])
           Logger.info "start"
@@ -99,10 +116,12 @@ module Cosmo
           instance.perform(*data[:args])
           message.ack
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "done" }
+          true
         rescue StandardError => e
           Logger.debug e
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail" }
-          handle_failure(message, data)
+          dropped = handle_failure(message, data)
+          false if dropped
         rescue Exception # rubocop:disable Lint/RescueException
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail" }
           raise
@@ -112,26 +131,46 @@ module Cosmo
         Logger.debug "processed message #{message.inspect}"
       end
 
-      def handle_failure(message, data) # rubocop:disable Metrics/AbcSize
+      def handle_failure(message, data) # rubocop:disable Naming/PredicateMethod
         current_attempt = message.metadata.num_delivered
         max_retries = data[:retry].to_i + 1
 
         if current_attempt < max_retries
-          # NATS will auto-retry based on max_deliver with exponential backoff
+          # NATS will auto-retry with delay (exponential backoff based on current attempt).
+          # When max_deliver is reached, NATS stops redelivering the message and marks it as "max deliveries exceeded".
+          # The message is effectively abandoned by NATS — it stays in the stream (consuming a slot) but will never be delivered again to that consumer.
           delay_ns = ((current_attempt**4) + 15) * 1_000_000_000
           message.nak(delay: delay_ns)
-          return
+          return false
         end
 
-        if data[:dead]
-          Client.instance.publish("jobs.dead.#{Utils::String.underscore(data[:class])}", message.data)
-          message.ack
-          Logger.debug "job moved #{data[:jid]} to DLQ"
-        else
-          message.term
-          Logger.debug "job dropped #{data[:jid]}"
+        data[:dead] ? move_message(message, data) : drop_message(message, data)
+        true
+      end
+
+      def drop_message(message, data)
+        message.term
+        Logger.debug "job dropped #{data[:jid]}"
+      end
+
+      def move_message(message, data = nil)
+        klass = data ? Utils::String.underscore(data[:class]) : "default"
+        headers = { "X-Stream" => message.metadata.stream, "X-Subject" => message.subject }
+        Client.instance.publish("jobs.dead.#{klass}", message.data, header: headers)
+        message.ack
+        Logger.debug "job moved #{data&.dig(:jid)} to DLQ"
+      end
+
+      def with_stats(message, &block)
+        API::Busy.instance.with(message) do
+          API::Counter.instance.with(&block)
         end
       end
+
+      def lock(stream_name, &)
+        @mutexes ||= Hash.new { |h, k| h[k] = Mutex.new }
+        @mutexes[stream_name].synchronize(&)
+      end
     end
   end
 end

data/lib/cosmo/job.rb

@@ -54,7 +54,7 @@ module Cosmo
       end
     end
 
-    attr_reader :jid
+    attr_accessor :jid
 
     def perform(...)
       raise NotImplementedError, "#{self.class}#perform must be implemented"

data/lib/cosmo/logger.rb

@@ -62,5 +62,9 @@ module Cosmo
     def self.instance
       @instance ||= ::Logger.new($stdout).tap { _1.formatter = SimpleFormatter.new }
     end
+
+    def self.instance=(logger)
+      @instance = logger
+    end
   end
 end

data/lib/cosmo/processor.rb

@@ -39,10 +39,16 @@ module Cosmo
       @running.true?
     end
 
-    def fetch_messages(subscription, batch_size:, timeout:)
+    def fetch(subscription, batch_size:, timeout:)
       subscription.fetch(batch_size, timeout:)
     rescue NATS::Timeout
       # No messages, continue
+    rescue StandardError => e
+      Logger.error "Snap! Error just happened"
+      Logger.error "#{e.class}: #{e.message}\n#{e.backtrace.join("\n")}"
+
+      backoff = ENV.fetch("COSMO_STREAMS_FETCH_BACKOFF", 5).to_f
+      sleep([timeout, backoff].max) # backoff before retry
     end
 
     def client

data/lib/cosmo/stream/data.rb

@@ -7,6 +7,7 @@ module Cosmo
     class Data
       DEFAULTS = {
         batch_size: 100,
+        fetch_timeout: 10.0,
         consumer: {
           ack_policy: "explicit",
           max_deliver: 1,

data/lib/cosmo/stream/processor.rb

@@ -19,21 +19,26 @@ module Cosmo
         setup_consumers
       end
 
-      def work_loop # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity
+      def work_loop # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+        shutdown = false
+
         while running?
+          break if shutdown
+
           @consumers.each do |(subscription, config, processor)|
             break unless running?
 
             begin
-              timeout = ENV.fetch("COSMO_STREAMS_FETCH_TIMEOUT", 0.1).to_f
               @pool.post do
+                timeout = convert_timeout(config[:fetch_timeout])
                 Logger.debug "fetching #{config.dig(:consumer, :subjects).inspect}, timeout=#{timeout}"
-                messages = fetch_messages(subscription, batch_size: config[:batch_size], timeout:)
+                messages = fetch(subscription, batch_size: config[:batch_size], timeout:)
                 Logger.debug "fetched (#{messages&.size.to_i}) messages"
                 process(messages, processor) if messages&.any?
                 Logger.debug "processed (#{messages&.size.to_i}) messages"
               end
             rescue Concurrent::RejectedExecutionError
+              shutdown = true
               break # pool doesn't accept new jobs, we are shutting down
             end
 
@@ -95,6 +100,16 @@ module Cosmo
       def dynamic_config
         Config.system[:streams].map { _1.default_options.merge(class: _1) }
       end
+
+      def convert_timeout(value)
+        timeout = value.to_f
+        if timeout <= 0
+          Logger.warn "Ignoring `fetch_timeout: #{timeout}` (causes high CPU usage) with #{Data::DEFAULTS[:fetch_timeout]}s instead"
+          timeout = Data::DEFAULTS[:fetch_timeout].to_f
+        end
+
+        timeout
+      end
     end
   end
 end

data/lib/cosmo/stream.rb

@@ -12,9 +12,9 @@ module Cosmo
     end
 
     module ClassMethods
-      def options(stream: nil, consumer_name: nil, batch_size: nil, start_position: nil, consumer: nil, publisher: nil) # rubocop:disable Metrics/ParameterLists
+      def options(stream: nil, consumer_name: nil, batch_size: nil, fetch_timeout: nil, start_position: nil, consumer: nil, publisher: nil) # rubocop:disable Metrics/ParameterLists
         register
-        default_options.merge!({ stream:, consumer_name:, batch_size:, start_position:, consumer:, publisher: }.compact)
+        default_options.merge!({ stream:, consumer_name:, batch_size:, fetch_timeout:, start_position:, consumer:, publisher: }.compact)
       end
 
       def publish(data, subject: nil, **options)

data/lib/cosmo/utils/overrides.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+Cosmo::Utils::Warnings.silence do
+  members = NATS::JetStream::API::StreamConfig.members + [:allow_msg_counter]
+  NATS::JetStream::API::StreamConfig = Struct.new(*members, keyword_init: true) do
+    def initialize(opts = {})
+      rem = opts.keys - members
+      opts.delete_if { |k| rem.include?(k) }
+      super
+    end
+  end
+
+  members = NATS::JetStream::PubAck.members + [:val]
+  NATS::JetStream::PubAck = Struct.new(*members, keyword_init: true)
+end

data/lib/cosmo/utils/warnings.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Utils
+    module Warnings
+      module_function
+
+      def silence
+        verbose = $VERBOSE
+        $VERBOSE = nil
+        yield
+      ensure
+        $VERBOSE = verbose
+      end
+    end
+  end
+end

data/lib/cosmo/utils.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+require "cosmo/utils/hash"
+require "cosmo/utils/json"
+require "cosmo/utils/string"
+require "cosmo/utils/signal"
+require "cosmo/utils/warnings"
+require "cosmo/utils/stopwatch"
+require "cosmo/utils/thread_pool"
+
+module Cosmo
+  module Utils
+  end
+end

data/lib/cosmo/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Cosmo
-  VERSION = "0.1.3"
+  VERSION = "0.2.0"
 end