rocketjob 5.4.1 → 6.0.0

data/lib/rocket_job/plugins/state_machine.rb

@@ -36,8 +36,8 @@ module RocketJob
           raise(ArgumentError, "Cannot supply both a method name and a block") if methods.size.positive? && block
           raise(ArgumentError, "Must supply either a method name or a block") unless methods.size.positive? || block
 
-          # TODO: Somehow get AASM to support options such as :if and :unless to be consistent with other callbacks
-          # For example:
+          # Limitation with AASM. It only supports guards on event transitions, not for callbacks.
+          # For example, AASM does not support callback options such as :if and :unless, yet Rails callbacks do.
           #    before_start :my_callback, unless: :encrypted?
           #    before_start :my_callback, if: :encrypted?
           event = aasm.state_machine.events[event_name]
@@ -51,7 +51,8 @@ module RocketJob
               # Validate methods are any of Symbol String Proc
               methods.each do |method|
                 unless method.is_a?(Symbol) || method.is_a?(String)
-                  raise(ArgumentError, "#{action}_#{event_name} currently does not support any options. Only Symbol and String method names can be supplied.")
+                  raise(ArgumentError,
+                        "#{action}_#{event_name} currently does not support any options. Only Symbol and String method names can be supplied.")
                 end
               end
               methods

data/lib/rocket_job/plugins/throttle_dependent_jobs.rb

@@ -0,0 +1,37 @@
+require "active_support/concern"
+module RocketJob
+  module Plugins
+    # Prevent this job from starting, or a batch slice from starting if the dependent jobs are running.
+    #
+    # Features:
+    # - Ensures dependent jobs won't run
+    # When the throttle has been exceeded all jobs of this class will be ignored until the
+    # next refresh. `RocketJob::Config::re_check_seconds` which by default is 60 seconds.
+    module ThrottleDependentJobs
+      extend ActiveSupport::Concern
+
+      included do
+        field :dependent_jobs, type: Array, class_attribute: true, user_editable: true, copy_on_restart: true
+
+        define_throttle :dependent_job_exists?
+        define_batch_throttle :dependent_job_exists? if respond_to?(:define_batch_throttle)
+      end
+
+      private
+
+      # Checks if there are any dependent jobs are running
+      def dependent_job_exists?
+        return false if dependent_jobs.blank?
+
+        jobs_count = RocketJob::Job.running.where(:_type.in => dependent_jobs).count
+        return false if jobs_count.zero?
+
+        logger.info(
+          message: "#{jobs_count} Dependent Jobs are running from #{dependent_jobs.join(', ')}",
+          metric:  "#{self.class.name}/dependent_jobs_throttle"
+        )
+        true
+      end
+    end
+  end
+end

data/lib/rocket_job/ractor_worker.rb

@@ -0,0 +1,42 @@
+module RocketJob
+  # Run each worker in its own "Ractor".
+  class RactorWorker < Worker
+    attr_reader :thread
+
+    def initialize(id:, server_name:)
+      super(id: id, server_name: server_name)
+      @shutdown = Concurrent::Event.new
+      @thread   = Ractor.new(name: "rocketjob-#{id}") { run }
+    end
+
+    def alive?
+      @thread.alive?
+    end
+
+    def backtrace
+      @thread.backtrace
+    end
+
+    def join(*args)
+      @thread.join(*args)
+    end
+
+    # Send each active worker the RocketJob::ShutdownException so that stops processing immediately.
+    def kill
+      @thread.raise(Shutdown, "Shutdown due to kill request for worker: #{name}") if @thread.alive?
+    end
+
+    def shutdown?
+      @shutdown.set?
+    end
+
+    def shutdown!
+      @shutdown.set
+    end
+
+    # Returns [true|false] whether the shutdown indicator was set
+    def wait_for_shutdown?(timeout = nil)
+      @shutdown.wait(timeout)
+    end
+  end
+end

data/lib/rocket_job/server/model.rb

@@ -28,7 +28,7 @@ module RocketJob
 
         # Current state
         #   Internal use only. Do not set this field directly
-        field :state, type: Symbol, default: :starting
+        field :state, type: Mongoid::StringifiedSymbol, default: :starting
 
         index({name: 1}, background: true, unique: true)
 

