cosmonats 0.1.0

data/lib/cosmo/cli.rb

@@ -0,0 +1,201 @@
+# frozen_string_literal: true
+
+require "yaml"
+require "optparse"
+
+module Cosmo
+  class CLI # rubocop:disable Metrics/ClassLength
+    def self.run
+      instance.run
+    end
+
+    def self.instance
+      @instance ||= new
+    end
+
+    def run
+      flags, command, _options = parse
+      load_config(flags[:config_file])
+      puts self.class.banner
+      require_files(flags[:require])
+      Engine.run(command)
+    end
+
+    private
+
+    def parse
+      flags = {}
+      parser = flags_parser(flags)
+      parser.order!
+
+      options = {}
+      command = ARGV.shift
+      parser = options_parser(command, options)
+      parser&.order!
+
+      [flags, command, options]
+    end
+
+    def load_config(path)
+      raise ConfigNotFoundError, path if path && !File.exist?(path)
+
+      unless path
+        # Try default path
+        default_path = File.expand_path(Config::DEFAULT_PATH)
+        path = default_path if File.exist?(default_path)
+      end
+
+      Config.load(path)
+    end
+
+    def require_files(path)
+      return unless path
+
+      if File.directory?(path)
+        files = Dir[File.expand_path("#{path}/*.rb")]
+        files.each { |f| require f }
+      else
+        require File.expand_path(path)
+      end
+    end
+
+    def flags_parser(flags) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+      OptionParser.new do |o| # rubocop:disable Metrics/BlockLength
+        o.banner = "Usage: cosmo [flags] [command] [options]"
+        o.separator ""
+        o.separator "Command:"
+        o.separator "  jobs      Run jobs"
+        o.separator "  streams   Run streams"
+        o.separator "  actions   Run actions"
+        o.separator ""
+        o.separator "Flags:"
+
+        o.on "-c", "--concurrency INT", Integer, "Threads to use" do |arg|
+          flags[:concurrency] = arg
+        end
+
+        o.on "-r", "--require PATH|DIR", "Location of files to require" do |arg|
+          flags[:require] = arg
+        end
+
+        o.on "-t", "--timeout NUM", Integer, "Shutdown timeout" do |arg|
+          flags[:timeout] = arg
+        end
+
+        o.on "-C", "--config PATH", "Path to config file" do |arg|
+          flags[:config_file] = arg
+        end
+
+        o.on "-S", "--setup", "Load config, create streams and exit" do
+          load_config(flags[:config_file])
+
+          Config[:streams].each do |name, config|
+            Client.instance.stream_info(name)
+          rescue NATS::JetStream::Error::NotFound
+            Client.instance.create_stream(name, config)
+          end
+
+          puts "Cosmo streams were created/updated"
+          exit(0)
+        end
+
+        o.on_tail "-v", "--version", "Print version" do
+          puts "Cosmo #{VERSION}"
+          exit(0)
+        end
+
+        o.on_tail "-h", "--help", "Show help" do
+          puts o
+          exit(0)
+        end
+      end
+    end
+
+    def options_parser(command, options) # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
+      case command
+      when "jobs"
+        OptionParser.new do |o|
+          o.banner = "Usage: cosmo jobs [options]"
+
+          o.on "--stream NAME", "Job's stream" do |arg|
+            options[:stream] = arg
+          end
+
+          o.on "--subject NAME", "Job's subject" do |arg|
+            options[:subject] = arg
+          end
+        end
+      when "streams"
+        OptionParser.new do |o|
+          o.banner = "Usage: cosmo streams [options]"
+
+          o.on "--stream NAME", "Specify stream name" do |arg|
+            options[:stream] = arg
+          end
+
+          o.on "--subject NAME", "Specify subject name" do |arg|
+            options[:subject] = arg
+          end
+
+          o.on "--consumer_name NAME", "Specify consumer name" do |arg|
+            options[:consumer_name] = arg
+          end
+
+          o.on "--batch_size NUM", Integer, "Number of messages in the batch" do |arg|
+            options[:batch_size] = arg
+          end
+        end
+      when "actions"
+        OptionParser.new do |o|
+          o.banner = "Usage: cosmo actions [options]"
+
+          o.on "-n", "--nop", "Do nothing and exit" do
+            exit(0)
+          end
+        end
+      end
+    end
+
+    # rubocop:disable Layout/TrailingWhitespace,Lint/IneffectiveAccessModifier
+    def self.banner
+      <<-TEXT
+                    .#%+:                                                  
+                     ==-.                                       +.         
+                       +:  .::::.                              :*-         
+                     .=%%%%%%%%%%%%%#-                                     
+                  .#%%%%%%%%##*+===+*#%%:                                  
+                :##%%%%#:  :-::...::::. -%.                                
+               +%%%**  :.             :+. -=.%:                         -  
+              *%%%%: .-%%%#             ++ ---%.    -=                :==- 
+              :%%%-  *%%%%               *- %:#+   =%%%.             .====:
+          .%@%+.##.  #%%:                -+ =-=-    *%%:  .            :=  
+          =*%%%=-#.                      :: =-:     *%-%%%%%#              
+          .%=##* #-                         %.     :%%-+++%%%:             
+           +=*+= +%.                      .*.   .=+.%%-#%%#*+:             
+            ===: =*%=                   .*+  *%+%%% -%*:%%%%%:       .     
+                 =***%*.            .:#*:  .%%*+%%%#.+%+. .          =     
+           -%#-:        -*########+   .:   +%%%=%%%*++:                    
+        .##:---.  :%%-  *: +%%%%%%%%%%%##. -#%%*= =-                       
+         *#:-: :%%%%+%%= --%%#. -#%%%%+=#-   =:            .:::::::::::::: 
+         :#- =%%%%%%%+++. +*%=-%%%%#-..       ..:::::::::::::..            
+            +***+:=***::  ==:.      ....::..                               
+         -%%%%%+%%-         ....               .     +##%%=  .#%%%%%%%#.   
+        .%%%%##=%%-               :+#%%%#. *%%%%+   -%%%%%= .%%%%%%%%%%%   
+         -+--         -=#%#+:   +%%%%%%%#. *%%%%%+  #%%%%%= =%%%*   %%%%.  
+           .-+###.  *%%%%%%%%+ -%%%%+.     *%%%%%%:*%%%%%%+ -%%%*   %%%%.  
+        .*%%%%%%%  %%%%+.=%%%# =%%%%-      *%%%%%%%%%%*%%%+ -%%%*   %%%%.  
+       =%%%%%#=:. :%%%#  .%%%#  #%%%%%%%:  +%%%-*%%%%:=%%%+ -%%%*   %%%%.  
+      -%%%%-      .%%%#   %%%%    .#%%%%%= +%%%- #%%= -%%%+ -%%%*  .%%%%.  
+      *%%%#       .%%%#.  %%%%       +%%%* +%%%-      -%%%+  #%%%###%%%*   
+      %%%%*        %%%#.  #%%%..****#####- =###-      -###+   =######*.    
+      #%%%#        #%%%*+*#%%* .########:  =#*=:                           
+      +%%%%-   .=+ .########=   :----.           .:::--====++++********### 
+      .#%#########:  :==-:          ..:--=====---::::..                    
+       .########+.         .:--=--::.                                      
+          :--.       .---:.                                                
+                 :.                                                        
+      TEXT
+    end
+    # rubocop:enable Layout/TrailingWhitespace,Lint/IneffectiveAccessModifier
+  end
+end

