cosmonats 0.2.0 → 0.4.0

data/lib/cosmo/client.rb

@@ -38,18 +38,63 @@ module Cosmo
       js.delete_stream(name, params)
     end
 
+    def update_stream(name, config)
+      js.update_stream(name: name, **config)
+    end
+
+    # Create/update a stream, falling back to create when there's no stream.
+    # @param name [String] Stream name
+    # @param config [Hash] Full desired stream configuration
+    def setup_stream(name, config)
+      update_stream(name, config)
+    rescue NATS::JetStream::Error::StreamNotFound
+      create_stream(name, config)
+    end
+
+    # Return all subjects in +stream_name+ that match +filter+ using NATS's
+    # subjects_filter on STREAM.INFO (requires NATS ≥ 2.9).
+    # @return [Array<String>]
+    def cron_subjects_in_stream(stream_name, filter)
+      payload = Utils::Json.dump({ subjects_filter: filter })
+      resp = nc.request("$JS.API.STREAM.INFO.#{stream_name}", payload)
+      data = Utils::Json.parse(resp.data, symbolize_names: false)
+      (data&.dig("state", "subjects") || {}).keys
+    rescue StandardError
+      []
+    end
+
     def list_streams
       response = nc.request("$JS.API.STREAM.LIST", "")
       data = Utils::Json.parse(response.data, symbolize_names: false)
       return [] if data.nil? || data["streams"].nil?
 
-      data["streams"].filter_map { _1.dig("config", "name") }
+      data["streams"]
+    end
+
+    def pause_stream(name)
+      config = stream_info(name).config.to_h
+      config[:metadata] ||= {}
+      config[:metadata][:"_cosmo.paused"] = "true"
+      update_stream(name, config)
+    end
+
+    def unpause_stream(name)
+      config = stream_info(name).config.to_h
+      config[:metadata] ||= {}
+      config[:metadata].delete(:"_cosmo.paused")
+      update_stream(name, config)
+    end
+
+    def stream_paused?(name)
+      stream_info(name).config.metadata&.[](:"_cosmo.paused") == "true"
+    rescue NATS::IO::Timeout
+      false
     end
 
     def list_consumers(stream_name)
       response = nc.request("$JS.API.CONSUMER.LIST.#{stream_name}", "")
-      data = Utils::Json.parse(response.data, symbolize_names: false)
-      data["consumers"]
+      data = Utils::Json.parse(response.data, default: {}, symbolize_names: false)
+      Array(data["consumers"])
     end
 
     def consumer_info(stream_name, consumer_name)
@@ -74,14 +119,39 @@ module Cosmo
       result["purged"] # number of messages purged
     end
 
-    def kv(name, **options)
+    def kv(name, allow_msg_ttl: false, **options)
       js.key_value(name)
     rescue NATS::KeyValue::BucketNotFoundError
-      js.create_key_value({ bucket: name }.merge(options))
+      allow_msg_ttl ? create_kv_with_msg_ttl(name, **options) : js.create_key_value({ bucket: name }.merge(options))
     end
 
     def close
       nc.close
     end
+
+    private
+
+    # NOTE: KV manager in nats-pure hardcodes the fields it copies into StreamConfig,
+    # so `allow_msg_ttl` is never forwarded via create_key_value. Send the raw stream-create API request instead.
+    def create_kv_with_msg_ttl(name, **options)
+      payload = Utils::Json.dump({
+        name: "KV_#{name}",
+        subjects: ["$KV.#{name}.>"],
+        storage: "file",
+        allow_direct: true,
+        allow_msg_ttl: true,
+        allow_rollup_hdrs: true,
+        max_msgs_per_subject: 1
+      }.merge(options))
+      resp = nc.request("$JS.API.STREAM.CREATE.KV_#{name}", payload)
+      result = Utils::Json.parse(resp.data, symbolize_names: false)
+      if result&.dig("error")
+        msg = result.dig("error", "description").to_s
+        # Two worker processes starting simultaneously can both attempt creation.
+        # If another process won the race, fall back to looking up the existing bucket.
+        raise NATS::JetStream::Error, msg unless msg.match?(/already in use|already exists/i)
+      end
+      js.key_value(name)
+    end
   end
 end

data/lib/cosmo/config.rb