data/lib/rocket_job/sliced/bzip2_output_slice.rb

@@ -7,36 +7,35 @@ module RocketJob
     # * The `bzip2` linux command line utility supports multiple embedded BZip2 stream,
     #   but some other custom implementations may not. They may only read the first slice and stop.
     # * It is only designed for use on output collections.
-    #
-    # To download the output when using this slice:
-    #
-    #   # Download the binary BZip2 streams into a single file
-    #   IOStreams.path(output_file_name).stream(:none).writer do |io|
-    #     job.download { |slice| io << slice[:binary] }
-    #   end
     class BZip2OutputSlice < ::RocketJob::Sliced::Slice
-      # This is a specialized binary slice for creating binary data from each slice
+      # This is a specialized binary slice for creating BZip2 binary data from each slice
       # that must be downloaded as-is into output files.
-      def self.binary?
-        true
+      def self.binary_format
+        :bz2
+      end
+
+      # Compress the supplied records with BZip2
+      def self.to_binary(records, record_delimiter = "\n")
+        return [] if records.blank?
+
+        lines = Array(records).join(record_delimiter) + record_delimiter
+        s     = StringIO.new
+        IOStreams::Bzip2::Writer.stream(s) { |io| io.write(lines) }
+        s.string
       end
 
       private
 
+      # Returns [Hash] the BZip2 compressed binary data in binary form when reading back from Mongo.
       def parse_records
-        records = attributes.delete("records")
-
         # Convert BSON::Binary to a string
-        @records = [{binary: records.data}]
+        @records = [attributes.delete("records").data]
       end
 
+      # Returns [BSON::Binary] the records compressed using BZip2 into a string.
       def serialize_records
-        return [] if @records.nil? || @records.empty?
-
-        lines = records.to_a.join("\n") + "\n"
-        s = StringIO.new
-        IOStreams::Bzip2::Writer.stream(s) { |io| io.write(lines) }
-        BSON::Binary.new(s.string)
+        # TODO: Make the line terminator configurable
+        BSON::Binary.new(self.class.to_binary(@records))
       end
     end
   end

data/lib/rocket_job/sliced/compressed_slice.rb

@@ -6,13 +6,10 @@ module RocketJob
       private
 
       def parse_records
-        records = attributes.delete("records")
-
         # Convert BSON::Binary to a string
-        binary_str = records.data
-
-        str      = Zlib::Inflate.inflate(binary_str)
-        @records = Hash.from_bson(BSON::ByteBuffer.new(str))["r"]
+        compressed_str   = attributes.delete("records").data
+        decompressed_str = Zlib::Inflate.inflate(compressed_str)
+        @records         = Hash.from_bson(BSON::ByteBuffer.new(decompressed_str))["r"]
       end
 
       def serialize_records

data/lib/rocket_job/sliced/encrypted_bzip2_output_slice.rb

@@ -0,0 +1,49 @@
+module RocketJob
+  module Sliced
+    # This is a specialized output serializer that renders each output slice as a single BZip2 compressed stream.
+    # BZip2 allows multiple output streams to be written into a single BZip2 file.
+    #
+    # Notes:
+    # * The `bzip2` linux command line utility supports multiple embedded BZip2 stream,
+    #   but some other custom implementations may not. They may only read the first slice and stop.
+    # * It is only designed for use on output collections.
+    class EncryptedBZip2OutputSlice < ::RocketJob::Sliced::Slice
+      # This is a specialized binary slice for creating BZip2 binary data from each slice
+      # that must be downloaded as-is into output files.
+      def self.binary_format
+        :bz2
+      end
+
+      private
+
+      # Returns [Hash] the BZip2 compressed binary data in binary form when reading back from Mongo.
+      def parse_records
+        # Convert BSON::Binary to a string
+        encrypted_str = attributes.delete("records").data
+
+        # Decrypt string
+        header = SymmetricEncryption::Header.new
+        header.parse(encrypted_str)
+        # Use the header that is present to decrypt the data, since its version could be different
+        decrypted_str = header.cipher.binary_decrypt(encrypted_str, header: header)
+
+        @records = [decrypted_str]
+      end
+
+      # Returns [BSON::Binary] the records compressed using BZip2 into a string.
+      def serialize_records
+        return [] if @records.nil? || @records.empty?
+
+        # TODO: Make the line terminator configurable
+        lines = records.to_a.join("\n") + "\n"
+        s     = StringIO.new
+        IOStreams::Bzip2::Writer.stream(s) { |io| io.write(lines) }
+
+        # Encrypt to binary without applying an encoding such as Base64
+        # Use a random_iv with each encryption for better security
+        data = SymmetricEncryption.cipher.binary_encrypt(s.string, random_iv: true, compress: false)
+        BSON::Binary.new(data)
+      end
+    end
+  end
+end

