cosmonats 0.6.0 → 0.7.0

data/lib/cosmo/api/kv.rb

@@ -3,6 +3,8 @@
 module Cosmo
   module API
     class KV
+      LIMIT = 25
+
       attr_reader :kv
 
       def initialize(name, options = nil)
@@ -39,17 +41,19 @@ module Cosmo
         kv.delete(key)
       end
 
-      def keys(subject = nil, limit: 25)
-        results = []
-        watcher = kv.watch(subject || ">", ignore_deletes: true, meta_only: true)
-        watcher.each do |entry|
-          break unless entry
+      # Walks the bucket and collects live keys. Pass +limit: nil+ to walk the whole
+      # bucket -- a bounded +limit+ stops the watch early, so it can never be used
+      # to count (see #count).
+      def keys(subject = nil, limit: LIMIT, offset: 0)
+        walk(subject, limit:, offset:, meta_only: true).map(&:first)
+      end
 
-          results << entry.key
-          break if results.size >= limit
-        end
-        watcher.stop
-        results
+      # Same walk as #keys, but carries each value back with its key. The watch
+      # already has the values in hand, so callers that need them get one pass
+      # instead of a #keys walk followed by a #get per key.
+      # @return [Array<Array(String, String)>] +[key, value]+ pairs
+      def entries(subject = nil, limit: LIMIT, offset: 0)
+        walk(subject, limit:, offset:, meta_only: false)
       end
 
       # Writes a KV-Operation tombstone (same issue as #delete on ttl buckets).
@@ -70,7 +74,7 @@ module Cosmo
       end
 
       def count
-        keys.size
+        keys(limit: nil).size
       rescue NATS::KeyValue::NoKeysFoundError, NATS::JetStream::Error::NotFound
         0
       end
@@ -78,6 +82,22 @@ module Cosmo
 
       private
 
+      def walk(subject, limit:, offset:, meta_only:)
+        results = []
+        skipped = 0
+        watcher = kv.watch(subject || ">", ignore_deletes: true, meta_only:)
+        watcher.each do |entry|
+          break unless entry
+
+          next skipped += 1 if skipped < offset.to_i
+
+          results << [entry.key, entry.value]
+          break if limit && results.size >= limit
+        end
+        watcher.stop
+        results
+      end
+
       # CAS = Compare-And-Swap: publish +value+ with a per-message Nats-TTL,
       # but only if the subject's current last sequence matches +last_seq+
       # (sent as the Nats-Expected-Last-Subject-Sequence header). Raises

data/lib/cosmo/api/registry.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module API
+    # Base for a KV bucket of JSON entries that live only while their writer keeps refreshing them:
+    # a writer re-puts its entries faster than the bucket's TTL, and an entry whose writer is gone expires.
+    # Subclasses define +BUCKET+ and +TTL+.
+    class Registry
+      LIMIT = 25
+
+      def self.instance
+        @instance ||= new
+      end
+
+      def initialize
+        @kv = bucket
+      end
+
+      # Writes or refreshes an entry, recreating the bucket if it was deleted under a running writer.
+      #
+      # @param key [String, Integer]
+      # @param value [Hash]
+      # @return [void]
+      def put(key, value)
+        @kv.set(sanitize(key), Utils::Json.dump(value))
+      rescue NATS::JetStream::Error::NoStreamResponse
+        @kv = bucket
+        @kv.set(sanitize(key), Utils::Json.dump(value))
+      end
+
+      # Removes an entry, leaving no tombstone behind.
+      #
+      # @param key [String, Integer]
+      # @return [void]
+      def remove(key)
+        @kv.erase(sanitize(key))
+      end
+
+      # @param page [Integer, nil]
+      # @param limit [Integer]
+      # @return [Array<Hash>]
+      def list(page: nil, limit: LIMIT)
+        offset = ([page.to_i, 1].max - 1) * limit
+        @kv.entries(limit:, offset:).filter_map { Utils::Json.parse(_1.last) }
+      end
+
+      # @return [Array<Hash>] every live entry
+      def all
+        @kv.entries(limit: nil).filter_map { Utils::Json.parse(_1.last) }
+      end
+
+      # @return [Integer] the number of live entries
+      def size
+        @kv.size
+      end
+
+      private
+
+      def bucket
+        KV.new(self.class::BUCKET, { ttl: self.class::TTL })
+      end
+
+      def sanitize(key)
+        key.to_s.gsub(/[^\w.=-]/, "_")
+      end
+    end
+  end
+end

