rocketjob 5.3.3 → 6.0.0.rc1

data/lib/rocket_job/plugins/job/throttle_running_jobs.rb

@@ -29,20 +29,28 @@ module RocketJob
           class_attribute :throttle_running_jobs
           self.throttle_running_jobs = nil
 
+          # Allow jobs to be throttled by group name instance of the job class name.
+          field :throttle_group, type: String, class_attribute: true, user_editable: true, copy_on_restart: true
+
           define_throttle :throttle_running_jobs_exceeded?
         end
 
         private
 
-        # Returns [Boolean] whether the throttle for this job has been exceeded
+        # Returns [true|false] whether the throttle for this job has been exceeded
         def throttle_running_jobs_exceeded?
-          return unless throttle_running_jobs&.positive?
+          return false unless throttle_running_jobs&.positive?
 
-          # Cannot use this class since it will include instances of parent job classes.
           RocketJob::Job.with(read: {mode: :primary}) do |conn|
-            conn.running.where("_type" => self.class.name, :id.ne => id).count >= throttle_running_jobs
+            query = throttle_running_jobs_base_query
+            throttle_group ? query["throttle_group"] = throttle_group : query["_type"] = self.class.name
+            conn.running.where(query).count >= throttle_running_jobs
           end
         end
+
+        def throttle_running_jobs_base_query
+          {:id.ne => id}
+        end
       end
     end
   end

data/lib/rocket_job/plugins/job/worker.rb

@@ -48,11 +48,11 @@ module RocketJob
         def perform_now
           raise(::Mongoid::Errors::Validations, self) unless valid?
 
-          worker = RocketJob::Worker.new(inline: true)
+          worker = RocketJob::Worker.new
           start if may_start?
           # Re-Raise exceptions
           rocket_job_work(worker, true) if running?
-          result
+          @rocket_job_output
         end
 
         def perform(*)
@@ -106,11 +106,6 @@ module RocketJob
               end
             end
 
-            if collect_output?
-              # Result must be a Hash, if not put it in a Hash
-              self.result = @rocket_job_output.is_a?(Hash) ? @rocket_job_output : {"result" => @rocket_job_output}
-            end
-
             if new_record? || destroyed?
               complete if may_complete?
             else

data/lib/rocket_job/plugins/restart.rb

@@ -91,8 +91,16 @@ module RocketJob
           logger.info("Job has expired. Not creating a new instance.")
           return
         end
-        attributes = rocket_job_restart_attributes.each_with_object({}) { |attr, attrs| attrs[attr] = send(attr) }
-        rocket_job_restart_create(attributes)
+        job_attrs              =
+          rocket_job_restart_attributes.each_with_object({}) { |attr, attrs| attrs[attr] = send(attr) }
+        job                    = self.class.new(job_attrs)
+
+        # Copy across input and output categories to new scheduled job so that all of the
+        # settings are remembered between instance. Example: slice_size
+        job.input_categories  = input_categories if respond_to?(:input_categories)
+        job.output_categories = output_categories if respond_to?(:output_categories)
+
+        rocket_job_restart_save(job)
       end
 
       def rocket_job_restart_abort
@@ -101,11 +109,10 @@ module RocketJob
 
       # Allow Singleton to prevent the creation of a new job if one is already running
       # Retry since the delete may not have persisted to disk yet.
-      def rocket_job_restart_create(attrs, retry_limit = 3, sleep_interval = 0.1)
+      def rocket_job_restart_save(job, retry_limit = 10, sleep_interval = 0.5)
         count = 0
         while count < retry_limit
-          job = self.class.create(attrs)
-          if job.persisted?
+          if job.save
             logger.info("Created a new job instance: #{job.id}")
             return true
           else

data/lib/rocket_job/plugins/state_machine.rb

@@ -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/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.rb

@@ -0,0 +1,36 @@
+module RocketJob
+  module Sliced
+    autoload :BZip2OutputSlice, "rocket_job/sliced/bzip2_output_slice"
+    autoload :CompressedSlice, "rocket_job/sliced/compressed_slice"
+    autoload :EncryptedSlice, "rocket_job/sliced/encrypted_slice"
+    autoload :Input, "rocket_job/sliced/input"
+    autoload :Output, "rocket_job/sliced/output"
+    autoload :Slice, "rocket_job/sliced/slice"
+    autoload :Slices, "rocket_job/sliced/slices"
+    autoload :Store, "rocket_job/sliced/store"
+
+    module Writer
+      autoload :Input, "rocket_job/sliced/writer/input"
+      autoload :Output, "rocket_job/sliced/writer/output"
+    end
+
+    # Returns [RocketJob::Sliced::Slices] for the relevant direction and category.
+    def self.factory(direction, category, job)
+      collection_name = "rocket_job.#{direction}s.#{job.id}"
+      collection_name << ".#{category.name}" unless category.name == :main
+
+      case direction
+      when :input
+        RocketJob::Sliced::Input.new(
+          collection_name: collection_name,
+          slice_class:     category.serializer_class,
+          slice_size:      category.slice_size
+        )
+      when :output
+        RocketJob::Sliced::Output.new(collection_name: collection_name, slice_class: category.serializer_class)
+      else
+        raise(ArgumentError, "Unknown direction: #{direction.inspect}")
+      end
+    end
+  end
+end

data/lib/rocket_job/sliced/bzip2_output_slice.rb

@@ -0,0 +1,43 @@
+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.
+    #
+    # 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
+      # that must be downloaded as-is into output files.
+      def self.binary?
+        true
+      end
+
+      private
+
+      def parse_records
+        records = attributes.delete("records")
+
+        # Convert BSON::Binary to a string
+        @records = [{binary: records.data}]
+      end
+
+      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)
+      end
+    end
+  end
+end

