dat-worker-pool 0.5.0 → 0.6.0

data/lib/dat-worker-pool/locked_object.rb

@@ -0,0 +1,60 @@
+require 'thread'
+
+class DatWorkerPool
+
+  class LockedObject
+    attr_reader :mutex
+
+    def initialize(object = nil)
+      @object = object
+      @mutex  = Mutex.new
+    end
+
+    def value
+      @mutex.synchronize{ @object }
+    end
+
+    def set(new_object)
+      @mutex.synchronize{ @object = new_object }
+    end
+
+    def with_lock(&block)
+      @mutex.synchronize{ block.call(@mutex, @object) }
+    end
+
+  end
+
+  class LockedArray < LockedObject
+    def initialize(array = nil)
+      super(array || [])
+    end
+
+    alias :values :value
+
+    def first;  @mutex.synchronize{ @object.first };  end
+    def last;   @mutex.synchronize{ @object.last };   end
+    def size;   @mutex.synchronize{ @object.size };   end
+    def empty?; @mutex.synchronize{ @object.empty? }; end
+
+    def push(new_item); @mutex.synchronize{ @object.push(new_item) }; end
+    def pop;            @mutex.synchronize{ @object.pop };            end
+
+    def shift;             @mutex.synchronize{ @object.shift };             end
+    def unshift(new_item); @mutex.synchronize{ @object.unshift(new_item) }; end
+
+    def delete(item); @mutex.synchronize{ @object.delete(item) }; end
+  end
+
+  class LockedSet < LockedObject
+    def initialize; super(Set.new); end
+
+    alias :values :value
+
+    def size;   @mutex.synchronize{ @object.size };   end
+    def empty?; @mutex.synchronize{ @object.empty? }; end
+
+    def add(item);    @mutex.synchronize{ @object.add(item) };    end
+    def remove(item); @mutex.synchronize{ @object.delete(item) }; end
+  end
+
+end

data/lib/dat-worker-pool/queue.rb

@@ -1,65 +1,88 @@
-require 'thread'
-
 class DatWorkerPool
 
-  class Queue
+  module Queue
 
-    attr_accessor :on_push_callbacks, :on_pop_callbacks
+    def self.included(klass)
+      klass.class_eval do
+        include InstanceMethods
+      end
+    end
 
-    def initialize
-      @work_items         = []
-      @shutdown           = false
-      @mutex              = Mutex.new
-      @condition_variable = ConditionVariable.new
+    module InstanceMethods
 
-      @on_pop_callbacks  = []
-      @on_push_callbacks = []
-    end
+      # overwrite this method to add custom logic for reading the current work
+      # items on the queue
+      def work_items
+        raise NotImplementedError
+      end
 
-    def start
-      @shutdown = false
-    end
+      def dwp_start
+        @dwp_running = true
+        start!
+      end
 
-    # * Wakes up any threads (`@condition_variable.broadcast`) who are sleeping
-    #   because of `pop`.
-    def shutdown
-      @shutdown = true
-      @mutex.synchronize{ @condition_variable.broadcast }
-    end
+      def dwp_signal_shutdown
+        @dwp_running = false
+      end
 
-    # * Add the work and wake up the first thread waiting from calling `pop`
-    #  (`@condition_variable.signal`).
-    def push(work_item)
-      raise "Unable to add work while shutting down" if @shutdown
-      @mutex.synchronize do
-        @work_items << work_item
-        @condition_variable.signal
+      def dwp_shutdown
+        self.dwp_signal_shutdown
+        shutdown!
       end
-      @on_push_callbacks.each(&:call)
-    end
 
-    # * Sleeps the current thread (`@condition_variable.wait(@mutex)`) until it
-    #   is signaled via `push` or `shutdown`.
-    def pop
-      return if @shutdown
-      item = @mutex.synchronize do
-        @condition_variable.wait(@mutex) while !@shutdown && @work_items.empty?
-        @work_items.shift
+      def running?
+        !!@dwp_running
       end
-      @on_pop_callbacks.each(&:call)
-      item
-    end
 
