rocketjob 3.5.2 → 4.0.0

data/lib/rocket_job/batch/model.rb

@@ -0,0 +1,157 @@
+require 'active_support/concern'
+module RocketJob
+  module Batch
+    # Model attributes
+    module Model
+      extend ActiveSupport::Concern
+
+      included do
+        #
+        # User definable attributes
+        #
+        # The following attributes are set when the job is created
+
+        # Number of records to include in each slice that is processed
+        # Note:
+        #   slice_size is only used by SlicedJob#upload & Sliced::Input#upload
+        #   When slices are supplied directly, their size is not modified to match this number
+        field :slice_size, type: Integer, default: 100, class_attribute: true, user_editable: true, copy_on_restart: true
+
+        # Whether to retain nil results.
+        #
+        # Only applicable if `collect_output` is `true`
+        # Set to `false` to prevent collecting output from the perform
+        # method when it returns `nil`.
+        field :collect_nil_output, type: Boolean, default: true, class_attribute: true
+
+        # Optional Array<Symbol> list of categories that this job can output to
+        #
+        # By using categories the output from #perform can be placed in different
+        # output collections, and therefore different output files
+        #
+        # Categories must be declared in advance to avoid a #perform method
+        # accidentally writing its results to an unknown category
+        field :output_categories, type: Array, default: [:main], class_attribute: true
+
+        # Optional Array<Symbol> list of categories that this job can load input data into
+        field :input_categories, type: Array, default: [:main], class_attribute: true
+
+        # The file name of the uploaded file, if any.
+        # Set by #upload if a file name was supplied, but can also be set explicitly.
+        # May or may not include the fully qualified path name.
+        field :upload_file_name, type: String
+
+        #
+        # Values that jobs can also update during processing
+        #
+
+        # Number of records in this job
+        # Note:
+        #   A record_count of nil means it has not been set and workers will
+        #   _not_ complete the job when processing slices.
+        #   This allows workers to start processing slices while slices are still
+        #   being uploaded
+        field :record_count, type: Integer
+
+        #
+        # Read-only attributes
+        #
+
+        # Breaks the :running state up into multiple sub-states:
+        #   :running -> :before -> :processing -> :after -> :complete
+        field :sub_state, type: Symbol
+
+        validates_presence_of :slice_size
+
+        validates_each :output_categories, :input_categories do |record, attr, value|
+          # Under some circumstances ActiveModel is passing in a nil value even though the
+          # attributes have default values
+          Array(value).each do |category|
+            record.errors.add(attr, 'must only contain Symbol values') unless category.kind_of?(Symbol)
+            record.errors.add(attr, 'must only consist of lowercase characters, digits, and _') unless category.to_s =~ /\A[a-z_0-9]+\Z/
+          end
+        end
+      end
+
+      # Returns [Integer] percent of records completed so far
+      # Returns 0 if the total record count has not yet been set
+      def percent_complete
+        return 100 if completed?
+        return 0 unless record_count.to_i > 0
+
+        # Approximate number of input records
+        input_records = input.count.to_f * slice_size
+        if input_records > record_count
+          # Sanity check in case slice_size is not being adhered to
+          99
+        else
+          ((1.0 - (input_records.to_f / record_count)) * 100).to_i
+        end
+      end
+
+      # Returns [Hash] status of this job
+      def status(time_zone = 'Eastern Time (US & Canada)')
+        h = {}
+        case
+        when queued?
+          h['queued_slices'] = input.queued.count
+        when running? || paused? || failed?
+          h['active_slices'] = worker_count
+          h['failed_slices'] = input.failed.count
+          h['queued_slices'] = input.queued.count
+          # Very high level estimated time left
+          if record_count && running? && (record_count > 0)
+            percent = percent_complete
+            if percent >= 5
+              secs                        = seconds.to_f
+              h['est_remaining_duration'] = RocketJob.seconds_as_duration((((secs / percent) * 100) - secs))
+            end
+          end
+        when completed?
+          secs                  = seconds.to_f
+          h['records_per_hour'] = ((record_count.to_f / secs) * 60 * 60).round if record_count && (record_count > 0) && (secs > 0.0)
+        end
+        h['output_slices'] = output.count if collect_output? && !completed?
+        h.merge!(super(time_zone))
+        h.delete('result')
+        # Worker name should be retrieved from the slices when processing
+        h.delete('worker_name') if sub_state == :processing
+        h
+      end
+
+      # Returns [Array<String>] names of workers currently working this job.
+      def worker_names
+        return [] unless running?
+
+        case sub_state
+        when :before, :after
+          worker_name
+        when :processing
+          input.running.collect { |slice| slice.worker_name }
+        else
+          []
+        end
+      end
+
+      # Returns [Integer] the number of workers currently working on this job.
+      def worker_count
+        return 0 unless running?
+        # Cache the number of workers for 1 second.
+        return @worker_count if @worker_count_last && (@worker_count_last == Time.now.to_i)
+
+        @worker_count      =
+          case sub_state
+          when :before, :after
+            1
+          when :processing
+            input.running.count
+          else
+            0
+          end
+        @worker_count_last = Time.now.to_i
+        @worker_count
+      end
+
+    end
+  end
+end