data/lib/rocket_job/sliced/encrypted_slice.rb

@@ -6,17 +6,15 @@ module RocketJob
       private
 
       def parse_records
-        records = attributes.delete("records")
-
         # Convert BSON::Binary to a string
-        binary_str = records.data
+        encrypted_str = attributes.delete("records").data
 
         header = SymmetricEncryption::Header.new
-        header.parse(binary_str)
+        header.parse(encrypted_str)
         # Use the header that is present to decrypt the data, since its version could be different
-        str = header.cipher.binary_decrypt(binary_str, header: header)
+        decrypted_str = header.cipher.binary_decrypt(encrypted_str, header: header)
 
-        @records = Hash.from_bson(BSON::ByteBuffer.new(str))["r"]
+        @records = Hash.from_bson(BSON::ByteBuffer.new(decrypted_str))["r"]
       end
 
       def serialize_records

data/lib/rocket_job/sliced/input.rb

@@ -1,16 +1,16 @@
 module RocketJob
   module Sliced
     class Input < Slices
-      def upload(on_first: nil, &block)
+      def upload(**args, &block)
         # Create indexes before uploading
         create_indexes
-        Writer::Input.collect(self, on_first: on_first, &block)
+        Writer::Input.collect(self, **args, &block)
       rescue Exception => e
         drop
         raise(e)
       end
 
-      def upload_mongo_query(criteria, *column_names, &block)
+      def upload_mongo_query(criteria, columns: [], slice_batch_size: nil, &block)
         options = criteria.options
 
         # Without a block extract the fields from the supplied criteria
@@ -18,23 +18,21 @@ module RocketJob
           # Criteria is returning old school :fields instead of :projections
           options[:projection] = options.delete(:fields) if options.key?(:fields)
         else
-          column_names = column_names.collect(&:to_s)
-          column_names << "_id" if column_names.size.zero?
-
-          fields = options.delete(:fields) || {}
-          column_names.each { |col| fields[col] = 1 }
+          columns = columns.blank? ? ["_id"] : columns.collect(&:to_s)
+          fields  = options.delete(:fields) || {}
+          columns.each { |col| fields[col] = 1 }
           options[:projection] = fields
 
           block =
-            if column_names.size == 1
-              column = column_names.first
+            if columns.size == 1
+              column = columns.first
               ->(document) { document[column] }
             else
-              ->(document) { column_names.collect { |c| document[c] } }
+              ->(document) { columns.collect { |c| document[c] } }
             end
         end
 
-        upload do |records|
+        upload(slice_batch_size: slice_batch_size) do |records|
           # Drop down to the mongo driver level to avoid constructing a Model for each document returned
           criteria.klass.collection.find(criteria.selector, options).each do |document|
             records << block.call(document)
@@ -42,58 +40,48 @@ module RocketJob
         end
       end
 
-      def upload_arel(arel, *column_names, &block)
+      def upload_arel(arel, columns: nil, slice_batch_size: nil, &block)
         unless block
-          column_names = column_names.empty? ? [:id] : column_names.collect(&:to_sym)
+          columns = columns.blank? ? [:id] : columns.collect(&:to_sym)
 
           block =
-            if column_names.size == 1
-              column = column_names.first
-              ->(model) { model.send(column) }
+            if columns.size == 1
+              column = columns.first
+              ->(model) { model.public_send(column) }
             else
-              ->(model) { column_names.collect { |c| model.send(c) } }
+              ->(model) { columns.collect { |c| model.public_send(c) } }
             end
           # find_each requires the :id column in the query
-          selection = column_names.include?(:id) ? column_names : column_names + [:id]
+          selection = columns.include?(:id) ? columns : columns + [:id]
           arel      = arel.select(selection)
         end
 
