cosmonats 0.6.0 → 0.7.0

data/lib/cosmo/heartbeat.rb

@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require "socket"
+require "shellwords"
+
+module Cosmo
+  # Publishes this process's identity and live state to API::Stats::Processes, so the Web UI lists
+  # every running worker, the idle ones included.
+  class Heartbeat
+    INTERVAL = 10
+
+    # @param engine [Engine]
+    # @param options [Hash] the parsed command options
+    def initialize(engine, options: {})
+      @engine = engine
+      @options = options
+      @started_at = Time.now.to_i
+    end
+
+    # Registers the process and keeps refreshing it every INTERVAL seconds.
+    #
+    # @return [self]
+    def start
+      beat
+      @thread = Thread.new do
+        loop do
+          sleep(INTERVAL)
+          beat
+        end
+      end
+      self
+    end
+
+    # Stops beating and removes the entry right away, instead of leaving it to expire.
+    #
+    # @return [void]
+    def stop
+      @thread&.kill&.join
+      @thread = nil
+      processes.unregister(identity)
+    rescue StandardError => e
+      Logger.debug "Heartbeat unregister error: #{e.class} #{e.message}"
+    end
+
+    # Writes the current state now, e.g. right after it changes.
+    #
+    # @return [void]
+    def beat
+      processes.register(identity, info)
+    rescue StandardError => e
+      Logger.debug "Heartbeat error: #{e.class} #{e.message}"
+    end
+
+    # @return [String]
+    def identity
+      @identity ||= "#{Socket.gethostname}-#{::Process.pid}"
+    end
+
+    # @return [Hash] everything the Web UI shows about this process
+    def info
+      static_info.merge(
+        state: @engine.state,
+        busy: @engine.busy,
+        subscriptions: @engine.subscriptions,
+        rss: rss,
+        nats: nats_info,
+        beat_at: Time.now.to_i
+      )
+    end
+
+    private
+
+    def static_info
+      @static_info ||= {
+        identity: identity,
+        hostname: Socket.gethostname,
+        pid: ::Process.pid,
+        cmdline: [File.basename($PROGRAM_NAME.to_s), *CLI.instance.argv].shelljoin,
+        options: @options,
+        concurrency: @engine.concurrency,
+        timeout: Config[:timeout],
+        http_port: Config.dig(:http, :port),
+        version: VERSION,
+        ruby: "#{RUBY_ENGINE} #{RUBY_VERSION}",
+        started_at: @started_at
+      }
+    end
+
+    def nats_info
+      client = Client.instance
+      connection = client.nc
+      server = connection.server_info
+      uri = connection.connected_server
+      {
+        name: client.name,
+        ip: server[:client_ip],
+        client_id: server[:client_id],
+        server: server[:server_name],
+        server_version: server[:version],
+        url: uri && "#{uri.host}:#{uri.port}",
+        rtt: rtt(connection),
+        reconnects: connection.stats[:reconnects]
+      }
+    end
+
+    # nats-pure has no RTT call, but a flush is a PING/PONG round trip.
+    def rtt(connection)
+      started = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+      connection.flush(1)
+      ((::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - started) * 1000).round(2)
+    rescue StandardError
+      nil
+    end
+
+    # Resident memory in KB: from procfs on Linux, from ps elsewhere.
+    def rss
+      status = "/proc/#{::Process.pid}/status"
+      return File.read(status)[/VmRSS:\s+(\d+)/, 1].to_i if File.exist?(status)
+
+      `ps -o rss= -p #{::Process.pid}`.to_i
+    rescue StandardError
+      nil
+    end
+
+    def processes
+      API::Stats::Processes.instance
+    end
+  end
+end