data/lib/cosmo/client.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require "nats/client"
+
+module Cosmo
+  class Client
+    def self.instance
+      @instance ||= Client.new
+    end
+
+    attr_reader :nc, :js
+
+    def initialize(nats_url: ENV.fetch("NATS_URL", "nats://localhost:4222"))
+      @nc = NATS.connect(nats_url)
+      @js = @nc.jetstream
+    end
+
+    def publish(subject, payload, **params)
+      js.publish(subject, payload, **params)
+    end
+
+    def subscribe(subject, consumer_name, config)
+      js.pull_subscribe(subject, consumer_name, config: config)
+    end
+
+    def stream_info(name)
+      js.stream_info(name)
+    end
+
+    def create_stream(name, config)
+      js.add_stream(name: name, **config)
+    end
+
+    def delete_stream(name, params = {})
+      js.delete_stream(name, params)
+    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") }
+    end
+
+    def get_message(name, seq)
+      js.get_msg(name, seq: seq)
+    end
+
+    def close
+      nc.close
+    end
+  end
+end

data/lib/cosmo/config.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+require "yaml"
+require "forwardable"
+
+module Cosmo
+  class Config
+    NANO = 1_000_000_000
+    DEFAULT_PATH = "config/cosmo.yml"
+
+    class << self
+      extend Forwardable
+
+      delegate %i[[] fetch dig to_h set load] => :instance
+    end
+
+    def self.parse_file(path)
+      YAML.load_file(path, aliases: true).tap { normalize!(_1) }
+    end
+
+    def self.normalize!(config) # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+      Utils::Hash.symbolize_keys!(config)
+
+      config[:consumers]&.each_key do |name|
+        config[:consumers][name].each do |stream_name, c|
+          next unless c
+
+          c[:subject] = format(c[:subject], { name: stream_name }) if c[:subject]
+          c[:subjects] = c[:subjects].map { |s| format(s, name: stream_name) } if c[:subjects]
+        end
+      end
+
+      config[:streams]&.each_key do |name|
+        c = config[:streams][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]
+      end
+    end
+
+    def self.deliver_policy(start_position)
+      case start_position
+      when "last", :last
+        { deliver_policy: "last" }
+      when "new", :new
+        { deliver_policy: "new" }
+      when Time
+        { deliver_policy: "by_start_time", opt_start_time: start_position.iso8601 }
+      when String
+        { deliver_policy: "by_start_time", opt_start_time: start_position }
+      else
+        { deliver_policy: "all" }
+      end
+    end
+
+    def self.instance
+      @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)
+    end
+
+    def set(...)
+      @config ||= {}
+      Utils::Hash.set(@config, ...)
+    end
+
+    def load(path = nil)
+      return unless path
+
+      @config = self.class.parse_file(path)
+    end
+  end
+end