-    def work_items
-      @mutex.synchronize{ @work_items }
-    end
+      def shutdown?
+        !self.running?
+      end
 
-    def empty?
-      @mutex.synchronize{ @work_items.empty? }
-    end
+      def dwp_push(*args)
+        raise "Unable to add work when shut down" if self.shutdown?
+        push!(*args)
+      end
+
+      def dwp_pop
+        return if self.shutdown?
+        pop!
+      end
+
+      private
+
+      # overwrite this method to add custom start logic; this is a no-op by
+      # default because we don't require a queue to have custom start logic
+      def start!; end
+
+      # overwrite this method to add custom shutdown logic; this is a no-op by
+      # default because we don't require a queue to have custom shutdown logic;
+      # more than likely you will want to use this to "wakeup" worker threads
+      # that are sleeping waiting to pop work from the queue (see the default
+      # queue for an example using mutexes and condition variables)
+      def shutdown!; end
+
+      # overwrite this method to add custom push logic; this doesn't have to be
+      # overwritten but if it isn't, you will not be able to add work items using
+      # the queue (and the `add_work` method on `DatWorkerPool` will not work);
+      # more than likely this should add work to the queue and "signal" the
+      # workers so they know to process it (see the default queue for an example
+      # using mutexes and condition variables)
+      def push!(*args)
+        raise NotImplementedError
+      end
+
+      # overwrite this method to add custom pop logic; this has to be overwritten
+      # or the workers will not be able to get work that needs to be processed;
+      # this is intended to sleep the worker threads (see the default queue for an
+      # example using mutexes and condition variables); if this returns `nil` the
+      # workers will ignore it and go back to sleep, `nil` is not a valid work
+      # item to process; also check if the queue is shutdown when waking up
+      # workers, you probably don't want to hand-off work while everything is
+      # shutting down
+      def pop!
+        raise NotImplementedError
+      end
 
-    def shutdown?
-      @shutdown
     end
 
   end

data/lib/dat-worker-pool/runner.rb

