rocketjob 3.5.2 → 4.0.0

data/lib/rocket_job/batch/tabular.rb

@@ -0,0 +1,56 @@
+module RocketJob
+  module Batch
+    # Format output results.
+    #
+    # Takes Batch::Results, Batch::Result, Hash, Array, or String and renders it for output.
+    #
+    # Example:
+    #
+    # tabular = Tabular.new(
+    #   main:       IOStreams::Tabular.new(columns: main_file_headers, format: tabular_output_format),
+    #   exceptions: IOStreams::Tabular.new(columns: exception_file_headers, format: tabular_output_format)
+    # )
+    #
+    # tabular.render(row)
+    class Tabular
+      autoload :Input, 'rocket_job/batch/tabular/input'
+      autoload :Output, 'rocket_job/batch/tabular/output'
+
+      def initialize(map)
+        @map = map
+      end
+
+      def [](category = :main)
+        @map[category] || raise("No tabular map defined for category: #{category.inspect}")
+      end
+
+      # Iterate over responses and format using Tabular
+      def render(row, category = :main)
+        if row.is_a?(Batch::Results)
+          results = Batch::Results.new
+          row.each { |result| results << render(result) }
+          results
+        elsif row.is_a?(Batch::Result)
+          row.value = self[row.category].render(row.value)
+          row
+        elsif row.blank?
+          nil
+        else
+          self[category].render(row)
+        end
+      end
+
+      def render_header(category = :main)
+        self[category].render_header
+      end
+
+      def requires_header?(category = :main)
+        self[category].requires_header?
+      end
+
+      def header?(category = :main)
+        self[category].header?
+      end
+    end
+  end
+end

data/lib/rocket_job/batch/tabular/input.rb