data/lib/cosmo/api/stats/busy.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require "socket"
+
+module Cosmo
+  module API
+    module Stats
+      # Jobs being executed right now, one entry per in-flight message, across all workers.
+      class Busy < Registry
+        TTL = 70
+        HEARTBEAT = 30
+        BUCKET = "cosmo_jobs_busy"
+
+        def initialize
+          super
+          @messages = {}
+        end
+
+        def with(message)
+          add(message)
+          yield
+        ensure
+          delete(message)
+        end
+
+        def add(message)
+          @thread ||= Thread.new { heartbeat_loop }
+          meta = message.metadata
+          seq = meta.sequence.stream
+          entry = { data: message.data, stream: meta.stream, worker: worker_id,
+                    started_at: Time.now.to_i, delivery: meta.num_delivered }
+          @messages[seq] = entry
+          put(seq, entry)
+        end
+
+        def delete(message)
+          seq = message.metadata.sequence.stream
+          @messages.delete(seq)
+          remove(seq)
+        end
+
+        def list(...)
+          super.map { _1.merge(data: Utils::Json.parse(_1[:data])) }
+        end
+
+        private
+
+        def heartbeat_loop
+          loop do
+            sleep(HEARTBEAT)
+            @messages.dup.each { |seq, entry| put(seq, entry) rescue StandardError }
+          rescue StandardError => e
+            Logger.debug "Busy heartbeat error: #{e.class} #{e.message}"
+          end
+        end
+
+        def worker_id
+          @worker_id ||= "#{Socket.gethostname}-#{::Process.pid}"
+        end
+      end
+    end
+  end
+end

data/lib/cosmo/api/stats/processes.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module API
+    module Stats
+      # Live worker processes, one entry per process, refreshed by Cosmo::Heartbeat.
+      class Processes < Registry
+        TTL = 60
+        BUCKET = "cosmo_processes"
+
+        # Sorted by host and pid before paging: every heartbeat rewrites an entry, so the bucket's own
+        # order shifts constantly and would move processes between pages.
+        #
+        # @param page [Integer, nil]
+        # @param limit [Integer]
+        # @return [Array<Hash>]
+        def list(page: nil, limit: LIMIT)
+          offset = ([page.to_i, 1].max - 1) * limit
+          all.sort_by { [_1[:hostname].to_s, _1[:pid].to_i] }.slice(offset, limit).to_a
+        end
+
+        # @param identity [String] a unique process id, e.g. +hostname-pid+
+        # @param info [Hash] the process details shown in the Web UI
+        # @return [void]
+        def register(identity, info)
+          put(identity, info)
+        end
+
+        # @param identity [String]
+        # @return [void]
+        def unregister(identity)
+          remove(identity)
+        end
+      end
+    end
+  end
+end

data/lib/cosmo/api/stats/totals.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module API
+    module Stats
+      # Processed and failed job totals across all workers.
+      class Totals < Counter
+        def self.instance
+          @instance ||= new("jobs")
+        end
+
+        # Counts the block's outcome: +true+ is processed, +false+ or a raise is failed.
+        #
+        # @return [void]
+        def with
+          result = yield
+          increment(:processed) if result == true
+          increment(:failed) if result == false
+        rescue Exception # rubocop:disable Lint/RescueException
+          increment(:failed)
+        end
+
+        # @return [Integer]
+        def processed
+          get(:processed)
+        end
+
+        # @return [Integer]
+        def failed
+          get(:failed)
+        end
+      end
+    end
+  end
+end

data/lib/cosmo/api/stats.rb

@@ -1,7 +1,8 @@
 # frozen_string_literal: true
 
-require "cosmo/api/counter"
-require "cosmo/api/busy"
+require "cosmo/api/stats/totals"
+require "cosmo/api/stats/busy"
+require "cosmo/api/stats/processes"
 
 module Cosmo
   module API
@@ -13,11 +14,11 @@ module Cosmo
       end
 
       def processed
-        Counter.instance.get(:processed)
+        Totals.instance.processed
       end
 
       def failed
-        Counter.instance.get(:failed)
+        Totals.instance.failed
       end
 
       def busy

data/lib/cosmo/api/stream.rb