@@ -4,7 +4,7 @@ require "yaml"
 require "forwardable"
 
 module Cosmo
-  class Config
+  class Config < ::Hash
     NANO = 1_000_000_000
     DEFAULT_PATH = "config/cosmo.yml"
 
@@ -31,11 +31,20 @@ module Cosmo
       end
 
       config[:setup]&.each_key do |type|
+        next if type == :cron
+
         config[:setup][type]&.each_key do |name|
           c = config[:setup][type][name]
           c[:max_age] = c[:max_age].to_i * NANO if c[:max_age]
           c[:duplicate_window] = c[:duplicate_window].to_i * NANO if c[:duplicate_window]
           c[:subjects] = c[:subjects].map { |s| format(s, name: name) } if c[:subjects]
+
+          next unless type == :jobs # Every jobs stream supports NATS 2.14 message scheduling.
+
+          c[:allow_msg_schedules] = true
+          cron_subject = "#{API::Cron::Entry::SUBJECT_PREFIX}.#{name}.>"
+          c[:subjects] = Array(c[:subjects])
+          c[:subjects] << cron_subject unless c[:subjects].include?(cron_subject)
         end
       end
     end
@@ -59,45 +68,18 @@ module Cosmo
       @instance ||= new
     end
 
-    def self.system
-      @system ||= {}
-    end
-
-    def initialize
-      @config = nil
-      @system = {}
-      @defaults = self.class.parse_file(File.expand_path("defaults.yml", __dir__))
-    end
-
-    def [](key)
-      dig(key)
-    end
-
-    def fetch(key, default = nil)
-      return @config.fetch(key, default) if @config && Utils::Hash.keys?(@config, key)
-
-      @defaults.fetch(key, default)
-    end
-
-    def dig(*keys)
-      return @config&.dig(*keys) if @config && Utils::Hash.keys?(@config, *keys)
-
-      @defaults.dig(*keys)
-    end
-
-    def to_h
-      Utils::Hash.merge(@defaults, @config)
+    def self.internal
+      @internal ||= {}
     end
 
     def set(...)
-      @config ||= {}
-      Utils::Hash.set(@config, ...)
+      Utils::Hash.set(self, ...)
     end
 
     def load(path = nil)
       return unless path
 
-      @config = self.class.parse_file(path)
+      replace(self.class.parse_file(path))
     end
   end
 end

data/lib/cosmo/engine.rb

@@ -25,7 +25,7 @@ module Cosmo
 
     def run(type, options)
       handler = Utils::Signal.trap(:INT, :TERM)
-      Logger.info "Starting processing, hit Ctrl-C to stop"
+      Logger.info "Starting processing, hit Ctrl-C to stop [concurrency=#{@concurrency}]"
 
       processor_classes = type && PROCESSORS.key?(type.to_sym) ? [PROCESSORS[type.to_sym]] : PROCESSORS.values
       @processors = processor_classes.map { _1.run(@pool, @running, options) }

data/lib/cosmo/job/data.rb

@@ -5,7 +5,7 @@ require "json"
 module Cosmo
   module Job
     class Data
-      DEFAULTS = { stream: :default, retry: 3, dead: true }.freeze
+      DEFAULTS = { stream: :default, retry: 3, dead: true, limit: nil }.freeze
 
       attr_reader :jid
 

data/lib/cosmo/job/limit.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Job
+    # Distributed concurrency limiter backed by NATS Key-Value with per-message TTL.
+    #
+    # Each unit of concurrency is a numbered KV slot:
+    #   "{concurrency_key}/0", "{concurrency_key}/1", ..., "{concurrency_key}/{limit-1}"
+    #
+    # Acquiring a slot is a single atomic `set` (CAS with last-revision=0).
+    # Only one worker can win a given slot; losers try the next number.
+    # When a job finishes the slot is deleted; if the worker crashes NATS
+    # expires it automatically via the per-message Nats-TTL header.
+    class Limit
+      BUCKET = "cosmo_jobs_limits"
+
+      def self.instance
+        @instance ||= new
+      end
+
+      def initialize
+        @kv = API::KV.new(BUCKET, allow_msg_ttl: true)
+      end
+
+      # Try to acquire one of the numbered slots for +key+.
+      #
+      # @param key      [String]  concurrency key
+      # @param jid      [String]  stored as the slot value for observability
+      # @param limit    [Integer] number of slots (0 … limit-1)
+      # @param duration [Integer] seconds before the slot is auto-expired by NATS
+      # @return [String, nil] the acquired slot key, or nil when all slots are taken
+      def acquire(key, jid:, limit:, duration:)
+        0.upto(limit - 1) do |i|
+          slot = "#{key}/#{i}"
+          @kv.set(slot, jid, ttl: duration)
+          return slot
+        rescue NATS::KeyValue::KeyWrongLastSequenceError
+          next # slot is live, try the next one
+        end
+        nil # all slots occupied
+      end
+
+      # Release a previously acquired slot.
+      def release(slot)
+        @kv.delete(slot)
+      rescue NATS::Error
+        # best effort — slot TTL will reclaim it if delete fails
+      end
+    end
+  end
+end