data/lib/rocket_job/batch/performance.rb

@@ -0,0 +1,99 @@
+require 'optparse'
+require 'csv'
+require 'yaml'
+module RocketJob
+  module Batch
+    class Performance
+      attr_accessor :count, :servers, :workers, :version, :ruby, :environment, :mongo_config, :compress, :encrypt, :slice_size
+
+      def initialize
+        @count        = 10_000_000
+        @environment  = ENV['RAILS_ENV'] || ENV['RACK_ENV'] || 'development'
+        @mongo_config = 'config/mongoid.yml'
+        @compress     = false
+        @encrypt      = false
+        @slice_size   = 1000
+      end
+
+      def run_test_case(count = self.count)
+        servers = RocketJob::Server.count
+        raise 'Please start workers before starting the performance test' if servers == 0
+
+        count_running_workers
+
+        puts "Loading job with #{count} records/lines"
+        args = {log_level: :warn, slice_size: slice_size}
+        if defined?(::RocketJob::Enterprise)
+          args[:compress] = compress
+          args[:encrypt]  = encrypt
+        end
+        job = RocketJob::Jobs::PerformanceJob.new(args)
+        job.upload do |writer|
+          count.times { |i| writer << i }
+        end
+        job.save!
+
+        puts 'Waiting for job to complete'
+        while (!job.reload.completed?)
+          sleep 3
+        end
+
+        duration = job.completed_at - job.started_at
+        {count: count, duration: duration, records_per_second: (count.to_f / duration).round(3), workers: workers, servers: servers, compress: compress, encrypt: encrypt}
+      end
+
+      # Export the Results hash to a CSV file
+      def export_results(results)
+        ruby    = defined?(JRuby) ? "jruby_#{JRUBY_VERSION}" : "ruby_#{RUBY_VERSION}"
+        version = RocketJob::VERSION
+
+        CSV.open("job_results_#{ruby}_v#{version}.csv", 'wb') do |csv|
+          csv << results.first.keys
+          results.each { |result| csv << result.values }
+        end
+      end
+
+      # Parse command line options
+      def parse(argv)
+        parser        = OptionParser.new do |o|
+          o.on('-c', '--count COUNT', 'Count of records to enqueue') do |arg|
+            self.count = arg.to_i
+          end
+          o.on('-m', '--mongo MONGO_CONFIG_FILE_NAME', 'Location of mongoid.yml config file') do |arg|
+            self.mongo_config = arg
+          end
+          o.on('-e', '--environment ENVIRONMENT', 'The environment to run the app on (Default: RAILS_ENV || RACK_ENV || development)') do |arg|
+            self.environment = arg
+          end
+          o.on('-z', '--compress', 'Turn on compression') do |arg|
+            self.compress = true
+          end
+          o.on('-E', '--encrypt', 'Turn on encryption') do |arg|
+            self.encrypt = true
+          end
+          o.on('-s', '--slice_size COUNT', 'Slice size') do |arg|
+            self.slice_size = arg.to_i
+          end
+        end
+        parser.banner = 'rocketjob_batch_perf <options>'
+        parser.on_tail '-h', '--help', 'Show help' do
+          puts parser
+          exit 1
+        end
+        parser.parse! argv
+      end
+
+      def count_running_workers
+        self.servers = 0
+        self.workers = 0
+        RocketJob::Server.running.each do |server|
+          next if server.zombie?
+          self.servers += 1
+          self.workers += server.heartbeat.workers
+        end
+        puts "Running: #{workers} workers, distributed across #{servers} servers"
+      end
+
+    end
+  end
+end

data/lib/rocket_job/batch/result.rb

@@ -0,0 +1,8 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    # Structure to hold results that need to be written to different output collections
+    Result = Struct.new(:category, :value)
+  end
+end

data/lib/rocket_job/batch/results.rb

@@ -0,0 +1,9 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    # For holding multiple categorized Result's
+    class Results < Array
+    end
+  end
+end

data/lib/rocket_job/batch/state_machine.rb