@@ -36,7 +36,7 @@ module Cosmo
 
       def total
         all_msgs = info[:state].messages.to_i
-        cron_count = Client.instance.cron_subjects_in_stream(name, "#{Cron::Entry::SUBJECT_PREFIX}.#{name}.>").size
+        cron_count = Client.instance.cron_subjects_in_stream(name, "#{Cron::Entry::SUBJECT_PREFIX}.>").size
         [all_msgs - cron_count, 0].max
       rescue NATS::Error
         0

data/lib/cosmo/api.rb

@@ -1,8 +1,9 @@
 # frozen_string_literal: true
 
-require "cosmo/api/stream"
-require "cosmo/api/counter"
 require "cosmo/api/kv"
+require "cosmo/api/counter"
+require "cosmo/api/registry"
+require "cosmo/api/stream"
 require "cosmo/api/stats"
 require "cosmo/api/cron"
 require "cosmo/api/batch"

data/lib/cosmo/batch.rb

@@ -39,6 +39,13 @@ module Cosmo
   # 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.
   #
+  # You can pass your own bid instead of letting one be generated - useful
+  # when something else (e.g. an ActiveRecord id) needs to match it:
+  #
+  #   Cosmo::Batch.new(bid: your_id)
+  #
+  # Same meaning as a plain #new either way - it's still a fresh batch.
+  #
   # == How this is actually stored in NATS
   #
   # Every batch has an id (bid, a random hex string). State lives
@@ -102,8 +109,8 @@ module Cosmo
 
     attr_reader :bid, :parent_id
 
-    def initialize(parent: nil)
-      @bid = SecureRandom.hex(8)
+    def initialize(bid: nil, parent: nil)
+      @bid = 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

data/lib/cosmo/cli.rb

@@ -4,7 +4,7 @@ require "yaml"
 require "optparse"
 
 module Cosmo
-  # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/BlockLength
+  # rubocop:disable-next Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/BlockLength
   class CLI
     def self.run
       instance.run
@@ -14,15 +14,22 @@ module Cosmo
       @instance ||= new
     end
 
+    # @return [Array<String>, nil] the arguments the process was started with, before parsing consumed them
+    attr_reader :argv
+
     def run
+      @argv = ARGV.dup
       flags, command, options = parse
       return run_setup(flags) if flags[:setup]
 
       load_config(flags)
+      Engine.processors_for(command).each { _1.validate_options!(options) }
       puts self.class.banner
       boot_application
       require_path(flags[:require])
       Engine.run(command, options)
+    rescue Error => e
+      abort e.message
     end
 
     private
@@ -53,6 +60,7 @@ module Cosmo
       Config.load(path)
       Config.set(:concurrency, flags[:concurrency]) if flags[:concurrency]
       Config.set(:timeout, flags[:timeout]) if flags[:timeout]
+      Config.set(:http, :port, flags[:http_port]) if flags[:http_port]
     end
 
     def run_setup(flags)
@@ -69,9 +77,11 @@ module Cosmo
           first_line ? print("Stream is ready: #{name}") : print(", #{name}")
           first_line = false
         end
+        puts unless first_line
       end
 
-      puts
+      API::Counter.setup!
+
       schedules = Config.dig(:setup, :cron)&.reduce(0) do |sum, (name, entry)|
         class_name = entry.delete(:class)
         API::Cron.instance.upsert!(**entry, name: name, class_name: class_name)
@@ -129,6 +139,10 @@ module Cosmo
           flags[:timeout] = arg
         end
 
+        o.on "-p", "--http-port INT", Integer, "Start HTTP server with /health endpoint on port" do |arg|
+          flags[:http_port] = arg
+        end
+
         o.on "-C", "--config PATH", "Path to config file" do |arg|
           flags[:config_file] = arg
         end
@@ -155,8 +169,16 @@ module Cosmo
         OptionParser.new do |o|
           o.banner = "Usage: cosmo jobs [options]"
 
-          o.on "--stream NAME", "Job's stream" do |arg|
-            options[:stream] = arg
+          o.on "--streams NAMES", "Only subscribe to these job streams (comma-separated), instead of all configured streams" do |arg|
+            options[:streams] = Array(options[:streams]) | arg.split(",")
+          end
+
+          o.on "--stream NAME", "Same as --streams, for a single job stream" do |arg|
+            options[:streams] = Array(options[:streams]) | [arg]
+          end
+
+          o.on "--[no-]scheduler", "Dispatch scheduled jobs (default: yes)" do |arg|
+            options[:scheduler] = arg
           end
 
           o.on "--subject NAME", "Job's subject" do |arg|
