milvaion-worker 1.0.0 → 1.1.0

data/lib/milvaion/options.rb

@@ -20,7 +20,9 @@ module Milvaion
       startup_timeout: 30,
       version: "1.0.0",
       instance_id: nil,
-      logger: nil
+      logger: nil,
+      transport: "rabbitmq",
+      kafka: nil
     }.freeze
 
     # worker_id:              identifies the worker application. Use the same value on every replica.
@@ -37,7 +39,9 @@ module Milvaion
     # version:                shown in the dashboard.
     # instance_id:            overrides the generated "{worker_id}-{8 hex}" id. Must be unique per process.
     # logger:                 a Logger. Defaults to one that writes to $stdout at INFO.
-    attr_reader :worker_id, :rabbitmq_url, :queue_type, :redis_url, :redis_key_prefix, :max_parallel_jobs,
+    # transport:              "rabbitmq" (default) or "kafka". The API and every worker of a deployment must use the same one.
+    # kafka:                  Kafka settings (a Hash, see Milvaion::Kafka::Settings::KEYS). Only used when transport is "kafka".
+    attr_reader :transport, :kafka, :worker_id, :rabbitmq_url, :queue_type, :redis_url, :redis_key_prefix, :max_parallel_jobs,
                 :execution_timeout, :heartbeat_interval, :job_heartbeat_interval, :shutdown_timeout,
                 :startup_timeout, :version, :instance_id, :logger
 
@@ -67,6 +71,17 @@ module Milvaion
       @version = values[:version]
       @instance_id = values[:instance_id].to_s.empty? ? "#{@worker_id}-#{SecureRandom.hex(4)}" : values[:instance_id]
       @logger = values[:logger] || Logger.new($stdout, progname: "milvaion", level: Logger::INFO)