@@ -0,0 +1,123 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    class Tabular
+      # For the simple case where all `input_categories` have the same format,
+      # If multiple input categories are used with different formats, then use IOStreams::Tabular directly
+      # instead of this plugin.
+      module Input
+        extend ActiveSupport::Concern
+
+        included do
+          field :tabular_input_header, type: Array, class_attribute: true, user_editable: true
+          field :tabular_input_format, type: Symbol, default: :csv, class_attribute: true, user_editable: true
+
+          # tabular_input_mode: [:line | :row | :record]
+          #   :line
+          #     Uploads the file a line (String) at a time for processing by workers.
+          #   :row
+          #     Parses each line from the file as an Array and uploads each array for processing by workers.
+          #   :record
+          #     Parses each line from the file into a Hash and uploads each hash for processing by workers.
+          #   See IOStream#each_line, IOStream#each_row, and IOStream#each_record.
+          field :tabular_input_mode, type: Symbol, default: :line, class_attribute: true, user_editable: true, copy_on_restart: true
+
+          validates_inclusion_of :tabular_input_format, in: IOStreams::Tabular.registered_formats
+          validates_inclusion_of :tabular_input_mode, in: %i[line row record]
+          validate :tabular_input_header_present
+
+          class_attribute :tabular_input_white_list
+          class_attribute :tabular_input_required
+          class_attribute :tabular_input_skip_unknown
+
+          # Cleanse all uploaded data by removing non-printable characters
+          # and any characters that cannot be converted to UTF-8
+          class_attribute :tabular_input_type
+
+          self.tabular_input_white_list   = nil
+          self.tabular_input_required     = nil
+          self.tabular_input_skip_unknown = true
+          self.tabular_input_type         = :text
+
+          before_perform :tabular_input_render
+        end
+
+        # Extract the header line during the upload.
+        #
+        # Overrides: RocketJob::Batch::IO#upload
+        #
+        # Notes:
+        # - When supplying a block the header must be set manually
+        def upload(file_name_or_io = nil, **args, &block)
+          if tabular_input_type == :text
+            args[:encoding]       = 'UTF-8'
+            args[:encode_cleaner] = :printable
+            args[:encode_replace] = ''
+          end
+
+          # If an input header is not required, then we don't extract it'
+          return super(file_name_or_io, stream_mode: tabular_input_mode, **args, &block) unless tabular_input.header?
+
+          # If the header is already set then it is not expected in the file
+          if tabular_input_header.present?
+            tabular_input_cleanse_header
+            return super(file_name_or_io, stream_mode: tabular_input_mode, **args, &block)
+          end
+
+          case tabular_input_mode
+          when :line
+            parse_header = -> (line) do
+              tabular_input.parse_header(line)
+              tabular_input_cleanse_header
+              self.tabular_input_header = tabular_input.header.columns
+            end
+            super(file_name_or_io, on_first: parse_header, stream_mode: tabular_input_mode, **args, &block)
+          when :row
+            set_header = -> (row) do
+              tabular_input.header.columns = row
+              tabular_input_cleanse_header
+              self.tabular_input_header = tabular_input.header.columns
+            end
+            super(file_name_or_io, on_first: set_header, stream_mode: tabular_input_mode, **args, &block)
+          when :record
+            super(file_name_or_io, stream_mode: tabular_input_mode, **args, &block)
+          else
+            raise(ArgumentError, "Invalid tabular_input_mode: #{stream_mode.inspect}")
+          end
+        end
+
+        private
+
+        # Shared instance used for this slice, by a single worker (thread)
+        def tabular_input
+          @tabular_input ||= IOStreams::Tabular.new(
+            columns:          tabular_input_header,
+            allowed_columns:  tabular_input_white_list,
+            required_columns: tabular_input_required,
+            skip_unknown:     tabular_input_skip_unknown,
+            format:           tabular_input_format
+          )
+        end
+
+        def tabular_input_render
+          @rocket_job_input = tabular_input.record_parse(@rocket_job_input) unless tabular_input_header.blank? && tabular_input.header?
+        end
+
+        # Cleanse custom input header if supplied.
+        def tabular_input_cleanse_header
+          ignored_columns = tabular_input.header.cleanse!
+          logger.warn('Stripped out invalid columns from custom header', ignored_columns) unless ignored_columns.empty?
+
+          self.tabular_input_header = tabular_input.header.columns
+        end
+
+        def tabular_input_header_present
+          return if tabular_input_header.present? || !tabular_input.header? || (tabular_input_mode == :record)
+
+          errors.add(:tabular_input_header, "is required when tabular_input_format is #{tabular_input_format.inspect}")
+        end
+      end
+    end
+  end
+end

data/lib/rocket_job/batch/tabular/output.rb

@@ -0,0 +1,59 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    class Tabular
+      # For the simple case where all `output_categories` have the same format,
+      # If multiple output categories are used with different formats, then use IOStreams::Tabular directly
+      # instead of this plugin.
+      module Output
+        extend ActiveSupport::Concern
+
+        included do
+          field :tabular_output_header, type: Array, class_attribute: true, user_editable: true, copy_on_restart: true
+          field :tabular_output_format, type: Symbol, default: :csv, class_attribute: true, user_editable: true, copy_on_restart: true
+
+          validates_inclusion_of :tabular_output_format, in: IOStreams::Tabular.registered_formats
+
+          after_perform :tabular_output_render
+        end
+
+        # Clear out cached tabular_output any time header or format is changed.
+        def tabular_output_header=(tabular_output_header)
+          super(tabular_output_header)
+          @tabular_output = nil
+        end
+
+        def tabular_output_format=(tabular_output_format)
+          super(tabular_output_format)
+          @tabular_output = nil
+        end
+
+        # Overrides: `RocketJob::Batch::IO#download` to add the `tabular_output_header`.
+        def download(file_name_or_io = nil, category: :main, **args, &block)
+          # No header required
+          return super(file_name_or_io, category: category, **args, &block) unless tabular_output.requires_header?(category)
+
+          header = tabular_output.render_header(category)
+          super(file_name_or_io, header_line: header, category: category, **args, &block)
+        end
+
+        private
+
+        # Delimited instance used for this slice, by a single worker (thread)
+        def tabular_output
+          @tabular_output ||= Tabular.new(
+            main: IOStreams::Tabular.new(columns: tabular_output_header, format: tabular_output_format)
+          )
+        end
+
+        # Render the output from the perform.
+        def tabular_output_render
+          return unless collect_output?
+
+          @rocket_job_output = tabular_output.render(@rocket_job_output)
+        end
+      end
+    end
+  end
+end

