workhorse 0.0.1 → 0.0.2

data/lib/workhorse/pool.rb

@@ -0,0 +1,51 @@
+module Workhorse
+  # Abstraction layer of a simple thread pool implementation used by the worker.
+  class Pool
+    attr_reader :mutex
+
+    def initialize(size)
+      @size = size
+      @executor = Concurrent::ThreadPoolExecutor.new(
+        min_threads: 0,
+        max_threads: @size,
+        max_queue: 0,
+        fallback_policy: :abort,
+        auto_terminate: false
+      )
+      @mutex = Mutex.new
+      @active_threads = Concurrent::AtomicFixnum.new(0)
+    end
+
+    # Posts a new work unit to the pool.
+    def post
+      mutex.synchronize do
+        if @active_threads.value >= @size
+          fail 'All threads are busy.'
+        end
+
+        active_threads = @active_threads
+
+        active_threads.increment
+
+        @executor.post do
+          begin
+            yield
+          ensure
+            active_threads.decrement
+          end
+        end
+      end
+    end
+
+    # Returns the number of idle threads.
+    def idle
+      @size - @active_threads.value
+    end
+
+    # Shuts down the pool
+    def shutdown
+      @executor.shutdown
+      @executor.wait_for_termination
+    end
+  end
+end

data/lib/workhorse/worker.rb