+
+      @transport = (values[:transport] || "rabbitmq").to_s.strip.downcase
+      @transport = "rabbitmq" if @transport.empty?
+      raise ArgumentError, %(transport must be "rabbitmq" or "kafka", got #{values[:transport].inspect}) unless %w[rabbitmq kafka].include?(@transport)
+
+      # Validated here, so a bad setting fails before anything connects.
+      @kafka = kafka? ? Kafka::Settings.new(values[:kafka], @max_parallel_jobs) : nil
+    end
+
+    def kafka?
+      transport == "kafka"
     end
 
     def cancellation_channel
@@ -77,6 +92,7 @@ module Milvaion
       "MILVAION_WORKER_ID" => [:worker_id, :string],
       "MILVAION_RABBITMQ_URL" => [:rabbitmq_url, :string],
       "MILVAION_RABBITMQ_QUEUE_TYPE" => [:queue_type, :string],
+      "MILVAION_TRANSPORT" => [:transport, :string],
       "MILVAION_REDIS_KEY_PREFIX" => [:redis_key_prefix, :string],
       "MILVAION_MAX_PARALLEL_JOBS" => [:max_parallel_jobs, :integer],
       "MILVAION_EXECUTION_TIMEOUT_SECONDS" => [:execution_timeout, :float],
@@ -85,8 +101,9 @@ module Milvaion
 
     # Reads MILVAION_WORKER_ID, MILVAION_RABBITMQ_URL, MILVAION_RABBITMQ_QUEUE_TYPE, MILVAION_REDIS_URL
     # ("false" disables Redis), MILVAION_REDIS_KEY_PREFIX, MILVAION_MAX_PARALLEL_JOBS,
-    # MILVAION_EXECUTION_TIMEOUT_SECONDS and MILVAION_HEARTBEAT_INTERVAL_SECONDS.
-    # Keyword arguments win over the environment.
+    # MILVAION_EXECUTION_TIMEOUT_SECONDS and MILVAION_HEARTBEAT_INTERVAL_SECONDS, and the transport
+    # from MILVAION_TRANSPORT and the MILVAION_KAFKA_* variables (see Milvaion::Kafka.options_from_env).
+    # Keyword arguments win over the environment; a kafka: Hash wins key by key.
     def self.from_env(env = ENV, **overrides)
       values = {}
 
@@ -112,7 +129,11 @@ module Milvaion
         values[:redis_url] = redis_url
       end
 
-      values.merge!(overrides)
+      kafka = Kafka.options_from_env(env)
+      kafka = kafka.merge(overrides[:kafka].to_h.transform_keys(&:to_sym)) if overrides[:kafka]
+      values[:kafka] = kafka unless kafka.empty?
+
+      values.merge!(overrides.except(:kafka))
       values[:worker_id] ||= ""
 
       new(**values)

data/lib/milvaion/processor.rb

@@ -3,8 +3,18 @@
 require "json"
 
 module Milvaion
-  # One message taken from a job queue. ack is called with no arguments, nack with the requeue flag.
-  Delivery = Struct.new(:body, :headers, :correlation_id, :redelivered, :ack, :nack, keyword_init: true)
+  # One message taken from a job queue or topic. ack is called with no arguments, nack with the
+  # requeue flag: nack(true) returns the message so it runs again, nack(false) dead-letters it.
+  #
+  # reject:         optional; dead-letters a message that cannot be processed, with the reason.
+  #                 Without it nack(false) is used.
+  # schedule_retry: optional; hands a failed job to the broker for a delayed retry instead of
+  #                 waiting in process (Kafka's retry topic). Called with a RetryRequest; returns once
+  #                 the broker has accepted the retry and raises otherwise.
+  Delivery = Struct.new(:body, :headers, :correlation_id, :redelivered, :ack, :nack, :reject, :schedule_retry, keyword_init: true)
+
+  # The next attempt of a failed job. retry_count is that of the next attempt; not_before a Time.
+  RetryRequest = Struct.new(:job_id, :occurrence_id, :retry_count, :max_retries, :not_before, keyword_init: true)
 
   # Abort signals of the jobs running on this instance, keyed by occurrence id.
   class CancellationRegistry
@@ -103,7 +113,7 @@ module Milvaion
         process(delivery, consumer)
       rescue StandardError => e
         logger.error("Unexpected error while processing a job delivery: #{e.class}: #{e.message}")
-        settle { delivery.nack.call(false) }
+        reject(delivery, "unexpected error while processing the job: #{e.class}: #{e.message}")
       ensure
         change_count(-1)
       end
@@ -192,7 +202,7 @@ module Milvaion
 
       if job.nil?
         logger.error("Received a message on #{consumer.queue} that is not a valid job. Sending it to the dead letter queue.")
-        settle { delivery.nack.call(false) }
+        reject(delivery, "the message is not a valid job")
         return
       end
 
@@ -200,7 +210,7 @@ module Milvaion
 
       if occurrence_id.nil?
         logger.error("Job #{job['id']} arrived without an occurrence id. Sending it to the dead letter queue.")
-        settle { delivery.nack.call(false) }
+        reject(delivery, "the message has no occurrence id")
         return
       end
 
@@ -421,8 +431,11 @@ module Milvaion
     end
 
     # Waits out the backoff, republishes the job with a higher retry count and acks the original.
-    # The original stays unacked meanwhile, so a crash during the wait gets it redelivered.
+    # The original stays unacked meanwhile, so a crash during the wait gets it redelivered. A broker
+    # that schedules retries itself (Kafka) gets the retry right away instead.
     def retry_later(delivery, consumer, run, delay, max_retries)
+      return schedule_retry(delivery, run, delay, max_retries) if delivery.schedule_retry
+
       interval = options.job_heartbeat_interval.to_f
       remaining = delay.to_f
 
@@ -444,13 +457,67 @@ module Milvaion
                                                              headers: { "CorrelationId" => run.occurrence_id, "x-retry-count" => run.retry_count + 1, "MaxRetries" => max_retries })
       rescue StandardError => e
         logger.error("Failed to republish job #{run.job['id']} for retry. Dead-lettering the message: #{e.message}")
-        settle { delivery.nack.call(false) }
+        reject(delivery, "the job could not be republished for retry: #{e.message}")
         return
       end
 
       settle { delivery.ack.call }
     end
 
