workhorse 0.0.1 → 0.0.2

data/lib/workhorse/daemon.rb

@@ -0,0 +1,164 @@
+module Workhorse
+  class Daemon
+    def initialize(count: 1, pidfile: nil, quiet: false, &block)
+      @count = count
+      @pidfile = pidfile
+      @quiet = quiet
+      @block = block
+
+      fail 'Count must be an integer > 0.' unless count.is_a?(Integer) && count > 0
+
+      if @pidfile.nil?
+        @pidfile = count > 1 ? 'tmp/pids/workhorse.%i.pid' : 'tmp/pids/workhorse.pid'
+      elsif @count > 1 && !@pidfile.include?('%s')
+        fail 'Pidfile must include placeholder "%s" for worker id when specifying a count > 1.'
+      elsif @count == 0 && @pidfile.include?('%s')
+        fail 'Pidfile must not include placeholder "%s" for worker id when specifying a count of 1.'
+      end
+    end
+
+    def start
+      code = 0
+
+      for_each_worker do |worker_id|
+        pid_file, pid = read_pid(worker_id)
+
+        if pid_file && pid
+          warn "Worker ##{worker_id}: Already started (PID #{pid})"
+          code = 1
+        elsif pid_file
+          File.delete pid_file
+          puts "Worker ##{worker_id}: Starting (stale pid file)"
+          start_worker worker_id
+        else
+          warn "Worker ##{worker_id}: Starting"
+          start_worker worker_id
+        end
+      end
+
+      return code
+    end
+
+    def stop
+      code = 0
+
+      for_each_worker do |worker_id|
+        pid_file, pid = read_pid(worker_id)
+
+        if pid_file && pid
+          puts "Worker ##{worker_id}: Stopping"
+          stop_worker pid_file, pid
+        elsif pid_file
+          File.delete pid_file
+          puts "Worker ##{worker_id}: Already stopped (stale PID file)"
+        else
+          warn "Worker ##{worker_id}: Already stopped"
+          code = 1
+        end
+      end
+
+      return code
+    end
+
+    def status(quiet: false)
+      code = 0
+
+      for_each_worker do |worker_id|
+        pid_file, pid = read_pid(worker_id)
+
+        if pid_file && pid
+          puts "Worker ##{worker_id}: Running" unless quiet
+        elsif pid_file
+          warn "Worker ##{worker_id}: Not running (stale PID file)" unless quiet
+          code = 1
+        else
+          warn "Worker ##{worker_id}: Not running" unless quiet
+          code = 1
+        end
+      end
+
+      return code
+    end
+
+    def watch
+      if defined?(Rails)
+        should_be_running = !File.exist?(Rails.root.join('tmp/stop.txt'))
+      else
+        should_be_running = true
+      end
+
+      if should_be_running && status(quiet: true) != 0
+        return start
+      else
+        return 0
+      end
+    end
+
+    def restart
+      stop
+      return start
+    end
+
+    private
+
+    def for_each_worker(&block)
+      1.upto(@count, &block)
+    end
+
+    def start_worker(worker_id)
+      pid = fork do
+        $0 = process_name(worker_id)
+        @block.call
+      end
+      IO.write(pid_file_for(worker_id), pid)
+    end
+
+    def stop_worker(pid_file, pid)
+      loop do
+        begin
+          Process.kill('TERM', pid)
+        rescue Errno::ESRCH
+          break
+        end
+
+        sleep 1
+      end
+
+      File.delete(pid_file)
+    end
+
+    def process_name(worker_id)
+      if defined?(Rails)
+        path = Rails.root
+      else
+        path = $PROGRAM_NAME
+      end
+
+      return "Workhorse Worker ##{worker_id}: #{path}"
+    end
+
+    def process?(pid)
+      return begin
+        Process.getpgid(pid)
+        true
+      rescue Errno::ESRCH
+        false
+      end
+    end
+
+    def pid_file_for(worker_id)
+      @pidfile % worker_id
+    end
+
+    def read_pid(worker_id)
+      file = pid_file_for(worker_id)
+
+      if File.exist?(file)
+        pid = IO.read(file).to_i
+        return file, process?(pid) ? pid : nil
+      else
+        return nil, nil
+      end
+    end
+  end
+end

data/lib/workhorse/daemon/shell_handler.rb