data/lib/rocket_job/batch/throttle.rb

@@ -0,0 +1,91 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    # Rocket Job Batch Throttling Framework.
+    #
+    # Example:
+    #   # Do not run any slices for this job when the MySQL slave delay exceeds 5 minutes.
+    #   class MyJob < RocketJob
+    #     include RocketJob::Batch
+    #
+    #     # Define a custom mysql throttle
+    #     # Prevents all slices from this job from running on the current server.
+    #     define_batch_throttle :mysql_throttle_exceeded?
+    #
+    #     def perform(record)
+    #       # ....
+    #     end
+    #
+    #     private
+    #
+    #     # Returns true if the MySQL slave delay exceeds 5 minutes
+    #     def mysql_throttle_exceeded?
+    #       status        = ActiveRecord::Base.connection.connection.select_one('show slave status')
+    #       seconds_delay = Hash(status)['Seconds_Behind_Master'].to_i
+    #       seconds_delay >= 300
+    #     end
+    #   end
+    module Throttle
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :rocket_job_batch_throttles
+        self.rocket_job_batch_throttles = []
+      end
+
+      module ClassMethods
+        # Add a new throttle.
+        #
+        # Parameters:
+        #   method_name: [Symbol]
+        #     Name of method to call to evaluate whether a throttle has been exceeded.
+        #     Note: Must return true or false.
+        #   filter: [Symbol|Proc]
+        #     Name of method to call to return the filter when the throttle has been exceeded.
+        #     Or, a block that will return the filter.
+        #     Default: :throttle_filter_class (Throttle all jobs of this class)
+        #
+        # Note: Throttles are executed in the order they are defined.
+        def define_batch_throttle(method_name, filter: :throttle_filter_class)
+          unless filter.is_a?(Symbol) || filter.is_a?(Proc)
+            raise(ArgumentError, "Filter for #{method_name} must be a Symbol or Proc")
+          end
+          if batch_throttle?(method_name)
+            raise(ArgumentError, "Cannot define #{method_name} twice, undefine previous throttle first")
+          end
+
+          self.rocket_job_batch_throttles += [ThrottleDefinition.new(method_name, filter)]
+        end
+
+        # Undefine a previously defined throttle
+        def undefine_batch_throttle(method_name)
+          rocket_job_batch_throttles.delete_if { |throttle| throttle.method_name == method_name }
+        end
+
+        # Has a throttle been defined?
+        def batch_throttle?(method_name)
+          rocket_job_batch_throttles.any? { |throttle| throttle.method_name == method_name }
+        end
+      end
+
+      private
+
+      ThrottleDefinition = Struct.new(:method_name, :filter)
+
+      # Returns the matching filter, or nil if no throttles were triggered.
+      def rocket_job_batch_evaluate_throttles(slice)
+        rocket_job_batch_throttles.each do |throttle|
+          throttle_exceeded = method(throttle.method_name).arity == 0 ? send(throttle.method_name) : send(throttle.method_name, slice)
+          next unless throttle_exceeded
+
+          logger.debug { "Batch Throttle: #{throttle.method_name} has been exceeded. #{self.class.name}:#{id}" }
+          filter = throttle.filter
+          return filter.is_a?(Proc) ? filter.call(self) : send(filter)
+        end
+        nil
+      end
+
+    end
+  end
+end