+    # Hands the next attempt to the broker, settles the original and keeps the occurrence's job
+    # heartbeat going until the attempt is due, since the occurrence stays Running.
+    def schedule_retry(delivery, run, delay, max_retries)
+      due = Time.now + delay
+
+      begin
+        delivery.schedule_retry.call(RetryRequest.new(job_id: run.job["id"], occurrence_id: run.occurrence_id, retry_count: run.retry_count + 1,
+                                                      max_retries: max_retries, not_before: due))
+      rescue StandardError => e
+        logger.error("Failed to schedule job #{run.job['id']} for retry. Dead-lettering the message: #{e.message}")
+        reject(delivery, "the job could not be scheduled for retry: #{e.message}")
+        return
+      end
+
+      settle { delivery.ack.call }
+
+      Thread.new do
+        Thread.current.name = "milvaion-retry-heartbeat"
+        heartbeat_while_waiting(run, due)
+      end
+    end
+
+    public
+
+    # Sends Running statuses every job_heartbeat_interval for an occurrence that waits for its next
+    # attempt, until due, until stop (an AbortSignal) is triggered or the worker stops.
+    def heartbeat_while_waiting(run, due, stop = nil)
+      interval = options.job_heartbeat_interval.to_f
+      return unless interval.positive?
+
+      loop do
+        remaining = due - Time.now
+        return if remaining <= 0 || stopping?
+
+        step = [remaining, interval].min
+        return if stop ? stop.wait(step) : @shutdown.wait(step)
+
+        run.status(Status::RUNNING, reliable: false) if due > Time.now && !stopping?
+      end
+    rescue StandardError => e
+      logger.debug("Job heartbeat while waiting for a retry failed: #{e.message}")
+    end
+
+    private
+
+    # Dead-letters a message that cannot be processed.
+    def reject(delivery, reason)
+      if delivery.reject
+        settle { delivery.reject.call(reason) }
+      else
+        settle { delivery.nack.call(false) }
+      end
+    end
+
     def dead_letter(delivery, run, exception, retry_count, max_retries)
       job = run.job
       body = JSON.generate({
@@ -465,11 +532,11 @@ module Milvaion
 
       begin
         publisher.publish(Protocol::EXCHANGE_DEAD_LETTER, Protocol::DEAD_LETTER_ROUTING_KEY, body, persistent: true, correlation_id: run.occurrence_id,
-                                                                                                  headers: { "CorrelationId" => run.occurrence_id, "x-retry-count" => retry_count, "MaxRetries" => max_retries })
+                                                                                                  key: run.occurrence_id, headers: { "CorrelationId" => run.occurrence_id, "x-retry-count" => retry_count, "MaxRetries" => max_retries })
       rescue StandardError => e
         # The queue dead-letters a rejected message itself, just without the exception details.
         logger.error("Failed to publish job #{job['id']} to the dead letter queue. Rejecting the message instead: #{e.message}")
-        settle { delivery.nack.call(false) }
+        reject(delivery, "the dead letter message could not be published: #{e.message}")
         return
       end
 
@@ -542,7 +609,7 @@ module Milvaion
       attempt = 0
 
       begin
-        @p.publisher.publish("", Protocol::QUEUE_STATUS_UPDATES, body, persistent: true)
+        @p.publisher.publish("", Protocol::QUEUE_STATUS_UPDATES, body, persistent: true, key: occurrence_id)
       rescue StandardError => e
         if attempt >= delays.size
           @p.logger.warn("Failed to publish status #{Status.name_of(status)} for occurrence #{occurrence_id}: #{e.message}")

data/lib/milvaion/worker/version.rb

@@ -3,6 +3,6 @@
 module Milvaion
   class Worker
     # Version of this SDK, reported to the dashboard as SdkVersion. A spec keeps it equal to the gemspec.
-    VERSION = "1.0.0"
+    VERSION = "1.1.0"
   end
 end

data/lib/milvaion/worker.rb

@@ -8,6 +8,7 @@ require "time"
 require_relative "worker/version"
 require_relative "protocol"
 require_relative "errors"
+require_relative "kafka"
 require_relative "options"
 require_relative "context"
 require_relative "log_buffer"
@@ -16,6 +17,7 @@ require_relative "thread_pool"
 require_relative "metrics"
 require_relative "job"
 require_relative "cancellation_listener"
+require_relative "kafka_transport"
 
 module Milvaion
   # A Milvaion worker. Register jobs with #job, then call #run (blocking) or #start / #stop.
@@ -130,9 +132,12 @@ module Milvaion
       end
     end
 
-    # Connects, registers this instance and starts consuming. Raises if RabbitMQ or Redis cannot be
+    # Connects, registers this instance and starts consuming. Raises if the broker or Redis cannot be
     # reached within startup_timeout.
     def start
+      # Before anything else, so a missing rdkafka gem fails with a clear message.
+      KafkaTransport.load_rdkafka! if options.kafka?
+
       @state_mutex.synchronize do
         raise Error, "Worker cannot be started from state #{@state}." unless @state == :idle
         raise Error, "Add at least one job before starting the worker." if @jobs.empty?
@@ -142,6 +147,11 @@ module Milvaion
 
       logger.info("Starting worker #{worker_id} (instance #{instance_id}) with #{@jobs.size} job type(s).")
 
+      if options.kafka?
+        @kafka = KafkaTransport.new(options, -> { register })
+        @publisher = @kafka
+      end
+
       @logs = LogBuffer.new(->(batch) { send_logs(batch) }, logger).start
       @processor = Processor.new(options, @jobs, @publisher, @logs, @cancellations)
       @processor.on_count_changed = -> { request_immediate_heartbeat }
@@ -155,7 +165,11 @@ module Milvaion
           logger.warn("Redis is disabled: jobs cannot be cancelled from the dashboard.")
         end
 
-        connect
+        if @kafka
+          @kafka.start(@processor, @pool)
+        else
+          connect
+        end
       rescue Exception # rubocop:disable Lint/RescueException -- clean up on any failure, then re-raise
         set_state(:stopped)
         close
@@ -338,7 +352,7 @@ module Milvaion
     end
 
     def register
-      @publisher.publish("", Protocol::QUEUE_WORKER_REGISTRATION, JSON.generate(build_registration), persistent: true)
+      @publisher.publish("", Protocol::QUEUE_WORKER_REGISTRATION, JSON.generate(build_registration), persistent: true, key: instance_id)
       send_heartbeat(false)
       logger.info("Registered #{worker_id} (instance #{instance_id}).")
     end
@@ -395,13 +409,14 @@ module Milvaion
                              "isStopping" => stopping
                            })
 