data/lib/cosmo/job/processor.rb

@@ -1,75 +1,36 @@
 # frozen_string_literal: true
 
+require "timeout"
+
 module Cosmo
   module Job
-    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
-
+    class Processor < ::Cosmo::Processor
       private
 
-      def run_loop
-        @work_thread = Thread.new { work_loop }
-        @schedule_thread = Thread.new { schedule_loop }
-      end
-
       def setup
+        # Initialize singletons before starting to process messages
+        API::Busy.instance
+        API::Counter.instance
+        Limit.instance
+
         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)
-          @weights += ([stream_name] * priority.to_i) if priority
-          subscription = client.subscribe(subject, consumer_name, config)
-          @consumers << [subscription, stream_name]
-        end
-      end
+          next if stream_name == :scheduled # scheduled jobs are handled in schedule_loop
 
-      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?
-
-            begin
-              timeout = ENV.fetch("COSMO_JOBS_FETCH_TIMEOUT", 0.1).to_f
-              @pool.post do
-                subscription = @consumers.find { |(_, sn)| sn == stream_name }&.first
-                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
-
-            break unless running?
-          end
+          @consumers << subscribe(stream_name, config)
         end
       end
 
       def schedule_loop # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/MethodLength, Metrics/AbcSize
+        config = Config.dig(:consumers, :jobs, :scheduled)
+        return unless config
+
+        subscription, = subscribe(:scheduled, config)
         while running?
           break unless running?
 
           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(subscription, batch_size: 100, timeout:)
           messages&.each do |message|
             headers = message.header.except("X-Stream", "X-Subject", "X-Execute-At", "Nats-Expected-Stream")
@@ -90,7 +51,7 @@ module Cosmo
         end
       end
 
-      def process(messages) # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
+      def process(messages, _) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         message = messages.first
         Logger.debug "received messages #{messages.inspect}"
         data = Utils::Json.parse(message.data)
@@ -107,30 +68,67 @@ 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
           sw = stopwatch
           Logger.with(jid: data[:jid])
           Logger.info "start"
           instance = worker_class.new
           instance.jid = data[:jid]
-          instance.perform(*data[:args])
+          if duration
+            Timeout.timeout(duration) { instance.perform(*data[:args]) }
+          else
+            instance.perform(*data[:args])
+          end
           message.ack
           Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "done" }
           true
+        rescue Timeout::Error
+          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail[timeout]" }
+          dropped = handle_failure(message, data)
+          false if dropped
         rescue StandardError => e
           Logger.debug e
-          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail" }
+          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail[error]" }
           dropped = handle_failure(message, data)
           false if dropped
         rescue Exception # rubocop:disable Lint/RescueException
-          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail" }
+          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
 
+      # Tries to acquire a concurrency slot for the job.
+      # Returns the slot key (String) on success, or false if all slots are
+      # taken (message is NAK'd with a delay equal to +duration+ 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: options[:duration] * Config::NANO)
+        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(message, data) # rubocop:disable Naming/PredicateMethod
         current_attempt = message.metadata.num_delivered
         max_retries = data[:retry].to_i + 1
@@ -139,7 +137,7 @@ module Cosmo
           # 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
+          delay_ns = ((current_attempt**4) + 15) * Config::NANO
           message.nak(delay: delay_ns)
           return false
         end
@@ -148,6 +146,15 @@ module Cosmo
         true
       end
 