@@ -0,0 +1,102 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    # State machine for sliced jobs
+    module StateMachine
+      extend ActiveSupport::Concern
+
+      included do
+        # Replace existing event and all callbacks for that event
+        aasm.state_machine.add_event(:retry, {}) do
+          # Retry before_batch
+          transitions from:  :failed, to: :queued,
+                      if:    -> { (sub_state == :before) || sub_state.nil? },
+                      after: :rocket_job_requeue_sub_state_before
+          # Retry perform and after_batch
+          transitions from:  :failed, to: :running,
+                      if:    -> { sub_state != :before },
+                      after: :rocket_job_requeue_sub_state_after
+        end
+
+        # Replace existing event and all callbacks for that event
+        aasm.state_machine.add_event(:requeue, {}) do
+          # Requeue perform
+          transitions from:  :running, to: :running,
+                      if:    -> server_name { sub_state == :processing },
+                      after: :rocket_job_requeue_sub_state_processing
+          # Requeue after_batch
+          transitions from:  :running, to: :running,
+                      if:    -> server_name { worker_on_server?(server_name) && (sub_state == :after) },
+                      after: :rocket_job_requeue_sub_state_after
+          # Requeue before_batch
+          transitions from:  :running, to: :queued,
+                      if:    -> server_name { worker_on_server?(server_name) && (sub_state == :before) },
+                      after: :rocket_job_requeue_sub_state_before
+        end
+
+        # Needed again here since the events have been overwritten above
+        before_retry :rocket_job_clear_exception
+
+        before_start :rocket_job_sub_state_before
+        before_complete :rocket_job_clear_sub_state
+        after_abort :cleanup!
+        after_retry :rocket_job_requeue_failed_slices
+        after_destroy :cleanup!
+      end
+
+      # Drop the input and output collections
+      def cleanup!
+        input_categories.each { |category| input(category).drop }
+        output_categories.each { |category| output(category).drop }
+      end
+
+      # A batch job can only be processed:
+      # - Whilst Queued (before processing).
+      # - During processing.
+      #
+      # I.e. Not during before_batch and after_batch.
+      def pausable?
+        queued? || paused? || running? && (sub_state == :processing)
+      end
+
+      private
+
+      # Is this job still being processed
+      def rocket_job_processing?
+        running? && (sub_state == :processing)
+      end
+
+      def rocket_job_sub_state_before
+        self.sub_state = :before unless self.sub_state
+      end
+
+      def rocket_job_clear_sub_state
+        self.sub_state = nil
+      end
+
+      # Called after a job in sub_state: :before is requeued
+      def rocket_job_requeue_sub_state_before
+        self.sub_state   = nil
+        self.started_at  = nil
+        self.worker_name = nil
+      end
+
+      def rocket_job_requeue_sub_state_after
+        self.sub_state   = :processing
+        self.worker_name = nil
+      end
+
+      def rocket_job_requeue_sub_state_processing(worker_name)
+        self.worker_name = nil
+        input.requeue_running(worker_name)
+      end
+
+      # Also retry failed slices when the job itself is re-tried
+      def rocket_job_requeue_failed_slices
+        input.requeue_failed
+      end
+
+    end
+  end
+end

data/lib/rocket_job/batch/statistics.rb

@@ -0,0 +1,88 @@
+require 'active_support/concern'
+
+module RocketJob
+  module Batch
+    # Allow statistics to be gathered while a batch job is running
+    module Statistics
+      extend ActiveSupport::Concern
+
+      class Stats
+        attr_reader :stats, :in_memory
+
+        # hash [Hash]
+        #   Update an `in-memory` copy of the stats instead of gathering them inside `stats`.
+        def initialize(hash = nil)
+          @in_memory = hash
+          @stats     = Hash.new(0) unless hash
+        end
+
+        def inc(hash)
+          hash.each_pair { |key, increment| inc_key(key, increment) }
+          self
+        end
+
+        def inc_key(key, increment = 1)
+          return if increment == 0
+          if in_memory
+            # For tests and in-process execution
+            inc_in_memory(key, increment)
+          elsif key && key != ''
+            stats["statistics.#{key}"] += increment
+          end
+          self
+        end
+
+        def empty?
+          stats.nil? || stats.empty?
+        end
+
+        private
+
+        # Navigates path and creates child hashes as needed at the end is reached
+        def inc_in_memory(key, increment)
+          paths = key.to_s.split('.')
+          last  = paths.pop
+          return unless last
+
+          target       = paths.inject(in_memory) { |target, key| target.key?(key) ? target[key] : target[key] = Hash.new(0) }
+          target[last] += increment
+        end
+      end
+
+      included do
+        field :statistics, type: Hash, default: -> { Hash.new(0) }
+
+        around_slice :statistics_capture
+      end
+
+      # Increment a statistic
+      def statistics_inc(key, increment = 1)
+        return if key.nil? || key == ''
+        # Being called within tests outside of a perform
+        @slice_statistics ||= Stats.new(new_record? ? statistics : nil)
+        key.is_a?(Hash) ? @slice_statistics.inc(key) : @slice_statistics.inc_key(key, increment)
+      end
+
+      private
+
+      # Capture the number of successful and failed tradelines
+      # as well as those with notices and alerts.
+      def statistics_capture
+        @slice_statistics = Stats.new(new_record? ? statistics : nil)
+        yield
+        collection.update_one({_id: id}, {'$inc' => @slice_statistics.stats}) unless @slice_statistics.empty?
+      end
+
+      # Overrides RocketJob::Batch::Logger#rocket_job_batch_log_payload
+      def rocket_job_batch_log_payload
+        h              = {
+          from:  aasm.from_state,
+          to:    aasm.to_state,
+          event: aasm.current_event
+        }
+        h[:statistics] = statistics.dup if statistics.present? && (completed? || failed?)
+        h
+      end
+    end
+  end
+end