data/lib/rocket_job/sliced/input.rb

@@ -5,7 +5,7 @@ module RocketJob
         # Create indexes before uploading
         create_indexes
         Writer::Input.collect(self, on_first: on_first, &block)
-      rescue StandardError => e
+      rescue Exception => e
         drop
         raise(e)
       end
@@ -73,7 +73,7 @@ module RocketJob
           count    += 1
         end
         count
-      rescue StandardError => e
+      rescue Exception => e
         drop
         raise(e)
       end
@@ -91,7 +91,7 @@ module RocketJob
           count  += 1
         end
         count
-      rescue StandardError => e
+      rescue Exception => e
         drop
         raise(e)
       end
@@ -139,7 +139,7 @@ module RocketJob
         document                 = all.queued.
                                    sort("_id" => 1).
                                    find_one_and_update(
-                                     {"$set" => {worker_name: worker_name, state: :running, started_at: Time.now}},
+                                     {"$set" => {worker_name: worker_name, state: "running", started_at: Time.now}},
                                      return_document: :after
                                    )
         document.collection_name = collection_name if document

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
@@ -94,6 +94,12 @@ module RocketJob
         end
       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
+      end
+
       # `records` array has special handling so that it can be modified in place instead of having
       # to replace the entire array every time. For example, when appending lines with `<<`.
       def records
@@ -133,18 +139,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

data/lib/rocket_job/sliced/slices.rb

@@ -42,10 +42,16 @@ module RocketJob
         slice
       end
 
+      # Returns whether this collection contains specialized binary slices for creating binary data from each slice
+      # that is then just downloaded as-is into output files.
+      def binary?
+        slice_class.binary?
+      end
+
       # Returns output slices in the order of their id
       # which is usually the order in which they were written.
-      def each
-        all.sort(id: 1).each { |document| yield(document) }
+      def each(&block)
+        all.sort(id: 1).each(&block)
       end
 
       # Insert a new slice into the collection
@@ -90,6 +96,17 @@ module RocketJob
         slice
       end
 
+      # Append to an existing slice if already present
+      def append(slice, input_slice)
+        existing_slice = all.where(id: input_slice.id).first
+        return insert(slice, input_slice) unless existing_slice
+
+        extra_records          = slice.is_a?(Slice) ? slice.records : slice
+        existing_slice.records = existing_slice.records + extra_records
+        existing_slice.save!
+        existing_slice
+      end
+
       alias << insert
 
       # Index for find_and_modify only if it is not already present
@@ -133,6 +150,7 @@ module RocketJob
       def last
         all.sort("_id" => -1).first
       end
+
       # rubocop:enable Style/RedundantSort
 
       # Returns [Array<Struct>] grouped exceptions by class name,

data/lib/rocket_job/sliced/writer/output.rb

@@ -1,30 +1,37 @@
 module RocketJob
   module Sliced
     module Writer
-      # Internal class for writing categorized results into output slices
-      class Output
+      class Null
         attr_reader :job, :categorized_records
-        attr_accessor :input_slice
-
-        # Collect output results and write to output collections
-        # iff job is collecting output
-        # Notes:
-        #   Nothing is saved if an exception is raised inside the block
-        def self.collect(job, input_slice = nil)
-          if job.collect_output?
-            writer = new(job, input_slice)
-            yield(writer)
-            writer.close
-          else
-            writer = NullWriter.new(job, input_slice)
-            yield(writer)
-          end
-        end
+        attr_accessor :input_slice, :append
 
-        def initialize(job, input_slice = nil)
+        def initialize(job, input_slice: nil, append: false)
           @job                 = job
           @input_slice         = input_slice
           @categorized_records = {}
+          @append              = append
+        end
+
+        def <<(_)
+          # noop
+        end
+
+        def close
+          # noop
+        end
+      end
+
+      # Internal class for writing categorized results into output slices
+      class Output < Null
+        # Collect output results and write to output collections
+        # iff job is collecting output
+        # Notes:
+        #   Partial slices are saved when an exception is raised inside the block
+        def self.collect(job, **args)
+          writer = job.output_categories.present? ? new(job, **args) : Null.new(job, **args)
+          yield(writer)
+        ensure
+          writer&.close
         end
 
         # Writes the supplied result, RocketJob::Batch::Result or RocketJob::Batch::Results
@@ -40,7 +47,8 @@ module RocketJob
         # Write categorized results to their relevant collections
         def close
           categorized_records.each_pair do |category, results|
-            job.output(category).insert(results, input_slice)
+            collection = job.output(category)
+            append ? collection.append(results, input_slice) : collection.insert(results, input_slice)
           end
         end
 
@@ -48,35 +56,16 @@ module RocketJob
 
         # Stores the categorized result from one result
         def extract_categorized_result(result)
-          category = :main
-          value    = result
+          named_category = :main
+          value          = result
           if result.is_a?(RocketJob::Batch::Result)
-            category = result.category
-            value    = result.value
-            raise(ArgumentError, "Invalid RocketJob Output Category: #{category}") if job.output_categories.exclude?(category)
+            named_category = result.category
+            value          = result.value
           end
-          (categorized_records[category] ||= []) << value unless value.nil? && !job.collect_nil_output?
+          (categorized_records[named_category] ||= []) << value unless value.nil? && !job.output_category(named_category).nils
         end
       end
 
-      class NullWriter
-        attr_reader :job, :categorized_records
-        attr_accessor :input_slice
-
-        def initialize(job, input_slice = nil)
-          @job                 = job
-          @input_slice         = input_slice
-          @categorized_records = {}
-        end
-
-        def <<(_)
-          # noop
-        end
-
-        def close
-          # noop
-        end
-      end
     end
   end
 end