data/lib/rocket_job/batch/throttle_running_slices.rb

@@ -0,0 +1,53 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    # Throttle the number of slices of a specific batch job that are processed at the same time.
+    #
+    # Example:
+    #   class MyJob < RocketJob
+    #     include RocketJob::Batch
+    #
+    #     # Maximum number of slices to process at the same time for each running instance.
+    #     self.throttle_running_slices = 25
+    #
+    #     def perform(record)
+    #       # ....
+    #     end
+    #   end
+    #
+    # It attempts to ensure that the number of workers do not exceed this number.
+    # This is not a hard limit and it is possible for the number of workers to
+    # slightly exceed this value at times. It can also occur that the number of
+    # slices running can drop below this number for a short period.
+    #
+    # This value can be modified while a job is running. The change will be picked
+    # up at the start of processing slices, or after processing a slice and
+    # `re_check_seconds` has been exceeded.
+    #
+    # 0 or nil : No limits in place
+    #
+    # Default: nil
+    module ThrottleRunningSlices
+      extend ActiveSupport::Concern
+
+      included do
+        field :throttle_running_slices, type: Integer, class_attribute: true, user_editable: true, copy_on_restart: true
+
+        validates :throttle_running_slices, numericality: {greater_than_or_equal_to: 0}, allow_nil: true
+
+        define_batch_throttle :throttle_running_slices_exceeded?, filter: :throttle_filter_id
+      end
+
+      private
+
+      # Returns [Boolean] whether the throttle for this job has been exceeded
+      def throttle_running_slices_exceeded?(slice)
+        throttle_running_slices &&
+          (throttle_running_slices != 0) &&
+          (input.running.where(:id.ne => slice.id).count >= throttle_running_slices)
+      end
+
+    end
+  end
+end

data/lib/rocket_job/batch/worker.rb