+      def subscribe(stream_name, config)
+        config = config.dup
+        config[:batch_size] = 1
+        config[:stream] = stream_name
+        consumer_name = "consumer-#{stream_name}"
+        subscription = client.subscribe(config[:subject], consumer_name, config.except(:subject, :priority, :stream, :batch_size))
+        [subscription, config, nil]
+      end
+
       def drop_message(message, data)
         message.term
         Logger.debug "job dropped #{data[:jid]}"
@@ -161,16 +168,28 @@ module Cosmo
         Logger.debug "job moved #{data&.dig(:jid)} to DLQ"
       end
 
+      def scheduler?
+        true
+      end
+
+      def consumers
+        @weights ||= @consumers.filter_map { |(_, c, _)| [c[:stream]] * [c[:priority].to_i, 1].max }.flatten
+        @weights.shuffle.map { |s| @consumers.find { |(_, c, _)| c[:stream] == s } }
+      end
+
+      def fetch_subjects(config)
+        config[:subject]
+      end
+
+      def fetch_timeout(_config)
+        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
-
-      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

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "cosmo/job/data"
+require "cosmo/job/limit"
 require "cosmo/job/processor"
 
 module Cosmo
@@ -10,8 +11,56 @@ module Cosmo
     end
 
     module ClassMethods
-      def options(stream: nil, retry: nil, dead: nil)
-        default_options.merge!({ stream:, retry:, dead: }.compact)
+      # @option config [Symbol]  :stream  NATS stream to publish to (default: :default)
+      # @option config [Integer] :retry   max delivery attempts before giving up (default: 3)
+      # @option config [Boolean] :dead    move to dead-letter stream after retries exhausted (default: true)
+      # @option config [Hash]    :limit   execution limits:
+      #
+      #   limit: { duration: 30 }
+      #   limit: { duration: 30, concurrency: 3 }
+      #   limit: { duration: 30, concurrency: { to: 3, key: ->(id) { id } } }
+      #
+      # @option config [Integer] :"limit[:duration]"    hard execution timeout in seconds. The job thread is
+      #   killed after this many seconds and counts as a failed attempt (retried with exponential backoff,
+      #   moved to DLQ after retries exhausted).
+      # @option config [Integer, Hash] :"limit[:concurrency]"  caps how many instances run at once across all
+      #   workers. Jobs that cannot acquire a slot are NAK'd with a delay equal to +duration+ so they are not
+      #   re-delivered until the slot is guaranteed free. Requires +duration+.
+      #   Pass an Integer for a class-wide cap, or <tt>{ to: N, key: ->(args) {} }</tt> to scope per key.
+      def options(**config)
+        if config[:limit] && config.dig(:limit, :concurrency) && !config.dig(:limit, :duration).to_i.positive?
+          raise ArgumentError, "limit: duration is required when concurrency is set"
+        end
+
+        default_options.merge!(config)
+      end
+      alias cosmo_options options
+
+      def limits_concurrency?
+        !!concurrency_options
+      end
+
+      # Returns a normalized concurrency config hash, or +nil+ when not configured.
+      # Always contains +:limit+, +:key+, and +:duration+.
+      def concurrency_options
+        value = default_options.dig(:limit, :concurrency)
+        duration = default_options.dig(:limit, :duration).to_i
+        return unless value
+
+        case value
+        when Integer then { limit: value, key: nil, duration: duration }
+        when Hash    then { limit: value.fetch(:to), key: value[:key], duration: duration }
+        end
+      end
+
+      # Derive the fully-scoped concurrency key for a given args array.
+      def concurrency_key(args)
+        config = concurrency_options
+        return unless config
+
+        base = Utils::String.underscore(name)
+        suffix = config[:key]&.call(*args)
+        suffix ? "#{base}/#{suffix}" : base
       end
 
       def perform(*args, async: true, **options)

data/lib/cosmo/logger.rb

@@ -60,7 +60,10 @@ module Cosmo
     end
 
     def self.instance
-      @instance ||= ::Logger.new($stdout).tap { _1.formatter = SimpleFormatter.new }
+      @instance ||= ::Logger.new($stdout).tap do |logger|
+        logger.formatter = SimpleFormatter.new
+        logger.level = ::Logger::Severity.coerce(ENV.fetch("COSMO_LOG_LEVEL", "info"))
+      end
     end
 
     def self.instance=(logger)