rocketjob 5.4.1 → 6.0.0

data/lib/rocket_job/sliced/slices.rb

@@ -42,16 +42,10 @@ 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
@@ -96,13 +90,33 @@ module RocketJob
         slice
       end
 
+      def insert_many(slices)
+        documents = slices.collect(&:as_document)
+        all.collection.insert_many(documents)
+      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
       def create_indexes
-        all.collection.indexes.create_one(state: 1, _id: 1) if all.collection.indexes.none? { |i| i["name"] == "state_1__id_1" }
-      rescue Mongo::Error::OperationFailure
-        all.collection.indexes.create_one(state: 1, _id: 1)
+        missing =
+          begin
+            all.collection.indexes.none? { |i| i["name"] == "state_1__id_1" }
+          rescue Mongo::Error::OperationFailure
+            true
+          end
+        all.collection.indexes.create_one({state: 1, _id: 1}, unique: true) if missing
       end
 
       # Forward additional methods.
@@ -139,6 +153,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/input.rb

@@ -12,43 +12,71 @@ module RocketJob
         #     Block to call on the first line only, instead of storing in the slice.
         #     Useful for extracting the header row
         #     Default: nil
-        def self.collect(input, **args)
-          writer = new(input, **args)
+        #
+        #   slice_size: [Integer]
+        #     Override the slice size when uploading for example ranges, where slice is the size
+        #     of the range itself.
+        #
+        #   slice_batch_size: [Integer]
+        #     The number of slices to batch up and to bulk load.
+        #     For smaller slices this significantly improves upload performance.
+        #     Note: If `slice_batch_size` is too high, it can exceed the maximum BSON block size.
+        def self.collect(data_store, **args)
+          writer = new(data_store, **args)
           yield(writer)
           writer.record_count
         ensure
-          writer&.close
+          writer&.flush
         end
 
-        def initialize(input, on_first: nil)
-          @on_first      = on_first
-          @batch_count   = 0
-          @record_count  = 0
-          @input         = input
-          @record_number = 1
-          @slice         = @input.new(first_record_number: @record_number)
+        def initialize(data_store, on_first: nil, slice_size: nil, slice_batch_size: nil)
+          @on_first         = on_first
+          @record_count     = 0
+          @data_store       = data_store
+          @slice_size       = slice_size || @data_store.slice_size
+          @slice_batch_size = slice_batch_size || 20
+          @batch            = []
+          @batch_count      = 0
+          new_slice
         end
 
         def <<(line)
-          @record_number += 1
           if @on_first
             @on_first.call(line)
             @on_first = nil
             return self
           end
           @slice << line
-          @batch_count  += 1
           @record_count += 1
-          if @batch_count >= @input.slice_size
-            @input.insert(@slice)
-            @batch_count = 0
-            @slice       = @input.new(first_record_number: @record_number)
+          if @slice.size >= @slice_size
+            save_slice
+            new_slice
           end
           self
         end
 
-        def close
-          @input.insert(@slice) if @slice.size.positive?
+        def flush
+          if @slice_batch_size
+            @batch << @slice if @slice.size.positive?
+            @data_store.insert_many(@batch)
+            @batch       = []
+            @batch_count = 0
+          elsif @slice.size.positive?
+            @data_store.insert(@slice)
+          end
+        end
+
+        def new_slice
+          @slice = @data_store.new(first_record_number: @record_count + 1)
+        end
+
+        def save_slice
+          return flush unless @slice_batch_size
+
+          @batch_count += 1
+          return flush if @batch_count >= @slice_batch_size
+
+          @batch << @slice
         end
       end
     end

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,33 +56,13 @@ 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?
-        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
+          (categorized_records[named_category] ||= []) << value unless value.nil? && !job.output_category(named_category).nils
         end
       end
     end

data/lib/rocket_job/sliced.rb

@@ -2,6 +2,7 @@ module RocketJob
   module Sliced
     autoload :BZip2OutputSlice, "rocket_job/sliced/bzip2_output_slice"
     autoload :CompressedSlice, "rocket_job/sliced/compressed_slice"
+    autoload :EncryptedBZip2OutputSlice, "rocket_job/sliced/encrypted_bzip2_output_slice"
     autoload :EncryptedSlice, "rocket_job/sliced/encrypted_slice"
     autoload :Input, "rocket_job/sliced/input"
     autoload :Output, "rocket_job/sliced/output"
@@ -13,79 +14,5 @@ module RocketJob
       autoload :Input, "rocket_job/sliced/writer/input"
       autoload :Output, "rocket_job/sliced/writer/output"
     end