-      @publisher.publish("", Protocol::QUEUE_WORKER_HEARTBEAT, body, persistent: false)
+      @publisher.publish("", Protocol::QUEUE_WORKER_HEARTBEAT, body, persistent: false, key: instance_id)
     rescue StandardError => e
       logger.debug("Heartbeat failed; will retry on the next interval: #{e.message}")
     end
 
     def send_logs(batch)
-      @publisher.publish("", Protocol::QUEUE_WORKER_LOGS, JSON.generate(batch), persistent: true)
+      first = batch.is_a?(Hash) ? Array(batch["logs"]).first : nil
+      @publisher.publish("", Protocol::QUEUE_WORKER_LOGS, JSON.generate(batch), persistent: true, key: first.is_a?(Hash) ? first["occurrenceId"] : nil)
     end
 
     def shutdown
@@ -414,6 +429,7 @@ module Milvaion
         nil
       end
       @consumers.clear
+      @kafka&.stop_consuming
 
       # Deliveries waiting for a pool thread have not started: give them straight back.
       @pool.close.each do |delivery|
@@ -430,6 +446,9 @@ module Milvaion
         @processor.wait_for_idle(10)
       end
 
+      # Kafka: commit what has finished and leave the group.
+      @kafka&.drained
+
       @heartbeat_mutex.synchronize do
         @heartbeat_stopped = true
         @heartbeat_changed.broadcast
@@ -453,6 +472,7 @@ module Milvaion
       @listener&.stop
       @logs&.close
       @publisher.detach
+      @kafka&.close
       @pool&.close
       @pool&.join(5)
 
@@ -486,7 +506,8 @@ module Milvaion
       end
     end
 
-    def publish(exchange, routing_key, body, persistent: true, correlation_id: nil, headers: nil)
+    # key is the Kafka record key; RabbitMQ has no use for it.
+    def publish(exchange, routing_key, body, persistent: true, correlation_id: nil, headers: nil, key: nil) # rubocop:disable Lint/UnusedMethodArgument
       @mutex.synchronize do
         channel = @channel
         raise ConnectionError, "Not connected to RabbitMQ." if channel.nil? || !channel.open?

data/milvaion-worker.gemspec

@@ -3,12 +3,12 @@
 Gem::Specification.new do |spec|
   spec.name = "milvaion-worker"
   # Kept equal to Milvaion::Worker::VERSION by a spec; the publish workflow checks it against the tag.
-  spec.version = "1.0.0"
+  spec.version = "1.1.0"
   spec.authors = ["Milvasoft"]
   spec.summary = "Ruby worker SDK for the Milvaion distributed job scheduler."
   spec.description = "Run Milvaion jobs in Ruby: registration, heartbeats, live logs, retries with " \
                      "exponential backoff, dead lettering, timeouts, cancellation and graceful shutdown " \
-                     "over RabbitMQ and Redis."
+                     "over RabbitMQ (or Apache Kafka with the rdkafka gem) and Redis."
   spec.homepage = "https://github.com/Milvasoft/milvaion"
   spec.license = "Apache-2.0"
   spec.required_ruby_version = ">= 3.1"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: milvaion-worker
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.1.0
 platform: ruby
 authors:
 - Milvasoft
@@ -66,7 +66,7 @@ dependencies:
         version: '6'
 description: 'Run Milvaion jobs in Ruby: registration, heartbeats, live logs, retries
   with exponential backoff, dead lettering, timeouts, cancellation and graceful shutdown
-  over RabbitMQ and Redis.'
+  over RabbitMQ (or Apache Kafka with the rdkafka gem) and Redis.'
 email:
 executables: []
 extensions: []
@@ -79,6 +79,8 @@ files:
 - lib/milvaion/context.rb
 - lib/milvaion/errors.rb
 - lib/milvaion/job.rb
+- lib/milvaion/kafka.rb
+- lib/milvaion/kafka_transport.rb
 - lib/milvaion/log_buffer.rb
 - lib/milvaion/metrics.rb
 - lib/milvaion/options.rb