@@ -0,0 +1,196 @@
+require 'set'
+require 'system_timer'
+require 'dat-worker-pool'
+require 'dat-worker-pool/locked_object'
+
+class DatWorkerPool
+
+  class Runner
+
+    attr_reader :num_workers, :worker_class, :worker_params
+    attr_reader :logger_proxy, :queue
+
+    def initialize(args)
+      @num_workers   = args[:num_workers]
+      @queue         = args[:queue]
+      @worker_class  = args[:worker_class]
+      @worker_params = args[:worker_params]
+
+      @logger_proxy = if args[:logger]
+        LoggerProxy.new(args[:logger])
+      else
+        NullLoggerProxy.new
+      end
+
+      @workers           = LockedArray.new
+      @available_workers = LockedSet.new
+    end
+
+    def workers
+      @workers.values
+    end
+
+    def start
+      log{ "Starting worker pool with #{@num_workers} worker(s)" }
+      @queue.dwp_start
+      @num_workers.times.each{ |n| build_worker(n + 1) }
+    end
+
+    # the workers should be told to shutdown before the queue because the queue
+    # shutdown will wake them up; a worker popping on a shutdown queue will
+    # always get `nil` back and will loop as fast as allowed until its shutdown
+    # flag is flipped, so shutting down the workers then the queue keeps them
+    # from looping as fast as possible; if any kind of standard error or the
+    # expected timeout error (assuming the workers take too long to shutdown) is
+    # raised, force a shutdown; this ensures we shutdown as best as possible
+    # instead of letting ruby kill the threads when the process exits;
+    # non-timeout errors will be re-raised so they can be caught and handled (or
+    # shown when the process exits)
+    def shutdown(timeout = nil, backtrace = nil)
+      log do
+        timeout_message = timeout ? "#{timeout} second(s)" : "none"
+        "Shutting down worker pool (timeout: #{timeout_message})"
+      end
+      begin
+        @workers.with_lock{ |m, ws| ws.each(&:dwp_signal_shutdown) }
+        @queue.dwp_signal_shutdown
+        OptionalTimeout.new(timeout) do
+          @queue.dwp_shutdown
+          wait_for_workers_to_shutdown
+        end
+      rescue StandardError => exception
+        force_workers_to_shutdown(exception, timeout, backtrace)
+        raise exception
+      rescue TimeoutInterruptError => exception
+        force_workers_to_shutdown(exception, timeout, backtrace)
+      end
+      log{ "Finished shutting down" }
+    end
+
+    def available_worker_count
+      @available_workers.size
+    end
+
+    def worker_available?
+      self.available_worker_count > 0
+    end
+
+    def make_worker_available(worker)
+      @available_workers.add(worker.object_id)
+    end
+
+    def make_worker_unavailable(worker)
+      @available_workers.remove(worker.object_id)
+    end
+
+    def log(&message_block)
+      @logger_proxy.runner_log(&message_block)
+    end
+
+    def worker_log(worker, &message_block)
+      @logger_proxy.worker_log(worker, &message_block)
+    end
+
+    private
+
+    def build_worker(number)
+      @workers.push(@worker_class.new(self, @queue, number).tap(&:dwp_start))
+    end
+
+    # use an until loop instead of each to join all the workers, while we are
+    # joining a worker a different worker can shutdown and remove itself from
+    # the `@workers` array; rescue when joining the workers, ruby will raise any
+    # exceptions that aren't handled by a thread when its joined, this allows
+    # all the workers to be joined
+    def wait_for_workers_to_shutdown
+      log{ "Waiting for #{@workers.size} workers to shutdown" }
+      while !(worker = @workers.first).nil?
+        begin
+          worker.dwp_join
+        rescue StandardError => exception
+          log{ "An error occurred while waiting for worker " \
+               "to shutdown ##{worker.dwp_number}" }
+        end
+        remove_worker(worker)
+        log{ "Worker ##{worker.dwp_number} shutdown" }
+      end
+    end
+
+    # use an until loop instead of each to join all the workers, while we are
+    # joining a worker a different worker can shutdown and remove itself from
+    # the `@workers` array; rescue when joining the workers, ruby will raise any
+    # exceptions that aren't handled by a thread when its joined, this ensures
+    # if the hard shutdown is raised and not rescued (for example, in the
+    # workers ensure), then it won't cause the forced shutdown to end
+    # prematurely
+    def force_workers_to_shutdown(orig_exception, timeout, backtrace)
+      log{ "Forcing #{@workers.size} workers to shutdown" }
+      error = build_forced_shutdown_error(orig_exception, timeout, backtrace)
+      while !(worker = @workers.first).nil?
+        worker.dwp_raise(error)
+        begin
+          worker.dwp_join
+        rescue StandardError => exception
+          log{ "An error occurred while waiting for worker " \
+               "to shutdown ##{worker.dwp_number} (forced)" }
+        rescue ShutdownError
+          # these are expected (because we raised them in the thread) so they
+          # don't need to be logged
+        end
+        remove_worker(worker)
+        log{ "Worker ##{worker.dwp_number} shutdown (forced)" }
+      end
+    end
+
+    # make sure the worker has been removed from the available workers, in case
+    # it errored before it was able to make itself unavailable
+    def remove_worker(worker)
+      self.make_worker_unavailable(worker)
+      @workers.delete(worker)
+    end
+
+    def build_forced_shutdown_error(orig_exception, timeout, backtrace)
+      if orig_exception.kind_of?(TimeoutInterruptError)
+        ShutdownError.new("Timed out shutting down (#{timeout} seconds).").tap do |e|
+          e.set_backtrace(backtrace) if backtrace
+        end
+      else
+        ShutdownError.new("Errored while shutting down: #{orig_exception.inspect}").tap do |e|
+          e.set_backtrace(orig_exception.backtrace)
+        end
+      end
+    end
+
+    # this needs to be an `Interrupt` to be sure we don't accidentally catch it
+    # when rescueing exceptions; in the shutdown methods we rescue any errors
+    # from `worker.join`, this will also rescue the timeout error if its a
+    # standard error and will keep it from doing a forced shutdown
+    TimeoutInterruptError = Class.new(Interrupt)
+
+    module OptionalTimeout
+      def self.new(seconds, &block)
+        if seconds
+          SystemTimer.timeout(seconds, TimeoutInterruptError, &block)
+        else
+          block.call
+        end
+      end
+    end
+
+    class LoggerProxy < Struct.new(:logger)
+      def runner_log(&message_block)
+        self.logger.debug("[DWP] #{message_block.call}")
+      end
+      def worker_log(worker, &message_block)
+        self.logger.debug("[DWP-#{worker.dwp_number}] #{message_block.call}")
+      end
+    end
+
+    class NullLoggerProxy
+      def runner_log(&block); end
+      def worker_log(worker, &block); end
+    end
+
+  end
+
+end