-
-    # Returns [RocketJob::Sliced::Slices] for the relevant type and category.
-    #
-    # Supports compress and encrypt with [true|false|Hash] values.
-    # When [Hash] they must specify whether the apply to the input or output collection types.
-    #
-    # Example, compress both input and output collections:
-    #   class MyJob < RocketJob::Job
-    #     include RocketJob::Batch
-    #     self.compress = true
-    #   end
-    #
-    # Example, compress just the output collections:
-    #   class MyJob < RocketJob::Job
-    #     include RocketJob::Batch
-    #     self.compress = {output: true}
-    #   end
-    #
-    # To use the specialized BZip output compressor, and the regular compressor for the input collections:
-    #   class MyJob < RocketJob::Job
-    #     include RocketJob::Batch
-    #     self.compress = {output: :bzip2, input: true}
-    #   end
-    def self.factory(type, category, job)
-      raise(ArgumentError, "Unknown type: #{type.inspect}") unless %i[input output].include?(type)
-
-      collection_name = "rocket_job.#{type}s.#{job.id}"
-      collection_name << ".#{category}" unless category == :main
-
-      args               = {collection_name: collection_name, slice_size: job.slice_size}
-      klass              = slice_class(type, job)
-      args[:slice_class] = klass if klass
-
-      if type == :input
-        RocketJob::Sliced::Input.new(args)
-      else
-        RocketJob::Sliced::Output.new(args)
-      end
-    end
-
-    private
-
-    # Parses the encrypt and compress options to determine which slice serializer to use.
-    # `encrypt` takes priority over any `compress` option.
-    def self.slice_class(type, job)
-      encrypt  = extract_value(type, job.encrypt)
-      compress = extract_value(type, job.compress)
-
-      if encrypt
-        case encrypt
-        when true
-          EncryptedSlice
-        else
-          raise(ArgumentError, "Unknown job `encrypt` value: #{compress}") unless compress.is_a?(Slices)
-          # Returns the supplied class to use for encryption.
-          encrypt
-        end
-      elsif compress
-        case compress
-        when true
-          CompressedSlice
-        when :bzip2
-          BZip2OutputSlice
-        else
-          raise(ArgumentError, "Unknown job `compress` value: #{compress}") unless compress.is_a?(Slices)
-          # Returns the supplied class to use for compression.
-          compress
-        end
-      end
-    end
-
-    def self.extract_value(type, value)
-      value.is_a?(Hash) ? value[type] : value
-    end
   end
 end

data/lib/rocket_job/subscribers/server.rb

@@ -17,7 +17,7 @@ module RocketJob
 
           supervisor.logger.info("Stopping Pool")
           supervisor.worker_pool.stop
-          unless supervisor.worker_pool.living_count == 0
+          unless supervisor.worker_pool.living_count.zero?
             supervisor.logger.info("Giving pool #{wait_timeout} seconds to terminate")
             sleep(wait_timeout)
           end

data/lib/rocket_job/thread_worker.rb