@@ -214,7 +236,7 @@ module Cosmo
       require path
     end
 
-    # rubocop:disable Layout/TrailingWhitespace,Lint/IneffectiveAccessModifier
+    # rubocop:disable-next Layout/TrailingWhitespace,Lint/IneffectiveAccessModifier
     def self.banner
       <<-TEXT
                     .#%+:                                                  
@@ -254,7 +276,5 @@ module Cosmo
                  :.                                                        
       TEXT
     end
-    # rubocop:enable Layout/TrailingWhitespace,Lint/IneffectiveAccessModifier
   end
-  # rubocop:enable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/BlockLength
 end

data/lib/cosmo/client.rb

@@ -1,21 +1,49 @@
 # frozen_string_literal: true
 
+require "socket"
 require "nats/client"
 require "cosmo/utils/overrides"
 
 module Cosmo
   class Client
+    # Seconds to wait for a JetStream API reply. Override with COSMO_JS_TIMEOUT.
+    JS_TIMEOUT = 5
+    # Seconds to wait for the TCP connect+handshake. Override with COSMO_CONNECT_TIMEOUT.
+    CONNECT_TIMEOUT = 2
+
     def self.instance
       @instance ||= Client.new
     end
 
-    attr_reader :nc, :js
+    # Labels the connection in `nats server report connections` and /connz, so a
+    # misbehaving process can be found without cross-referencing host IPs. Derived from
+    # the running program, which separates a web process from a `cosmo` worker on the
+    # same host. Override with COSMO_CLIENT_NAME.
+    def self.default_name
+      ENV.fetch("COSMO_CLIENT_NAME") { "cosmo-#{program_name}-#{Socket.gethostname}-#{Process.pid}" }
+    end
+
+    def self.js_timeout
+      ENV.fetch("COSMO_JS_TIMEOUT", JS_TIMEOUT).to_i
+    end
+
+    def self.connect_timeout
+      ENV.fetch("COSMO_CONNECT_TIMEOUT", CONNECT_TIMEOUT).to_i
+    end
+
+    # $PROGRAM_NAME is "puma 7.2.1 (tcp://...)" under Puma and a path under the CLI.
+    def self.program_name
+      File.basename($PROGRAM_NAME.to_s.split.first.to_s, ".*").gsub(/[^\w.-]/, "")
+    end
+
+    attr_reader :nc, :js, :name
 
-    def initialize(nats_url: ENV.fetch("NATS_URL", "nats://localhost:4222"))
-      Logger.debug "Connecting to NATS server at #{nats_url}..."
-      @nc = NATS.connect(nats_url)
+    def initialize(nats_url: ENV.fetch("NATS_URL", "nats://localhost:4222"), name: self.class.default_name)
+      @name = name
+      Logger.debug "Connecting to NATS server at #{nats_url} as #{@name}..."
+      @nc = NATS.connect(nats_url, name: @name, connect_timeout: self.class.connect_timeout)
       Logger.debug "Connection established"
-      @js = @nc.jetstream
+      @js = @nc.jetstream(timeout: self.class.js_timeout)
     end
 
     def publish(subject, payload, **params)

data/lib/cosmo/config.rb

@@ -47,10 +47,14 @@ module Cosmo
           c[:duplicate_window] = to_ns(Utils::Duration.parse(c[:duplicate_window])) 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.
+          # The scheduled stream carries every cron template, so it alone enables NATS message
+          # scheduling. A schedule can only target a subject its own stream covers, and NATS
+          # rejects `discard: new` on any stream with scheduling enabled - keeping both here
+          # leaves the streams that actually run jobs free to choose their own discard policy.
+          next unless type == :jobs && name.to_s == Job::StreamFilter::SCHEDULED.to_s
 
           c[:allow_msg_schedules] = true
-          cron_subject = "#{API::Cron::Entry::SUBJECT_PREFIX}.#{name}.>"
+          cron_subject = "#{API::Cron::Entry::SUBJECT_PREFIX}.>"
           c[:subjects] = Array(c[:subjects])
           c[:subjects] << cron_subject unless c[:subjects].include?(cron_subject)
         end

data/lib/cosmo/engine.rb