data/lib/cosmo/defaults.yml

@@ -0,0 +1,69 @@
+timeout: 25
+max_retries: &max_retries 3
+concurrency: &concurrency 1
+
+consumers:
+  jobs:
+    critical:
+      <<: &config
+        ack_policy: explicit      # each individual message must be acknowledged
+        max_deliver: *max_retries # max number of times a message delivery will be attempted
+        max_ack_pending: 3        # maximum number of messages w/o ack
+        ack_wait: 60              # duration server waits for ack of message once it's delivered
+        subject: jobs.%{name}.>
+      priority: 50
+    high:
+      <<: *config
+      priority: 30
+    default:
+      <<: *config
+      priority: 15
+    low:
+      <<: *config
+      priority: 5
+    scheduled:
+      <<: *config
+      max_deliver: 1
+      max_ack_pending: 100
+      ack_wait: 10
+
+streams:
+  critical:
+    <<: &config
+      storage: file
+      retention: workqueue
+      duplicate_window: 120 # 2m
+      discard: old
+      allow_direct: true
+      subjects:
+        - jobs.%{name}.>
+    description: Very critical priority jobs
+  high:
+    <<: *config
+    description: Higher priority jobs
+  default:
+    <<: *config
+    description: Default priority jobs
+  low:
+    <<: *config
+    description: Lower priority jobs
+  scheduled:
+    <<: *config
+    description: Scheduled jobs
+  dead:
+    <<: *config
+    retention: limits
+    max_msgs: 10000
+    max_age: 604800 # 7d
+    description: Broken jobs (DLQ)
+
+development:
+  verbose: false
+  concurrency: 1
+
+staging:
+  verbose: true
+  concurrency: 3
+
+production:
+  concurrency: 3

data/lib/cosmo/engine.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "concurrent-ruby"
+
+module Cosmo
+  class Engine
+    PROCESSORS = {
+      jobs: Job::Processor,
+      streams: Stream::Processor
+    }.freeze
+
+    def self.run(...)
+      instance.run(...)
+    end
+
+    def self.instance
+      @instance ||= new
+    end
+
+    def initialize
+      @concurrency = Config.fetch(:concurrency, 1)
+      @pool = Utils::ThreadPool.new(@concurrency)
+      @running = Concurrent::AtomicBoolean.new
+    end
+
+    def run(type)
+      handler = Utils::Signal.trap(:INT, :TERM)
+      Logger.info "Starting processing, hit Ctrl-C to stop"
+
+      @processors = type && PROCESSORS.key?(type.to_sym) ? [PROCESSORS[type.to_sym]] : PROCESSORS.values
+      @processors = @processors.map { _1.run(@pool, @running) }
+
+      signal = handler.wait
+      Logger.info "Shutting down... (#{signal} received)"
+      shutdown
+    end
+
+    def shutdown
+      @running.make_false
+      @pool.shutdown
+      Logger.info "Pausing to allow jobs to finish..."
+      @pool.wait_for_termination(Config[:timeout])
+      Logger.info "Bye!"
+    end
+  end
+end