@@ -0,0 +1,46 @@
+require "concurrent"
+module RocketJob
+  # ThreadWorker
+  #
+  # A worker runs on a single operating system thread.
+  # Is usually started under a Rocket Job server process.
+  class ThreadWorker < Worker
+    attr_reader :thread
+
+    def initialize(id:, server_name:)
+      super(id: id, server_name: server_name)
+      @shutdown = Concurrent::Event.new
+      @thread   = Thread.new { 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/throttle_definitions.rb

@@ -1,6 +1,6 @@
 module RocketJob
   class ThrottleDefinitions
-    attr_reader :throttles
+    attr_accessor :throttles
 
     def initialize
       @throttles = []
@@ -35,5 +35,11 @@ module RocketJob
       end
       nil
     end
+
+    def deep_dup
+      new_defination           = dup
+      new_defination.throttles = throttles.map(&:dup)
+      new_defination
+    end
   end
 end

data/lib/rocket_job/version.rb

@@ -1,3 +1,3 @@
 module RocketJob
-  VERSION = "5.4.1".freeze
+  VERSION = "6.0.0".freeze
 end

data/lib/rocket_job/worker.rb

@@ -1,5 +1,3 @@
-require "concurrent"
-require "forwardable"
 module RocketJob
   # Worker
   #
@@ -7,12 +5,9 @@ module RocketJob
   # Is usually started under a Rocket Job server process.
   class Worker
     include SemanticLogger::Loggable
-    include ActiveSupport::Callbacks
-
-    define_callbacks :running
 
     attr_accessor :id, :current_filter
-    attr_reader :thread, :name, :inline, :server_name
+    attr_reader :name, :server_name
 
     # Raised when a worker is killed so that it shutdown immediately, yet cleanly.
     #
@@ -21,59 +16,41 @@ module RocketJob
     class Shutdown < RuntimeError
     end
 
-    def self.before_running(*filters, &blk)
-      set_callback(:running, :before, *filters, &blk)
-    end
-
-    def self.after_running(*filters, &blk)
-      set_callback(:running, :after, *filters, &blk)
-    end
-
-    def self.around_running(*filters, &blk)
-      set_callback(:running, :around, *filters, &blk)
-    end
-
-    def initialize(id: 0, server_name: "inline:0", inline: false)
+    def initialize(id: 0, server_name: "inline:0")
       @id             = id
       @server_name    = server_name
-      @shutdown       = Concurrent::Event.new
       @name           = "#{server_name}:#{id}"
       @re_check_start = Time.now
       @current_filter = Config.filter || {}
-      @thread         = Thread.new { run } unless inline
-      @inline         = inline
     end
 
     def alive?
-      inline ? true : @thread.alive?
+      true
     end
 
     def backtrace
-      inline ? Thread.current.backtrace : @thread.backtrace
+      Thread.current.backtrace
     end
 
-    def join(*args)
-      @thread.join(*args) unless inline
+    def join(*_args)
+      true
     end
 
-    # Send each active worker the RocketJob::ShutdownException so that stops processing immediately.
     def kill
-      return true if inline
-
-      @thread.raise(Shutdown, "Shutdown due to kill request for worker: #{name}") if @thread.alive?
+      true
     end
 
     def shutdown?
-      @shutdown.set?
+      false
     end
 
     def shutdown!
-      @shutdown.set
+      true
     end
 
     # Returns [true|false] whether the shutdown indicator was set
-    def wait_for_shutdown?(timeout = nil)
-      @shutdown.wait(timeout)
+    def wait_for_shutdown?(_timeout = nil)
+      false
     end
 
     # Process jobs until it shuts down
@@ -146,6 +123,8 @@ module RocketJob
 
         # Should this job be throttled?
         next if job.fail_on_exception! { throttled_job?(job) }
+        # Job failed during throttle execution?
+        next if job.failed?
 
         # Start this job!
         job.fail_on_exception! { job.start!(name) }
@@ -171,27 +150,14 @@ module RocketJob
     # Applies the current filter to exclude filtered jobs.
     #
     # Returns nil if no jobs are available for processing.
-    if Mongoid::VERSION.to_f >= 7.1
-      def find_and_assign_job
-        SemanticLogger.silence(:info) do
-          scheduled = RocketJob::Job.where(run_at: nil).or(:run_at.lte => Time.now)
-          working   = RocketJob::Job.queued.or(state: :running, sub_state: :processing)
-          query     = RocketJob::Job.and(working, scheduled)
-          query     = query.and(current_filter) unless current_filter.blank?
-          update    = {"$set" => {"worker_name" => name, "state" => "running"}}
-          query.sort(priority: 1, _id: 1).find_one_and_update(update, bypass_document_validation: true)
-        end
-      end
-    else
-      def find_and_assign_job
-        SemanticLogger.silence(:info) do
-          scheduled = {"$or" => [{run_at: nil}, {:run_at.lte => Time.now}]}
-          working   = {"$or" => [{state: :queued}, {state: :running, sub_state: :processing}]}
-          query     = RocketJob::Job.and(working, scheduled)
-          query     = query.where(current_filter) unless current_filter.blank?
-          update    = {"$set" => {"worker_name" => name, "state" => "running"}}
-          query.sort(priority: 1, _id: 1).find_one_and_update(update, bypass_document_validation: true)
-        end
+    def find_and_assign_job
+      SemanticLogger.silence(:info) do
+        scheduled = RocketJob::Job.where(run_at: nil).or(:run_at.lte => Time.now)
+        working   = RocketJob::Job.queued.or(state: "running", sub_state: "processing")
+        query     = RocketJob::Job.and(working, scheduled)
+        query     = query.and(current_filter) unless current_filter.blank?
+        update    = {"$set" => {"worker_name" => name, "state" => "running"}}
+        query.sort(priority: 1, _id: 1).find_one_and_update(update, bypass_document_validation: true)
       end
     end