data/lib/cosmo/http_server/handler.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "stringio"
+
+module Cosmo
+  class HTTPServer
+    # Minimal WEBrick servlet adapting requests to a Rack app.
+    class Handler < ::WEBrick::HTTPServlet::AbstractServlet
+      def initialize(server, app)
+        super(server)
+        @app = app
+      end
+
+      def service(req, res)
+        status, headers, body = @app.call(env_for(req))
+        res.status = status
+        write_headers(res, headers)
+        body.each { res.body << _1 }
+      ensure
+        body.close if body.respond_to?(:close)
+      end
+
+      private
+
+      def env_for(req)
+        env = req.meta_vars.compact
+        env.merge!(
+          Rack::SCRIPT_NAME => "",
+          Rack::PATH_INFO => req.path,
+          Rack::QUERY_STRING => req.query_string.to_s,
+          Rack::RACK_INPUT => StringIO.new(req.body.to_s).tap(&:binmode),
+          Rack::RACK_ERRORS => $stderr,
+          Rack::RACK_URL_SCHEME => req.ssl? ? "https" : "http"
+        )
+      end
+
+      def write_headers(res, headers)
+        headers.each do |key, value|
+          key.downcase == "set-cookie" ? Array(value).each { res.cookies << _1 } : res[key] = Array(value).join(", ")
+        end
+      end
+    end
+  end
+end

data/lib/cosmo/http_server/health.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Cosmo
+  class HTTPServer
+    # Rack middleware answering GET /health, every other request goes down the stack.
+    # Responds 200 when the engine is running and NATS is connected, 503 otherwise.
+    class Health
+      PATH = "/health"
+      HEADERS = { "content-type" => "application/json", "cache-control" => "no-store" }.freeze
+
+      def initialize(app, check: nil)
+        @app = app
+        @check = check || method(:default_check)
+      end
+
+      def call(env)
+        return @app.call(env) unless env[Rack::PATH_INFO] == PATH
+        return [405, HEADERS.merge("allow" => "GET, HEAD"), []] unless %w[GET HEAD].include?(env[Rack::REQUEST_METHOD])
+
+        checks = @check.call
+        healthy = checks.values.all?
+        body = Utils::Json.dump({ status: healthy ? "ok" : "unavailable", checks: checks })
+        [healthy ? 200 : 503, HEADERS.dup, [body]]
+      end
+
+      private
+
+      def default_check
+        { engine: Engine.instance.running?, nats: nats_connected? }
+      end
+
+      def nats_connected?
+        Client.instance.nc.connected?
+      rescue StandardError
+        false
+      end
+    end
+  end
+end

data/lib/cosmo/http_server/ping.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Cosmo
+  class HTTPServer
+    # Rack middleware answering GET /ping with 200 and no checks, every other request goes down the stack.
+    # Useful as a liveness probe: it only proves the process is up and serving HTTP.
+    class Ping
+      PATH = "/ping"
+      HEADERS = { "content-type" => "text/plain", "cache-control" => "no-store" }.freeze
+
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        return @app.call(env) unless env[Rack::PATH_INFO] == PATH
+        return [405, HEADERS.merge("allow" => "GET, HEAD"), []] unless %w[GET HEAD].include?(env[Rack::REQUEST_METHOD])
+
+        [200, HEADERS.dup, ["pong"]]
+      end
+    end
+  end
+end

data/lib/cosmo/http_server.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+begin
+  require "rack"
+  require "webrick"
+rescue LoadError => e
+  raise LoadError, "Cosmo::HTTPServer requires `rack` and `webrick` gems, add them to your Gemfile (#{e.message})"
+end
+
+require "json"
+require "cosmo/http_server/handler"
+require "cosmo/http_server/health"
+require "cosmo/http_server/ping"
+
+module Cosmo
+  # WEBrick server running the Rack app in a background thread.
+  class HTTPServer
+    DEFAULT_HOST = "0.0.0.0"
+    NOT_FOUND = [404, { "content-type" => "application/json" }, ['{"error":"not found"}']].freeze
+
+    # Routing is done by middlewares, anything unmatched falls through to 404.
+    def self.app
+      Rack::Builder.app do
+        use Health
+        use Ping
+        run ->(_env) { NOT_FOUND }
+      end
+    end
+
+    def initialize(app = self.class.app, port:, host: DEFAULT_HOST)
+      @app = app
+      @host = host
+      @port = port
+    end
+
+    def start
+      @server = ::WEBrick::HTTPServer.new(
+        BindAddress: @host,
+        Port: @port,
+        Logger: ::WEBrick::Log.new(IO::NULL),
+        AccessLog: []
+      )
+      @server.mount("/", Handler, @app)
+      @thread = Thread.new { @server.start }
+      Logger.info "HTTP server listening on #{@host}:#{port}"
+      self
+    end
+
+    def stop
+      return unless @server
+
+      @server.shutdown
+      @thread&.join
+      @server = @thread = nil
+    end
+
+    # Actual bound port, useful when started with port 0.
+    def port
+      @server ? @server.listeners.first.addr[1] : @port
+    end
+  end
+end