@@ -13,38 +13,137 @@ module Cosmo
       instance.run(...)
     end
 
+    # Processor classes a +type+ starts: the matching one, or all of them when it names none.
+    #
+    # @param type [String, Symbol, nil]
+    # @return [Array<Class>]
+    def self.processors_for(type)
+      type && PROCESSORS.key?(type.to_sym) ? [PROCESSORS[type.to_sym]] : PROCESSORS.values
+    end
+
     def self.instance
       @instance ||= new
     end
 
+    attr_reader :concurrency
+
     def initialize
       @concurrency = Config.fetch(:concurrency, 1)
       @pool = Utils::ThreadPool.new(@concurrency)
       @running = Concurrent::AtomicBoolean.new
+      @quiet = Concurrent::AtomicBoolean.new
+      @http_server = nil
+      @heartbeat = nil
     end
 
     def run(type, options)
-      handler = Utils::Signal.trap(:INT, :TERM)
+      handler = Utils::Signal.trap(:INT, :TERM, :TSTP, :CONT, :USR1, :TTIN)
       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) }
+      @processors = self.class.processors_for(type).map { _1.run(@pool, @running, options, quiet: @quiet) }
       if @running.false?
         Logger.warn "Shutting down... (No processors are running)"
         return
       end
 
-      signal = handler.wait
+      @heartbeat = Heartbeat.new(self, options:).start
+      start_http_server
+
+      signal = handle_shutdown(handler)
       Logger.info "Shutting down... (#{signal} received)"
       shutdown
     end
 
+    def running?
+      @running.true?
+    end
+
+    # @return [String] +running+, +quiet+ (no new work fetched) or +stopping+
+    def state
+      return "stopping" unless running?
+
+      @quiet.true? ? "quiet" : "running"
+    end
+
+    # @return [Integer] threads processing messages right now, across all processors
+    def busy
+      Array(@processors).sum(&:busy)
+    end
+
+    # @return [Hash{Symbol => Array<String>}] what each running processor pulls from, keyed by processor type
+    def subscriptions
+      Array(@processors).to_h { [PROCESSORS.key(_1.class), _1.subscriptions] }
+    end
+
     def shutdown
       @running.make_false
+      @heartbeat&.beat
+      @http_server&.stop
       @pool.shutdown
       Logger.info "Pausing to allow jobs to finish..."
       @pool.wait_for_termination(Config[:timeout])
+      @heartbeat&.stop
       Logger.info "Bye!"
     end
+
+    private
+
+    # HTTPServer is autoloaded, so it's referenced only after the port is checked
+    # to keep its optional dependencies (rack, webrick) unloaded otherwise.
+    def start_http_server
+      port = Config.dig(:http, :port)
+      return unless port
+
+      host = Config.dig(:http, :host) || HTTPServer::DEFAULT_HOST
+      @http_server = HTTPServer.new(port: port, host: host).start
+    end
+
+    def handle_shutdown(handler)
+      loop do
+        signal = handler.wait
+        case signal.to_s
+        when "TSTP" then quiet
+        when "CONT" then unquiet
+        when "USR1" then drain_and_exit(handler)
+        when "TTIN" then dump_threads
+        else return signal
+        end
+      end
+    end
+
+    def quiet
+      return unless @quiet.make_true
+
+      Logger.info "Received TSTP, no new jobs will be fetched; finishing in-flight work"
+      @heartbeat&.beat
+    end
+
+    def unquiet
+      return unless @quiet.make_false
+
+      Logger.info "Received CONT, resuming normal fetching"
+      @heartbeat&.beat
+    end
+
+    def drain_and_exit(handler)
+      return unless @quiet.make_true
+
+      Logger.info "Received USR1, no new jobs will be fetched, exiting when work drains"
+      @heartbeat&.beat
+      Thread.new do
+        @pool.wait_idle
+        handler.push(:TERM)
+      end
+    end
+
+    def dump_threads
+      threads = Thread.list
+      Logger.warn "Received TTIN, dumping backtraces of #{threads.size} threads"
+      threads.each do |thread|
+        backtrace = thread.backtrace || ["<no backtrace available>"]
+        tid = (thread.object_id ^ ::Process.pid).to_s(36)
+        Logger.warn "Thread tid=#{tid} name=#{thread.name} [#{thread.status}]\n#{backtrace.join("\n")}"
+      end
+    end
   end
 end