data/lib/dat-worker-pool/version.rb

@@ -1,3 +1,3 @@
 class DatWorkerPool
-  VERSION = "0.5.0"
+  VERSION = "0.6.0"
 end

data/lib/dat-worker-pool/worker.rb

@@ -1,91 +1,270 @@
 require 'thread'
-require 'dat-worker-pool'
+require 'dat-worker-pool/runner'
 
 class DatWorkerPool
 
-  class Worker
-
-    attr_accessor :on_work, :on_error_callbacks
-    attr_accessor :on_start_callbacks, :on_shutdown_callbacks
-    attr_accessor :on_sleep_callbacks, :on_wakeup_callbacks
-    attr_accessor :before_work_callbacks, :after_work_callbacks
-
-    def initialize(queue)
-      @queue = queue
-      @on_work = proc{ |worker, work_item| }
-      @on_error_callbacks    = []
-      @on_start_callbacks    = []
-      @on_shutdown_callbacks = []
-      @on_sleep_callbacks    = []
-      @on_wakeup_callbacks   = []
-      @before_work_callbacks = []
-      @after_work_callbacks  = []
-
-      @shutdown = false
-      @thread   = nil
-    end
+  module Worker
 
-    def start
-      @thread ||= Thread.new{ work_loop }
+    def self.included(klass)
+      klass.class_eval do
+        extend ClassMethods
+        include InstanceMethods
+      end
     end
 
-    def shutdown
-      @shutdown = true
-    end
+    module ClassMethods
 
-    def running?
-      !!(@thread && @thread.alive?)
-    end
+      def on_start_callbacks;       @on_start_callbacks       ||= []; end
+      def on_shutdown_callbacks;    @on_shutdown_callbacks    ||= []; end
+      def on_available_callbacks;   @on_available_callbacks   ||= []; end
+      def on_unavailable_callbacks; @on_unavailable_callbacks ||= []; end
+      def on_error_callbacks;       @on_error_callbacks       ||= []; end
+      def before_work_callbacks;    @before_work_callbacks    ||= []; end
+      def after_work_callbacks;     @after_work_callbacks     ||= []; end
 
-    def join(*args)
-      @thread.join(*args) if running?
-    end
+      def on_start(&block);       self.on_start_callbacks       << block; end
+      def on_shutdown(&block);    self.on_shutdown_callbacks    << block; end
+      def on_available(&block);   self.on_available_callbacks   << block; end
+      def on_unavailable(&block); self.on_unavailable_callbacks << block; end
+      def on_error(&block);       self.on_error_callbacks       << block; end
+      def before_work(&block);    self.before_work_callbacks    << block; end
+      def after_work(&block);     self.after_work_callbacks     << block; end
 
-    def raise(*args)
-      @thread.raise(*args) if running?
-    end
+      def prepend_on_start(&block);       self.on_start_callbacks.unshift(block);       end
+      def prepend_on_shutdown(&block);    self.on_shutdown_callbacks.unshift(block);    end
+      def prepend_on_available(&block);   self.on_available_callbacks.unshift(block);   end
+      def prepend_on_unavailable(&block); self.on_unavailable_callbacks.unshift(block); end
+      def prepend_on_error(&block);       self.on_error_callbacks.unshift(block);       end
+      def prepend_before_work(&block);    self.before_work_callbacks.unshift(block);    end
+      def prepend_after_work(&block);     self.after_work_callbacks.unshift(block);     end
 