data/lib/cosmo/job/failure.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Job
+    # Describes why a job died, as NATS headers carried alongside the payload in the
+    # dead letter queue. The worker that raised is long gone by the time the dead jobs
+    # page is rendered, so the cause has to travel with the message itself.
+    class Failure
+      HEADER_LIMIT = 1024
+      BACKTRACE_LINES = 10
+      BACKTRACE_SEPARATOR = " | "
+
+      # @param exception [Exception]
+      # @return [Hash{String => String}] headers describing the failure
+      def self.headers(exception)
+        new(exception).headers
+      end
+
+      # @param exception [Exception]
+      def initialize(exception)
+        @exception = exception
+      end
+
+      # @return [Hash{String => String}]
+      def headers
+        headers = { "X-Error-Class" => name, "X-Error-Message" => message }
+        headers["X-Error-Backtrace"] = backtrace unless backtrace.empty?
+        headers
+      end
+
+      private
+
+      def name
+        @exception.class.name.to_s
+      end
+
+      def message
+        value = single_line(@exception.message)
+        value.empty? ? name : value
+      end
+
+      def backtrace
+        @backtrace ||= single_line(Array(@exception.backtrace).first(BACKTRACE_LINES).join(BACKTRACE_SEPARATOR))
+      end
+
+      # NATS header values are a single line, so whitespace collapses and the value is capped.
+      def single_line(value)
+        value.to_s.gsub(/\s+/, " ").strip[0, HEADER_LIMIT].to_s
+      end
+    end
+  end
+end

data/lib/cosmo/job/processor.rb

@@ -5,27 +5,62 @@ 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
 
         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
+        Config.dig(:consumers, :jobs, StreamFilter::SCHEDULED)
       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 +68,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))
@@ -60,16 +100,11 @@ module Cosmo
         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
@@ -115,9 +150,9 @@ 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
@@ -154,7 +189,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
@@ -207,7 +242,7 @@ 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
@@ -217,16 +252,44 @@ module Cosmo
         Logger.debug "job dropped #{data[:jid]}"
       end
 
-      def move_message(message, data = nil)
+      # 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
+        move_message(message, data, error)
+      end
+
+      def move_message(message, data = nil, exception = nil)
         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
@@ -243,8 +306,8 @@ module Cosmo
       end
 
       def with_stats(message, &block)
-        API::Busy.instance.with(message) do
-          API::Counter.instance.with(&block)
+        API::Stats::Busy.instance.with(message) do
+          API::Stats::Totals.instance.with(&block)
         end
       end
 
@@ -253,15 +316,14 @@ module Cosmo
       # @param message [NATS::Msg]
       # @param duration [Float, nil]
       #
-      # rubocop:disable Lint/UnusedMethodArgument
+      # rubocop:disable-next Lint/UnusedMethodArgument
       def perform_job(job_instance, data:, message:, duration: nil)
         if duration
-          Timeout.timeout(duration) { job_instance.perform(*data[:args]) }
+          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
       end
-      # rubocop:enable Lint/UnusedMethodArgument
     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,8 @@
 # frozen_string_literal: true
 
 require "cosmo/job/data"
+require "cosmo/job/failure"
+require "cosmo/job/stream_filter"
 require "cosmo/job/limit"
 require "cosmo/job/processor"