@@ -0,0 +1,288 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    module Worker
+      extend ActiveSupport::Concern
+
+      included do
+        # While working on a slice, the current slice is available via this reader
+        attr_reader :rocket_job_slice, :rocket_job_record_number
+
+        private
+
+        attr_writer :rocket_job_slice, :rocket_job_record_number
+      end
+
+      # Processes records in each available slice for this job. Slices are processed
+      # one at a time to allow for concurrent calls to this method to increase
+      # throughput. Processing will continue until there are no more jobs available
+      # for this job.
+      #
+      # Returns [true|false] whether this job should be excluded from the next lookup
+      #
+      # Slices are destroyed after their records are successfully processed
+      #
+      # Results are stored in the output collection if `collect_output?`
+      # `nil` results from workers are kept if `collect_nil_output`
+      #
+      # If an exception was thrown the entire slice of records is marked as failed.
+      #
+      # If the mongo_ha gem has been loaded, then the connection to mongo is
+      # automatically re-established and the job will resume anytime a
+      # Mongo connection failure occurs.
+      #
+      # Thread-safe, can be called by multiple threads at the same time
+      def rocket_job_work(worker, re_raise_exceptions = false, filter = {})
+        raise 'Job must be started before calling #rocket_job_work' unless running?
+        start_time = Time.now
+        if sub_state != :processing
+          rocket_job_handle_callbacks(worker, re_raise_exceptions)
+          return false unless running?
+        end
+
+        while !worker.shutdown?
+          if slice = input.next_slice(worker.name)
+            # Grab a slice before checking the throttle to reduce concurrency race condition.
+            if new_filter = rocket_job_batch_evaluate_throttles(slice)
+              # Restore retrieved slice so that other workers can process it later.
+              slice.set(worker_name: nil, state: :queued, started_at: nil)
+              self.class.send(:rocket_job_merge_filter, filter, new_filter)
+              return true
+            end
+
+            SemanticLogger.named_tagged(slice: slice.id.to_s) do
+              rocket_job_process_slice(slice, re_raise_exceptions)
+            end
+          else
+            break if record_count && rocket_job_batch_complete?(worker.name)
+            logger.debug 'No more work available for this job'
+            self.class.send(:rocket_job_merge_filter, filter, throttle_filter_id)
+            return true
+          end
+
+          # Allow new jobs with a higher priority to interrupt this job
+          break if (Time.now - start_time) >= Config.instance.re_check_seconds
+        end
+        false
+      end
+
+      # Prior to a job being made available for processing it can be processed one
+      # slice at a time.
+      #
+      # For example, to extract the header row which would be in the first slice.
+      #
+      # Returns [Integer] the number of records processed in the slice
+      #
+      # Note: The slice will be removed from processing when this method completes
+      def work_first_slice(&block)
+        raise '#work_first_slice can only be called from within before_batch callbacks' unless sub_state == :before
+        # TODO Make these settings configurable
+        count        = 0
+        wait_seconds = 5
+        while (slice = input.first).nil?
+          break if count > 10
+          logger.info "First slice has not arrived yet, sleeping for #{wait_seconds} seconds"
+          sleep wait_seconds
+          count += 1
+        end
+
+        if slice = input.first
+          SemanticLogger.named_tagged(slice: slice.id.to_s) do
+            # TODO Persist that the first slice is being processed by this worker
+            slice.start
+            rocket_job_process_slice(slice, true, &block)
+          end
+        else
+          # No records processed
+          0
+        end
+      end
+
+      # Returns [Array<ActiveWorker>] All workers actively working on this job
+      def rocket_job_active_workers(server_name = nil)
+        servers = []
+        case sub_state
+        when :before, :after
+          unless server_name && !worker_on_server?(server_name)
+            servers << ActiveWorker.new(worker_name, started_at, self) if running?
+          end
+        when :processing
+          query = input.running
+          query = query.where(worker_name: /\A#{server_name}/) if server_name
+          query.each do |slice|
+            servers << ActiveWorker.new(slice.worker_name, slice.started_at, self)
+          end
+        end
+        servers
+      end
+
+      private
+
+      # Process a single slice from Mongo
+      # Once the slice has been successfully processed it will be removed from the input collection
+      # Returns [Integer] the number of records successfully processed
+      def rocket_job_process_slice(slice, re_raise_exceptions)
+        slice_record_number       = 0
+        @rocket_job_record_number = slice.first_record_number || 0
+        @rocket_job_slice         = slice
+        run_callbacks :slice do
+          RocketJob::Sliced::Writer::Output.collect(self, slice) do |writer|
+            slice.each do |record|
+              slice_record_number += 1
+              SemanticLogger.named_tagged(record: @rocket_job_record_number) do
+                if _perform_callbacks.empty?
+                  @rocket_job_output = block_given? ? yield(record) : perform(record)
+                else
+                  # Allows @rocket_job_input to be modified by before/around callbacks
+                  @rocket_job_input = record
+                  # Allow callbacks to fail, complete or abort the job
+                  if running?
+                    if block_given?
+                      run_callbacks(:perform) { @rocket_job_output = yield(@rocket_job_input) }
+                    else
+                      # Allows @rocket_job_output to be modified by after/around callbacks
+                      run_callbacks(:perform) { @rocket_job_output = perform(@rocket_job_input) }
+                    end
+                  end
+                end
+                writer << @rocket_job_output
+              end
+              # JRuby says self.rocket_job_record_number= is private and cannot be accessed
+              @rocket_job_record_number += 1
+            end
+          end
+          @rocket_job_input = @rocket_job_slice = @rocket_job_output = nil
+        end
+
+        # On successful completion remove the slice from the input queue
+        # TODO Option to complete slice instead of destroying it to retain input data
+        slice.destroy
+        slice_record_number
+      rescue Exception => exc
+        slice.fail!(exc, slice_record_number)
+        raise exc if re_raise_exceptions
+        slice_record_number > 0 ? slice_record_number - 1 : 0
+      end
+
+      # Checks for completion and runs after_batch if defined
+      # Returns true if the job is now complete/aborted/failed
+      def rocket_job_batch_complete?(worker_name)
+        return true unless running?
+        return false unless record_count
+
+        # Only failed slices left?
+        input_count  = input.count
+        failed_count = input.failed.count
+        if (failed_count > 0) && (input_count == failed_count)
+          # Reload to pull in any counters or other data that was modified.
+          reload unless new_record?
+          if may_fail?
+            fail_job = true
+            unless new_record?
+              # Fail job iff no other worker has already finished it
+              # Must set write concern to at least 1 since we need the nModified back
+              result   = self.class.with(write: {w: 1}) do |query|
+                query.
+                  where(id: id, state: :running, sub_state: :processing).
+                  update({'$set' => {state: :failed, worker_name: worker_name}})
+              end
+              fail_job = false unless result.modified_count > 0
+            end
+            if fail_job
+              message        = "#{failed_count} slices failed to process"
+              self.exception = JobException.new(message: message)
+              fail!(worker_name, message)
+            end
+          end
+          return true
+        end
+
+        # Any work left?
+        return false if input_count > 0
+
+        # If the job was not saved to the queue, do not save any changes
+        if new_record?
+          rocket_job_batch_run_after_callbacks(false)
+          return true
+        end
+
+        # Complete job iff no other worker has already completed it
+        # Must set write concern to at least 1 since we need the nModified back
+        result = self.class.with(write: {w: 1}) do |query|
+          query.
+            where(id: id, state: :running, sub_state: :processing).
+            update('$set' => {sub_state: :after, worker_name: worker_name})
+        end
+
+        # Reload to pull in any counters or other data that was modified.
+        reload
+        if result.modified_count > 0
+          rocket_job_batch_run_after_callbacks(false)
+        else
+          # Repeat cleanup in case this worker was still running when the job was aborted
+          cleanup! if aborted?
+        end
+        true
+      end
+
+      # Run the before_batch callbacks
+      # Saves the current state before and after running callbacks if callbacks present
+      def rocket_job_batch_run_before_callbacks
+        unless _before_batch_callbacks.empty?
+          self.sub_state = :before
+          save! unless new_record? || destroyed?
+          logger.measure_info(
+            'before_batch',
+            metric:             "#{self.class.name}/before_batch",
+            log_exception:      :full,
+            on_exception_level: :error,
+            silence:            log_level
+          ) do
+            run_callbacks(:before_batch)
+          end
+        end
+        self.sub_state = :processing
+        save! unless new_record? || destroyed?
+      end
+
+      # Run the after_batch callbacks
+      # Saves the current state before and after running callbacks if callbacks present
+      def rocket_job_batch_run_after_callbacks(save_before = true)
+        unless _after_batch_callbacks.empty?
+          self.sub_state = :after
+          save! if save_before && !new_record? && !destroyed?
+          logger.measure_info(
+            'after_batch',
+            metric:             "#{self.class.name}/after_batch",
+            log_exception:      :full,
+            on_exception_level: :error,
+            silence:            log_level
+          ) do
+            run_callbacks(:after_batch)
+          end
+        end
+        if new_record? || destroyed?
+          complete if may_complete?
+        else
+          may_complete? ? complete! : save!
+        end
+      end
+
+      # Handle before and after callbacks
+      def rocket_job_handle_callbacks(worker, re_raise_exceptions)
+        rocket_job_fail_on_exception!(worker.name, re_raise_exceptions) do
+          # If this is the first worker to pickup this job
+          if sub_state == :before
+            rocket_job_batch_run_before_callbacks
+            # Check for 0 record jobs
+            rocket_job_batch_complete?(worker.name) if running?
+          elsif sub_state == :after
+            rocket_job_batch_run_after_callbacks
+          end
+        end
+      end
+
+    end
+  end
+end