-    private
-
-    # * Rescue `ShutdownError` but don't do anything with it. We want to handle
-    #   the error but we just want it to cause the worker to exit its work loop.
-    #   If the `ShutdownError` isn't rescued, it will be raised when the worker
-    #   is joined.
-    def work_loop
-      @on_start_callbacks.each{ |p| p.call(self) }
-      loop do
-        break if @shutdown
-        fetch_and_do_work
-      end
-    rescue ShutdownError
-    ensure
-      @on_shutdown_callbacks.each{ |p| p.call(self) }
-      @thread = nil
     end
 
-    # * Rescue `ShutdownError` but re-raise it after calling the error
-    #   callbacks. This ensures it causes the work loop to exit (see
-    #   `work_loop`).
-    def fetch_and_do_work
-      @on_sleep_callbacks.each{ |p| p.call(self) }
-      work_item = @queue.pop
-      @on_wakeup_callbacks.each{ |p| p.call(self) }
-      do_work(work_item) if work_item
-    rescue ShutdownError => exception
-      handle_exception(exception, work_item)
-      raise exception
-    rescue StandardError => exception
-      handle_exception(exception, work_item)
-    end
+    module InstanceMethods
+
+      attr_reader :dwp_number
+
+      def initialize(runner, queue, number)
+        @dwp_runner, @dwp_queue, @dwp_number = runner, queue, number
+        @dwp_running = false
+        @dwp_thread  = nil
+      end
+
+      def dwp_start
+        @dwp_running = true
+        @dwp_thread ||= Thread.new{ dwp_work_loop }
+      end
+
+      def dwp_signal_shutdown
+        @dwp_running = false
+      end
+
+      def dwp_running?
+        !!@dwp_running
+      end
+
+      def dwp_shutdown?
+        !self.dwp_running?
+      end
+
+      # this is needed because even if the running flag has been set to false
+      # (meaning the worker has been shutdown) the thread may still be alive
+      # because its `work` is taking a long time or its still trying to shut
+      # down
+      def dwp_thread_alive?
+        !!(@dwp_thread && @dwp_thread.alive?)
+      end
+
+      def dwp_join(*args)
+        @dwp_thread.join(*args) if self.dwp_thread_alive?
+      end
+
+      def dwp_raise(*args)
+        @dwp_thread.raise(*args) if self.dwp_thread_alive?
+      end
+
+      private
+
+      # Helpers
+      def number; @dwp_number;               end
+      def params; @dwp_runner.worker_params; end
+      def queue;  @dwp_runner.queue;         end
+
+      # overwrite this method to add custom work logic; this has to be
+      # overwritten or the workers will not know how to handle a work item
+      def work!(work_item)
+        raise NotImplementedError
+      end
+
+      # rescue `ShutdownError` but re-raise it after calling the on-error
+      # callbacks, this ensures it causes the loop to exit
+      def dwp_work_loop
+        dwp_setup
+        while self.dwp_running?
+          begin
+            if !(work_item = queue.dwp_pop).nil?
+              begin
+                dwp_make_unavailable
+                dwp_work(work_item)
+              rescue ShutdownError => exception
+                dwp_handle_exception(exception, work_item)
+                Thread.current.raise exception
+              rescue StandardError => exception
+                dwp_handle_exception(exception, work_item)
+              ensure
+                dwp_make_available
+              end
+            end
+          rescue StandardError => exception
+            dwp_handle_exception(exception, work_item)
+          end
+        end
+      ensure
+        dwp_teardown
+      end
+
+      def dwp_setup
+        dwp_log{ "Starting" }
+        begin
+          dwp_run_callback 'on_start'
+          dwp_make_available
+        rescue StandardError => exception
+          dwp_handle_exception(exception)
+          Thread.current.raise exception
+        end
+      end
+
+      # this is a separate method so the test runner can call it individually
+      def dwp_make_unavailable
+        @dwp_runner.make_worker_unavailable(self)
+        dwp_run_callback 'on_unavailable'
+        dwp_log{ "Unavailable" }
+      end
+
+      # this is a separate method so the test runner can call it individually
+      def dwp_make_available
+        @dwp_runner.make_worker_available(self)
+        dwp_run_callback 'on_available'
+        dwp_log{ "Available" }
+      end
+
+      # this is a separate method so the test runner can call it individually
+      def dwp_work(work_item)
+        dwp_log{ "Working, item: #{work_item.inspect}" }
+        dwp_run_callback('before_work', work_item)
+        work!(work_item)
+        dwp_run_callback('after_work', work_item)
+      end
+
+      def dwp_teardown
+        begin
+          dwp_make_unavailable
+          dwp_run_callback 'on_shutdown'
+        rescue StandardError => exception
+          dwp_handle_exception(exception)
+        end
+        dwp_log{ "Shutdown" }
+        @dwp_running = false
+        @dwp_thread  = nil
+      end
+
+      def dwp_handle_exception(exception, work_item = nil)
+        begin
+          dwp_log_exception(exception)
+          dwp_run_callback('on_error', exception, work_item)
+        rescue StandardError => on_error_exception
+          # errors while running on-error callbacks are logged but otherwise
+          # ignored to keep the worker from crashing, ideally these should be
+          # caught by the on-error callbacks themselves and never get here
+          dwp_log_exception(on_error_exception)
+        end
+      end
+
+      def dwp_run_callback(callback, *args)
+        (self.class.send("#{callback}_callbacks") || []).each do |callback|
+          self.instance_exec(*args, &callback)
+        end
+      end
+
+      def dwp_log(&message_block)
+        @dwp_runner.worker_log(self, &message_block)
+      end
+
+      def dwp_log_exception(exception)
+        dwp_log{ "#{exception.class}: #{exception.message}" }
+        (exception.backtrace || []).each{ |l| dwp_log{ l } }
+      end
 