-        upload { |records| arel.find_each { |model| records << block.call(model) } }
+        upload(slice_batch_size: slice_batch_size) { |records| arel.find_each { |model| records << block.call(model) } }
       end
 
-      def upload_integer_range(start_id, last_id)
-        # Create indexes before uploading
-        create_indexes
-        count = 0
-        while start_id <= last_id
-          end_id = start_id + slice_size - 1
-          end_id = last_id if end_id > last_id
-          create!(records: [[start_id, end_id]])
-          start_id += slice_size
-          count    += 1
+      def upload_integer_range(start_id, last_id, slice_batch_size: 1_000)
+        # Each "record" is actually a range of Integers which makes up each slice
+        upload(slice_size: 1, slice_batch_size: slice_batch_size) do |records|
+          while start_id <= last_id
+            end_id = start_id + slice_size - 1
+            end_id = last_id if end_id > last_id
+            records << [start_id, end_id]
+            start_id += slice_size
+          end
         end
-        count
-      rescue Exception => e
-        drop
-        raise(e)
       end
 
-      def upload_integer_range_in_reverse_order(start_id, last_id)
-        # Create indexes before uploading
-        create_indexes
-        end_id = last_id
-        count  = 0
-        while end_id >= start_id
-          first_id = end_id - slice_size + 1
-          first_id = start_id if first_id.negative? || (first_id < start_id)
-          create!(records: [[first_id, end_id]])
-          end_id -= slice_size
-          count  += 1
+      def upload_integer_range_in_reverse_order(start_id, last_id, slice_batch_size: 1_000)
+        # Each "record" is actually a range of Integers which makes up each slice
+        upload(slice_size: 1, slice_batch_size: slice_batch_size) do |records|
+          end_id = last_id
+          while end_id >= start_id
+            first_id = end_id - slice_size + 1
+            first_id = start_id if first_id.negative? || (first_id < start_id)
+            records << [first_id, end_id]
+            end_id -= slice_size
+          end
         end
-        count
-      rescue Exception => e
-        drop
-        raise(e)
       end
 
       # Iterate over each failed record, if any
@@ -137,11 +125,11 @@ module RocketJob
         # TODO: Will it perform faster without the id sort?
         # I.e. Just process on a FIFO basis?
         document                 = all.queued.
-                                   sort("_id" => 1).
-                                   find_one_and_update(
-                                     {"$set" => {worker_name: worker_name, state: :running, started_at: Time.now}},
-                                     return_document: :after
-                                   )
+          sort("_id" => 1).
+          find_one_and_update(
+            {"$set" => {worker_name: worker_name, state: "running", started_at: Time.now}},
+            return_document: :after
+          )
         document.collection_name = collection_name if document
         document
       end

data/lib/rocket_job/sliced/slice.rb

@@ -33,7 +33,7 @@ module RocketJob
       #
 
       # Current state, as set by AASM
-      field :state, type: Symbol, default: :queued
+      field :state, type: Mongoid::StringifiedSymbol, default: :queued
 
       # When processing started on this slice
       field :started_at, type: Time
@@ -95,9 +95,13 @@ module RocketJob
       end
 
       # Returns whether this is a specialized binary slice for creating binary data from each slice
-      # that is then just downloaded as-is into output files.
-      def self.binary?
-        false
+      # that is downloaded without conversion into output files.
+      def self.binary_format
+      end
+
+      # For binary formats only, format the supplied records into the binary format for this slice
+      def self.to_binary(_records)
+        raise NotImplementedError
       end
 
       # `records` array has special handling so that it can be modified in place instead of having
@@ -139,18 +143,10 @@ module RocketJob
 
       # Returns [Hash] the slice as a Hash for storage purposes
       # Compresses / Encrypts the slice according to the job setting
-      if ::Mongoid::VERSION.to_i >= 6
-        def as_attributes
-          attrs            = super
-          attrs["records"] = serialize_records if @records
-          attrs
-        end
-      else
-        def as_document
-          attrs            = super
-          attrs["records"] = serialize_records if @records
-          attrs
-        end
+      def as_attributes
+        attrs            = super
+        attrs["records"] = serialize_records if @records
+        attrs
       end
 
       def inspect