@@ -0,0 +1,144 @@
+module Workhorse
+  class Worker
+    LOG_LEVELS = %i[fatal error warn info debug].freeze
+    SHUTDOWN_SIGNALS = %w[TERM INT].freeze
+
+    attr_reader :queues
+    attr_reader :state
+    attr_reader :pool_size
+    attr_reader :polling_interval
+    attr_reader :mutex
+    attr_reader :logger
+
+    # Instantiates and starts a new worker with the given arguments and then
+    # waits for its completion (i.e. an interrupt).
+    def self.start_and_wait(*args)
+      worker = new(*args)
+      worker.start
+      worker.wait
+    end
+
+    # Instantiates a new worker. The worker is not automatically started.
+    #
+    # @param queues [Array] The queues you want this worker to process. If an
+    #   empty array is given, any queues will be processed. Queues need to be
+    #   specified as a symbol. To also process jobs without a queue, supply
+    #   `nil` within the array.
+    # @param pool_size [Integer] The number of jobs that will be processed
+    #   simultaneously. If this parameter is not given, it will be set to the
+    #   number of given queues + 1.
+    # @param polling_interval [Integer] Interval in seconds the database will
+    #   be polled for new jobs. Set this as high as possible to avoid
+    #   unnecessary database load Set this as high as possible to avoid
+    #   unnecessary database load.
+    # @param auto_terminate [Boolean] Whether to automatically shut down the
+    #   worker properly on INT and TERM signals.
+    # @param quiet [Boolean] If this is set to `false`, the worker will also log
+    #   to STDOUT.
+    # @param logger [Logger] An optional logger the worker will append to. This
+    #   can be any instance of ruby's `Logger` but is commonly set to
+    #   `Rails.logger`.
+    def initialize(queues: [], pool_size: nil, polling_interval: 5, auto_terminate: true, quiet: true, logger: nil)
+      @queues = queues
+      @pool_size = pool_size || queues.size + 1
+      @polling_interval = polling_interval
+      @auto_terminate = auto_terminate
+      @state = :initialized
+      @quiet = quiet
+
+      @mutex = Mutex.new
+      @pool = Pool.new(@pool_size)
+      @poller = Workhorse::Poller.new(self)
+      @logger = logger
+
+      fail 'Polling interval must be an integer.' unless @polling_interval.is_a?(Integer)
+    end
+
+    def log(text, level = :info)
+      text = "[Job worker #{id}] #{text}"
+      puts text unless @quiet
+      return unless logger
+      fail "Log level #{level} is not available. Available are #{LOG_LEVELS.inspect}." unless LOG_LEVELS.include?(level)
+      logger.send(level, text.strip)
+    end
+
+    def id
+      @id ||= "#{Socket.gethostname}.#{Process.pid}.#{SecureRandom.hex(3)}"
+    end
+
+    # Starts the worker. This call is not blocking - call {wait} for this
+    # purpose.
+    def start
+      mutex.synchronize do
+        assert_state! :initialized
+        log 'Starting up'
+        @state = :running
+        @poller.start
+        log 'Started up'
+
+        trap_termination if @auto_terminate
+      end
+    end
+
+    def assert_state!(state)
+      fail "Expected worker to be in state #{state} but current state is #{self.state}." unless self.state == state
+    end
+
+    # Shuts down worker and DB poller. Jobs currently beeing processed are
+    # properly finished before this method returns. Subsequent calls to this
+    # method are ignored.
+    def shutdown
+      mutex.synchronize do
+        return if @state == :shutdown
+        assert_state! :running
+        log 'Shutting down'
+        @state = :shutdown
+
+        @poller.shutdown
+        @pool.shutdown
+        log 'Shut down'
+      end
+    end
+
+    # Waits until the worker is shut down. This only happens if shutdown gets
+    # called - either by another thread or by enabling `auto_terminate` and
+    # receiving a respective signal. Use this method to let worker run
+    # undefinitely.
+    def wait
+      assert_state! :running
+      @poller.wait
+    end
+
+    def idle
+      @pool.idle
+    end
+
+    def perform(db_job)
+      mutex.synchronize do
+        assert_state! :running
+        log "Posting job #{db_job.id} to thread pool"
+
+        @pool.post do
+          begin
+            Workhorse::Performer.new(db_job, self).perform
+          rescue => e
+            log %(#{e.message}\n#{e.backtrace.join("\n")}), :error
+          end
+        end
+      end
+    end
+
+    private
+
+    def trap_termination
+      SHUTDOWN_SIGNALS.each do |signal|
+        Signal.trap(signal) do
+          Thread.new do
+            log "\nCaught #{signal}, shutting worker down..."
+            shutdown
+          end
+        end
+      end
+    end
+  end
+end

data/test/lib/db_schema.rb

@@ -0,0 +1,20 @@
+ActiveRecord::Schema.define do
+  self.verbose = false
+
+  create_table :jobs, force: true do |t|
+    t.string :state, null: false, default: 'waiting'
+    t.string :queue, null: true
+    t.text :handler, null: false
+
+    t.string :locked_by
+    t.datetime :locked_at
+
+    t.datetime :started_at
+
+    t.datetime :succeeded_at
+    t.datetime :failed_at
+    t.text :last_error
+
+    t.timestamps null: false
+  end
+end

data/test/lib/jobs.rb

@@ -0,0 +1,46 @@
+class BasicJob
+  class_attribute :results
+  self.results = Concurrent::Array.new
+
+  def initialize(some_param: nil, sleep_time: 1)
+    @some_param = some_param
+    @sleep_time = sleep_time
+  end
+
+  def perform
+    results << @some_param
+    sleep @sleep_time
+  end
+end
+
+class DbConnectionTestJob
+  class_attribute :db_connections
+  self.db_connections = Concurrent::Array.new
+
+  def perform
+    db_connections << ActiveRecord::Base.connection.object_id
+  end
+end
+
+class DummyRailsOpsOp
+  class_attribute :results
+  self.results = Concurrent::Array.new
+
+  def self.run!(params = {})
+    new(params).run!
+  end
+
+  def initialize(params = {})
+    @params = params
+  end
+
+  def run!
+    perform
+  end
+
+  private
+
+  def perform
+    results << @params
+  end
+end

data/test/lib/test_helper.rb

@@ -0,0 +1,29 @@
+require 'minitest/autorun'
+require 'active_record'
+require 'mysql2'
+require 'benchmark'
+require 'jobs'
+
+class WorkhorseTest < ActiveSupport::TestCase
+  def setup
+    Workhorse::DbJob.delete_all
+  end
+
+  protected
+
+  def work(time = 2, options = {})
+    options[:pool_size] ||= 5
+    options[:polling_interval] ||= 1
+
+    w = Workhorse::Worker.new(options)
+    w.start
+    sleep time
+    w.shutdown
+  end
+end
+
+ActiveRecord::Base.logger = Logger.new('debug.log')
+ActiveRecord::Base.establish_connection adapter: 'mysql2', database: 'workhorse', username: 'travis', password: '', pool: 10, host: :localhost
+
+require 'db_schema'
+require 'workhorse'

data/test/workhorse/enqueuer_test.rb

@@ -0,0 +1,42 @@
+require 'test_helper'
+
+class Workhorse::EnqueuerTest < WorkhorseTest
+  def test_basic
+    assert_equal 0, Workhorse::DbJob.all.count
+    Workhorse::Enqueuer.enqueue BasicJob.new
+    assert_equal 1, Workhorse::DbJob.all.count
+
+    db_job = Workhorse::DbJob.first
+    assert_equal 'waiting', db_job.state
+    assert_equal Marshal.dump(BasicJob.new), db_job.handler
+    assert_nil db_job.locked_by
+    assert_nil db_job.queue
+    assert_nil db_job.locked_at
+    assert_nil db_job.started_at
+    assert_nil db_job.last_error
+    assert_not_nil db_job.created_at
+    assert_not_nil db_job.updated_at
+  end
+
+  def test_with_queue
+    assert_equal 0, Workhorse::DbJob.all.count
+    Workhorse::Enqueuer.enqueue BasicJob.new, queue: :q1
+    assert_equal 1, Workhorse::DbJob.all.count
+
+    db_job = Workhorse::DbJob.first
+    assert_equal 'q1', db_job.queue
+  end
+
+  def test_op
+    Workhorse::Enqueuer.enqueue_op DummyRailsOpsOp, { foo: :bar }, queue: :q1
+
+    w = Workhorse::Worker.new(queues: [:q1])
+    w.start
+    sleep 1.5
+    w.shutdown
+
+    assert_equal 'succeeded', Workhorse::DbJob.first.state
+
+    assert_equal [{ foo: :bar }], DummyRailsOpsOp.results
+  end
+end

data/test/workhorse/performer_test.rb

@@ -0,0 +1,18 @@
+require 'test_helper'
+
+class Workhorse::WorkerTest < WorkhorseTest
+  # This test makes sure that concurrent jobs always work in different database
+  # connections.
+  def test_db_connections
+    w = Workhorse::Worker.new polling_interval: 1, pool_size: 5
+    5.times do
+      Workhorse::Enqueuer.enqueue DbConnectionTestJob.new
+    end
+    w.start
+    sleep 2
+    w.shutdown
+
+    assert_equal 5, DbConnectionTestJob.db_connections.count
+    assert_equal 5, DbConnectionTestJob.db_connections.uniq.count
+  end
+end

data/test/workhorse/poller_test.rb

@@ -0,0 +1,13 @@
+require 'test_helper'
+
+class Workhorse::PollerTest < WorkhorseTest
+  def test_interruptable_sleep
+    w = Workhorse::Worker.new(polling_interval: 60)
+    w.start
+    sleep 2
+
+    Timeout.timeout(1.5) do
+      w.shutdown
+    end
+  end
+end

data/test/workhorse/pool_test.rb

@@ -0,0 +1,72 @@
+require 'test_helper'
+
+class Workhorse::PoolTest < WorkhorseTest
+  def test_idle
+    with_pool 5 do |p|
+      assert_equal 5, p.idle
+
+      4.times do |_i|
+        p.post do
+          sleep 1
+        end
+      end
+
+      sleep 0.5
+      assert_equal 1, p.idle
+
+      sleep 1
+      assert_equal 5, p.idle
+    end
+  end
+
+  def test_overflow
+    with_pool 5 do |p|
+      5.times { p.post { sleep 1 } }
+
+      exception = assert_raises do
+        p.post { sleep 1 }
+      end
+
+      assert_equal 'All threads are busy.', exception.message
+    end
+  end
+
+  def test_work
+    with_pool 5 do |p|
+      counter = Concurrent::AtomicFixnum.new(0)
+
+      5.times do
+        p.post do
+          sleep 1
+          counter.increment
+        end
+      end
+
+      sleep 1.2
+
+      assert_equal 5, counter.value
+
+      2.times do
+        p.post do
+          sleep 1
+          counter.increment
+        end
+      end
+
+      sleep 1.2
+
+      assert_equal 7, counter.value
+    end
+  end
+
+  private
+
+  def with_pool(size)
+    p = Workhorse::Pool.new(size)
+    begin
+      yield(p)
+    ensure
+      p.shutdown
+    end
+  end
+end