-    def do_work(work_item)
-      @before_work_callbacks.each{ |p| p.call(self, work_item) }
-      @on_work.call(self, work_item)
-      @after_work_callbacks.each{ |p| p.call(self, work_item) }
     end
 
-    def handle_exception(exception, work_item = nil)
-      @on_error_callbacks.each{ |p| p.call(self, exception, work_item) }
+    module TestHelpers
+
+      def test_runner(worker_class, options = nil)
+        TestRunner.new(worker_class, options)
+      end
+
+      class TestRunner
+        attr_reader :worker_class, :worker
+        attr_reader :queue, :dwp_runner
+
+        def initialize(worker_class, options = nil)
+          @worker_class = worker_class
+
+          @queue = options[:queue] || begin
+            require 'dat-worker-pool/default_queue'
+            DatWorkerPool::DefaultQueue.new
+          end
+
+          @dwp_runner = DatWorkerPool::Runner.new({
+            :num_workers   => MIN_WORKERS,
+            :logger        => options[:logger],
+            :queue         => @queue,
+            :worker_class  => @worker_class,
+            :worker_params => options[:params]
+          })
+
+          @worker = worker_class.new(@dwp_runner, @queue, 1)
+        end
+
+        def run(work_item)
+          self.start
+          self.make_unavailable
+          self.work(work_item)
+          self.make_available
+          self.shutdown
+        end
+
+        def work(work_item)
+          self.worker.instance_eval{ dwp_work(work_item) }
+        end
+
+        def error(exception, work_item = nil)
+          run_callback('on_error', self.worker, exception, work_item)
+        end
+
+        def start
+          run_callback('on_start', self.worker)
+        end
+
+        def shutdown
+          run_callback('on_shutdown', self.worker)
+        end
+
+        def make_unavailable
+          self.worker.instance_eval{ dwp_make_unavailable }
+        end
+
+        def make_available
+          self.worker.instance_eval{ dwp_make_available }
+        end
+
+        private
+
+        def run_callback(callback, worker, *args)
+          self.worker.instance_eval{ dwp_run_callback(callback, *args) }
+        end
+      end
+
     end
 
   end