data/lib/cosmo/job/data.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Cosmo
+  module Job
+    class Data
+      DEFAULTS = { stream: :default, retry: 3, dead: true }.freeze
+
+      attr_reader :jid
+
+      def initialize(class_name, args, options = nil)
+        @class_name = class_name
+        @args = args
+        @options = Hash(options)
+        validate!
+
+        @at = @options[:at].to_i if @options[:at]
+        @at ||= Time.now.to_i + @options[:in].to_i if @options[:in]
+        @subject = @options[:subject] if @options[:subject]
+
+        @jid = SecureRandom.hex(12)
+      end
+
+      def stream(target: false)
+        return @options[:stream] if target
+
+        @at ? :scheduled : @options[:stream]
+      end
+
+      def subject(target: false)
+        ["jobs", stream(target:).to_s, Utils::String.underscore(@class_name)]
+      end
+
+      def as_json
+        {
+          jid: jid,
+          class: @class_name,
+          args: @args,
+          retry: retries,
+          dead: dead
+        }
+      end
+
+      def to_json(*_args)
+        Utils::Json.dump(as_json)
+      end
+
+      def to_args
+        headers = { "Nats-Msg-Id" => jid }
+        if @at
+          headers.merge!("X-Execute-At" => @at.to_i,
+                         "X-Stream" => stream(target: true),
+                         "X-Subject" => subject(target: true).join("."))
+        end
+        [@subject || subject.join("."), to_json, { stream: stream, header: headers }]
+      end
+
+      private
+
+      def validate!
+        raise ArgumentError, "stream is not provided" unless @options[:stream]
+      end
+
+      def retries
+        @options[:retry].nil? ? DEFAULTS[:retry] : @options[:retry]
+      end
+
+      def dead
+        @options[:dead].nil? ? DEFAULTS[:dead] : @options[:dead]
+      end
+    end
+  end
+end

data/lib/cosmo/job/processor.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+module Cosmo
+  module Job
+    class Processor < ::Cosmo::Processor
+      def initialize(pool, running)
+        super
+        @weights = []
+      end
+
+      private
+
+      def run_loop
+        Thread.new { work_loop }
+        Thread.new { schedule_loop }
+      end
+
+      def setup
+        jobs_config = Config.dig(:consumers, :jobs)
+        jobs_config&.each do |stream_name, config|
+          consumer_name = "consumer-#{stream_name}"
+          subject = config.delete(:subject)
+          priority = config.delete(:priority)
+          @weights += ([stream_name] * priority.to_i) if priority
+          @consumers[stream_name] = client.subscribe(subject, consumer_name, config)
+        end
+      end
+
+      def work_loop
+        while running?
+          @weights.shuffle.each do |stream_name|
+            break unless running?
+
+            begin
+              timeout = ENV.fetch("COSMO_JOBS_FETCH_TIMEOUT", 0.1).to_f
+              @pool.post { fetch_messages(stream_name, batch_size: 1, timeout:) }
+            rescue Concurrent::RejectedExecutionError
+              break # pool doesn't accept new jobs, we are shutting down
+            end
+
+            break unless running?
+          end
+        end
+      end
+
+      def schedule_loop # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
+        while running?
+          break unless running?
+
+          timeout = ENV.fetch("COSMO_JOBS_SCHEDULER_FETCH_TIMEOUT", 5).to_f
+          fetch_messages(:scheduled, batch_size: 100, timeout:) do |messages|
+            now = Time.now.to_i
+            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")
+              headers["Nats-Expected-Stream"] = stream
+              execute_at = execute_at.to_i
+
+              if now >= execute_at
+                client.publish(subject, message.data, headers: headers)
+                message.ack
+              else
+                delay_ns = (execute_at - now) * 1_000_000_000
+                message.nak(delay: delay_ns)
+              end
+            end
+          end
+
+          break unless running?
+        end
+      end
+
+      def process(messages) # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
+        message = messages.first
+        Logger.debug "received messages #{messages.inspect}"
+        data = Utils::Json.parse(message.data)
+        unless data
+          Logger.debug ArgumentError.new("malformed payload")
+          return
+        end
+
+        worker_class = Utils::String.safe_constantize(data[:class])
+        unless worker_class
+          Logger.debug ArgumentError.new("#{data[:class]} class not found")
+          return
+        end
+
+        begin
+          sw = stopwatch
+          Logger.with(jid: data[:jid])
+          Logger.info "start"
+          instance = worker_class.new
+          instance.jid = data[:jid]
+          instance.perform(*data[:args])
+          message.ack
+          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "done" }
+        rescue StandardError => e
+          Logger.debug e
+          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail" }
+          handle_failure(message, data)
+        rescue Exception # rubocop:disable Lint/RescueException
+          Logger.with(elapsed: sw.elapsed_seconds) { Logger.info "fail" }
+          raise
+        end
+      ensure
+        Logger.without(:jid)
+        Logger.debug "processed message #{message.inspect}"
+      end
+
+      def handle_failure(message, data) # rubocop:disable Metrics/AbcSize
+        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
+          delay_ns = ((current_attempt**4) + 15) * 1_000_000_000
+          message.nak(delay: delay_ns)
+          return
+        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]}"
+        end
+      end
+    end
+  end
+end