@@ -0,0 +1,54 @@
+module Workhorse
+  class Daemon::ShellHandler
+    def self.run(options = {}, &block)
+      unless ARGV.count == 1
+        usage
+        exit 99
+      end
+
+      lockfile_path = options.delete(:lockfile) || 'workhorse.lock'
+      lockfile = File.open(lockfile_path, 'a')
+      lockfile.flock(File::LOCK_EX || File::LOCK_NB)
+
+      daemon = Workhorse::Daemon.new(options, &block)
+
+      begin
+        case ARGV.first
+        when 'start'
+          exit daemon.start
+        when 'stop'
+          exit daemon.stop
+        when 'status'
+          exit daemon.status
+        when 'watch'
+          exit daemon.watch
+        when 'restart'
+          exit daemon.restart
+        when 'usage'
+          usage
+          exit 99
+        else
+          usage
+        end
+
+        exit 0
+      rescue => e
+        warn "#{e.message}\n#{e.backtrace.join("\n")}"
+        exit 99
+      ensure
+        lockfile.flock(File::LOCK_UN)
+      end
+    end
+
+    def self.usage
+      warn <<~USAGE
+        Usage: #{$PROGRAM_NAME} start|stop|status|watch|restart|usage
+
+        Exit status:
+         0  if OK,
+         1  if at least one worker has an unexpected status,
+         99 on all other errors.
+      USAGE
+    end
+  end
+end

data/lib/workhorse/db_job.rb

@@ -0,0 +1,67 @@
+module Workhorse
+  class DbJob < ActiveRecord::Base
+    STATE_WAITING   = :waiting
+    STATE_LOCKED    = :locked
+    STATE_STARTED   = :started
+    STATE_SUCCEEDED = :succeeded
+    STATE_FAILED    = :failed
+
+    self.table_name = 'jobs'
+
+    def mark_locked!(worker_id)
+      if changed?
+        fail "Dirty jobs can't be locked."
+      end
+
+      if locked_at
+        fail "Job #{id} is already locked by #{locked_by.inspect}."
+      end
+
+      self.locked_at = Time.now
+      self.locked_by = worker_id
+      self.state     = STATE_LOCKED
+      save!
+    end
+
+    def mark_started!
+      assert_state! STATE_LOCKED
+
+      self.started_at = Time.now
+      self.state      = STATE_STARTED
+      save!
+    end
+
+    def mark_failed!(exception)
+      assert_state! STATE_LOCKED, STATE_STARTED
+
+      self.failed_at  = Time.now
+      self.last_error = %(#{exception.message}\n#{exception.backtrace.join("\n")})
+      self.state      = STATE_FAILED
+      save!
+    end
+
+    def mark_succeeded!
+      assert_state! STATE_STARTED
+
+      self.succeeded_at = Time.now
+      self.state        = STATE_SUCCEEDED
+      save!
+    end
+
+    def assert_state!(*states)
+      unless states.include?(state.to_sym)
+        fail "Job #{id} is not in state #{states.inspect} but in state #{state.inspect}."
+      end
+    end
+
+    def assert_locked_by!(worker_id)
+      assert_state! STATE_WAITING
+
+      if locked_by.nil?
+        fail "Job #{id} is not locked by any worker."
+      elsif locked_by != worker_id
+        fail "Job #{id} is locked by another worker (#{locked_by})."
+      end
+    end
+  end
+end

data/lib/workhorse/enqueuer.rb

@@ -0,0 +1,22 @@
+module Workhorse
+  class Enqueuer
+    # Enqueue any object that is serializable and has a `perform` method
+    def self.enqueue(job, queue: nil)
+      return DbJob.create!(
+        queue: queue,
+        handler: Marshal.dump(job)
+      )
+    end
+
+    # Enqueue an ActiveJob job
+    def self.enqueue_active_job(job)
+      enqueue job, queue: job.queue_name
+    end
+
+    # Enqueue the execution of an operation by its class and params
+    def self.enqueue_op(cls, params, queue: nil)
+      job = Workhorse::Jobs::RunRailsOp.new(cls, params)
+      enqueue job, queue: queue
+    end
+  end
+end

data/lib/workhorse/jobs/run_rails_op.rb

@@ -0,0 +1,12 @@
+module Workhorse::Jobs
+  class RunRailsOp
+    def initialize(cls, params = {})
+      @cls = cls
+      @params = params
+    end
+
+    def perform
+      @cls.run!(@params)
+    end
+  end
+end

data/lib/workhorse/performer.rb

@@ -0,0 +1,91 @@
+module Workhorse
+  class Performer
+    attr_reader :worker
+
+    def initialize(db_job, worker)
+      @db_job = db_job
+      @worker = worker
+      @started = false
+    end
+
+    def perform
+      fail 'Performer can only run once.' if @started
+      @started = true
+      perform!
+    end
+
+    private
+
+    def perform!
+      Thread.current[:workhorse_current_performer] = self
+
+      ActiveRecord::Base.connection_pool.with_connection do
+        if defined?(Rails) && Rails.application && Rails.application.respond_to?(:executor)
+          Rails.application.executor.wrap do
+            perform_wrapped
+          end
+        else
+          perform_wrapped
+        end
+      end
+    rescue => e
+      # ---------------------------------------------------------------
+      # Mark job as failed
+      # ---------------------------------------------------------------
+      log %(#{e.message}\n#{e.backtrace.join("\n")}), :error
+
+      Workhorse.tx_callback.call do
+        log 'Mark failed', :debug
+        @db_job.mark_failed!(e)
+      end
+    ensure
+      Thread.current[:workhorse_current_performer] = nil
+    end
+
+    def perform_wrapped
+      # ---------------------------------------------------------------
+      # Mark job as started
+      # ---------------------------------------------------------------
+      Workhorse.tx_callback.call do
+        log 'Marking as started', :debug
+        @db_job.mark_started!
+      end
+
+      # ---------------------------------------------------------------
+      # Deserialize and perform job
+      # ---------------------------------------------------------------
+      log 'Performing', :info
+
+      if Workhorse.perform_jobs_in_tx
+        Workhorse.tx_callback.call do
+          deserialized_job.perform
+        end
+      else
+        deserialized_job.perform
+      end
+
+      log 'Successfully performed', :info
+
+      # ---------------------------------------------------------------
+      # Mark job as succeeded
+      # ---------------------------------------------------------------
+      Workhorse.tx_callback.call do
+        log 'Mark succeeded', :debug
+        @db_job.mark_succeeded!
+      end
+    end
+
+    def log(text, level = :info)
+      text = "[#{@db_job.id}] #{text}"
+      worker.log text, level
+    end
+
+    def deserialized_job
+      # The source is safe as long as jobs are always enqueued using
+      # Workhorse::Enqueuer so it is ok to use Marshal.load.
+      # rubocop: disable Security/MarshalLoad
+      Marshal.load(@db_job.handler)
+      # rubocop: enable Security/MarshalLoad
+    end
+  end
+end

data/lib/workhorse/poller.rb

@@ -0,0 +1,119 @@
+module Workhorse
+  class Poller
+    attr_reader :worker
+
+    def initialize(worker)
+      @worker = worker
+      @running = false
+    end
+
+    def running?
+      @running
+    end
+
+    def start
+      fail 'Poller is already running.' if running?
+      @running = true
+
+      @thread = Thread.new do
+        begin
+          loop do
+            poll
+            break unless running?
+            sleep
+          end
+        rescue => e
+          worker.log %(Poller stopped with exception:\n#{e.message}\n#{e.backtrace.join("\n")})
+        end
+      end
+    end
+
+    def shutdown
+      fail 'Poller is not running.' unless running?
+      @running = false
+      wait
+    end
+
+    def wait
+      @thread.join
+    end
+
+    private
+
+    def sleep
+      remaining = worker.polling_interval
+
+      while running? && remaining > 0
+        Kernel.sleep 1
+        remaining -= 1
+      end
+    end
+
+    def poll
+      Workhorse.tx_callback.call do
+        idle = worker.idle
+        worker.log "Polling DB for jobs (#{idle} available threads)...", :debug
+
+        unless idle.zero?
+          jobs = queued_db_jobs(idle)
+          jobs.each do |job|
+            worker.log "Marking job #{job.id} as locked", :debug
+            job.mark_locked!(worker.id)
+            worker.perform job
+          end
+        end
+      end
+    end
+
+    def queued_db_jobs(limit)
+      table = Workhorse::DbJob.arel_table
+      is_oracle = ActiveRecord::Base.connection.adapter_name == 'OracleEnhanced'
+
+      # ---------------------------------------------------------------
+      # Lock all queued jobs that are not complete
+      # ---------------------------------------------------------------
+      Workhorse::DbJob.connection.execute(
+        Workhorse::DbJob.select('null').where(
+          table[:queue].not_eq(nil)
+          .and(table[:state].eq(:waiting))
+        ).lock.to_sql
+      )
+
+      # ---------------------------------------------------------------
+      # Select jobs to execute
+      # ---------------------------------------------------------------
+
+      # Fetch all waiting jobs of the correct queues
+      select = table.project(Arel.sql('*')).where(table[:state].eq(:waiting))
+
+      # Restrict queues that are currently in progress
+      bad_queries_select = table.project(table[:queue])
+                                .where(table[:state].in(%i[locked running]))
+                                .distinct
+      select = select.where(table[:queue].not_in(bad_queries_select))
+
+      # Restrict queues to "open" ones
+      if worker.queues.any?
+        where = table[:queue].in(worker.queues.reject(&:nil?))
+        if worker.queues.include?(nil)
+          where = where.or(table[:queue].eq(nil))
+        end
+        select = select.where(where)
+      end
+
+      # Order by creation date
+      select = select.order(table[:created_at].asc)
+
+      # Limit number of records
+      if is_oracle
+        select = select.where(Arel.sql('ROWNUM').lteq(limit))
+      else
+        select = select.take(limit)
+      end
+
+      select = select.lock
+
+      return Workhorse::DbJob.find_by_sql(select.to_sql)
+    end
+  end
+end