zk 1.4.2 → 1.5.0

data/lib/zk/pool.rb

@@ -91,6 +91,13 @@ module ZK
         end
       end
 
+      # @private
+      def wait_until_closed
+        @mutex.synchronize do
+          @checkin_cond.wait_until { @state == :closed }
+        end
+      end
+
       # yields next available connection to the block
       #
       # raises PoolIsShuttingDownException immediately if close_all! has been
@@ -139,7 +146,7 @@ module ZK
         end
 
         def assert_open!
-          raise Exceptions::PoolIsShuttingDownException unless open? 
+          raise Exceptions::PoolIsShuttingDownException, "pool is shutting down" unless open? 
         end
 
     end # Base
@@ -195,7 +202,7 @@ module ZK
 
           @pool << connection
 
-          @checkin_cond.signal
+          @checkin_cond.broadcast
         end
       end
 
@@ -206,7 +213,8 @@ module ZK
 
       def checkout(blocking=true) 
         raise ArgumentError, "checkout does not take a block, use .with_connection" if block_given?
-        @mutex.synchronize do
+        @mutex.lock
+        begin
           while true
             assert_open!
 
@@ -236,7 +244,9 @@ module ZK
             else
               return false
             end
-          end # while 
+          end # while
+        ensure
+          @mutex.unlock rescue nil
         end
       end
 

data/lib/zk/subscription.rb

@@ -18,9 +18,9 @@ module ZK
 
       def initialize(parent, block)
         raise ArgumentError, "block must repsond_to?(:call)" unless block.respond_to?(:call)
-        @parent = parent
+        @parent   = parent
         @callable = block
-        reopen_after_fork!
+        @mutex    = Monitor.new
       end
 
       def unregistered?
@@ -29,11 +29,20 @@ module ZK
 
       # calls unregister on parent, then sets parent to nil
       def unregister
-        return false unless @parent
-        @parent.unregister(self)
-        @parent = nil
+        obj = nil
+
+        synchronize do
+          return false unless @parent
+          obj, @parent = @parent, nil
+        end
+
+        obj.unregister(self)
+      end
+
+      # an alias for unregister
+      def unsubscribe
+        unregister
       end
-      alias unsubscribe unregister
 
       # @private
       def call(*args)
@@ -50,34 +59,6 @@ module ZK
           @mutex.synchronize { yield }
         end
     end
-
-    module ActorStyle
-      extend Concern
-
-      included do
-        alias_method_chain :unsubscribe, :threaded_callback
-        alias_method_chain :callable, :threaded_callback_wrapper
-        alias_method_chain :reopen_after_fork!, :threaded_refresh
-      end
-
-      def unsubscribe_with_threaded_callback
-        synchronize do
-          @threaded_callback && @threaded_callback.shutdown
-          unsubscribe_without_threaded_callback
-        end
-      end
-
-      def reopen_after_fork_with_threaded_refresh!
-        reopen_after_fork_without_threaded_refresh!
-        @threaded_callback = ThreadedCallback.new(@callable)
-      end
-    
-      def callable_with_threaded_callback_wrapper(*args)
-        synchronize do 
-          @threaded_callback ||= ThreadedCallback.new(@callable) 
-        end
-      end
-    end
   end
 end
 

data/lib/zk/threaded_callback.rb

@@ -5,37 +5,58 @@ module ZK
   # for background processing.
   class ThreadedCallback
     include ZK::Logging
+    include ZK::Exceptions
 
     attr_reader :callback
 
     def initialize(callback=nil, &blk)
       @callback = callback || blk
-      @mutex = Monitor.new
 
-      @mutex.synchronize do
-        @running = true
-        reopen_after_fork!
-      end
+      @state  = :paused
+      reopen_after_fork!
     end
 
     def running?
-      @mutex.synchronize { @running }
+      @mutex.synchronize { @state == :running }
+    end
+
+    # @private
+    def alive?
+      @thread && @thread.alive?
     end
 
     # how long to wait on thread shutdown before we return
-    def shutdown(timeout=2)
-      @mutex.synchronize do
-        @running = false
-        @queue.push(KILL_TOKEN)
-        return unless @thread 
-        unless @thread.join(2)
-          logger.error { "#{self.class} timed out waiting for dispatch thread, callback: #{callback.inspect}" }
-        end
+    def shutdown(timeout=5)
+#       logger.debug { "#{self.class}##{__method__}" }
+
+      @mutex.lock
+      begin
+        return true if @state == :shutdown
+
+        @state = :shutdown
+        @cond.broadcast
+      ensure
+        @mutex.unlock rescue nil
       end
+
+      return true unless @thread 
+
+      unless @thread.join(timeout) == @thread
+        logger.error { "#{self.class} timed out waiting for dispatch thread, callback: #{callback.inspect}" }
+        return false
+      end
+
+      true
     end
 
     def call(*args)
-      @queue.push(args)
+      @mutex.lock
+      begin
+        @array << args
+        @cond.broadcast
+      ensure
+        @mutex.unlock rescue nil
+      end
     end
 
     # called after a fork to replace a dead delivery thread
@@ -44,27 +65,90 @@ module ZK
     #
     # @private
     def reopen_after_fork!
-      return unless @running
-      return if @thread and @thread.alive?
-      @mutex = Monitor.new
-      @queue = Queue.new
-      @thread = spawn_dispatch_thread()
+#       logger.debug { "#{self.class}##{__method__}" }
+
+      unless @state == :paused
+        raise InvalidStateError, "state should have been :paused, not: #{@state.inspect}"
+      end
+
+      if @thread 
+        raise InvalidStateError, "WTF?! did you fork in a callback? my thread was alive!" if @thread.alive?
+        @thread = nil
+      end
+
+      @mutex  = Mutex.new
+      @cond   = ConditionVariable.new
+      @array  = []
+      resume_after_fork_in_parent
+    end
+
+    # shuts down the event delivery thread, but keeps the queue so we can continue
+    # delivering queued events when {#resume_after_fork_in_parent} is called
+    def pause_before_fork_in_parent
+      @mutex.lock
+      begin
+        raise InvalidStateError, "@state was not :running, @state: #{@state.inspect}" if @state != :running
+        return if @state == :paused 
+
+        @state = :paused
+        @cond.broadcast
+      ensure
+        @mutex.unlock rescue nil
+      end
+
+      return unless @thread
+
+      @thread.join
+      @thread = nil
+    end
+
+    def resume_after_fork_in_parent
+      @mutex.lock
+      begin
+        raise InvalidStateError, "@state was not :paused, @state: #{@state.inspect}" if @state != :paused
+        raise InvalidStateError, "@thread was not nil! #{@thread.inspect}" if @thread 
+
+        @state = :running
+#         logger.debug { "#{self.class}##{__method__} spawning dispatch thread" }
+        spawn_dispatch_thread
+      ensure
+        @mutex.unlock rescue nil
+      end
     end
 
     protected
+      # intentionally *not* synchronized
       def spawn_dispatch_thread
-        Thread.new do
-          while running?
-            args = @queue.pop
-            break if args == KILL_TOKEN
-            begin
-              callback.call(*args)
-            rescue Exception => e
-              logger.error { "error caught in handler for path: #{path.inspect}, interests: #{interests.inspect}" }
-              logger.error { e.to_std_format }
+        @thread = Thread.new(&method(:dispatch_thread_body))
+      end
+
+      def dispatch_thread_body
+        Thread.current.abort_on_exception = true
+        while true
+          args = nil
+
+          @mutex.lock
+          begin
+            @cond.wait(@mutex) while @array.empty? and @state == :running
+
+            if @state != :running
+#               logger.warn { "ThreadedCallback, state is #{@state.inspect}, returning" } 
+              return 
             end
+
+            args = @array.shift
+          ensure
+            @mutex.unlock rescue nil
+          end
+            
+          begin
+            callback.call(*args)
+          rescue Exception => e
+            logger.error { e.to_std_format }
           end
         end
+#       ensure
+#         logger.debug { "#{self.class}##{__method__} returning" }
       end
   end
 end

data/lib/zk/threadpool.rb

@@ -2,6 +2,7 @@ module ZK
   # a simple threadpool for running blocks of code off the main thread
   class Threadpool
     include Logging
+    include Exceptions
 
     DEFAULT_SIZE = 5
 
@@ -18,15 +19,28 @@ module ZK
       @size = size || self.class.default_size
 
       @threadpool = []
-      @threadqueue = ::Queue.new
+      @state = :new
+      @queue = []
 
-      @mutex = Monitor.new
+      @mutex = Mutex.new
+      @cond  = ConditionVariable.new
 
       @error_callbacks = []
 
       start!
     end
 
+    # are all of our threads alive?
+    # returns false if there are no running threads
+    def alive?
+      @mutex.lock
+      begin
+        !@threadpool.empty? and @threadpool.all?(&:alive?)
+      ensure
+        @mutex.unlock rescue nil
+      end
+    end
+
     # Queue an operation to be run on an internal threadpool. You may either
     # provide an object that responds_to?(:call) or pass a block. There is no
     # mechanism for retrieving the result of the operation, it is purely
@@ -36,45 +50,103 @@ module ZK
     def defer(callable=nil, &blk)
       callable ||= blk
 
-      # XXX(slyphon): do we care if the threadpool is not running?
-#       raise Exceptions::ThreadpoolIsNotRunningException unless running?
       raise ArgumentError, "Argument to Threadpool#defer must respond_to?(:call)" unless callable.respond_to?(:call)
 
-      @threadqueue << callable
+      @mutex.lock
+      begin
+        @queue << callable
+        @cond.broadcast
+      ensure
+        @mutex.unlock rescue nil
+      end
+
       nil
     end
 
     def running?
-      @mutex.synchronize { @running }
+      @mutex.lock
+      begin
+        @state == :running
+      ensure
+        @mutex.unlock rescue nil
+      end
     end
 
     # returns true if the current thread is one of the threadpool threads
     def on_threadpool?
-      tp = @mutex.synchronize { @threadpool.dup }
-      tp and tp.respond_to?(:include?) and tp.include?(Thread.current)
+      tp = nil
+
+      @mutex.synchronize do
+        return false unless @threadpool # you can't dup nil
+        tp = @threadpool.dup
+      end
+
+      tp.respond_to?(:include?) and tp.include?(Thread.current)
     end
 
     # starts the threadpool if not already running
     def start!
       @mutex.synchronize do
-        return false if @running
-        @running = true
+        return false if @state == :running
+        @state = :running
         spawn_threadpool
       end
+
       true
     end
 
     # like the start! method, but checks for dead threads in the threadpool
     # (which will happen after a fork())
+    #
+    # This will reset the state of the pool and any blocks registered will be
+    # lost
+    #
+    #
     # @private
     def reopen_after_fork!
-      return false unless @running
-      @mutex = Monitor.new
-      @threadqueue = Queue.new
+      # ok, we know that only the child process calls this, right?
+      return false unless (@state == :running) or (@state == :paused)
+      logger.debug { "#{self.class}##{__method__}" }
+
+      @state = :running
+      @mutex = Mutex.new
+      @cond  = ConditionVariable.new
+      @queue = []
       prune_dead_threads
       spawn_threadpool
     end
 
+    # @private
+    def pause_before_fork_in_parent
+      threads = nil
+
+      @mutex.lock
+      begin
+        raise InvalidStateError, "invalid state, expected to be :running, was #{@state.inspect}" if @state != :running
+        return false if @state == :paused
+        @state = :paused
+        @cond.broadcast   # wake threads, let them die
+        threads = @threadpool.slice!(0, @threadpool.length)
+      ensure
+        @mutex.unlock rescue nil
+      end
+
+      join_all(threads)
+      true
+    end
+
+    # @private
+    def resume_after_fork_in_parent
+      @mutex.lock
+      begin
+        raise InvalidStateError, "expected :paused, was #{@state.inspect}" if @state != :paused
+      ensure
+        @mutex.unlock rescue nil
+      end
+
+      start!
+    end
+
     # register a block to be called back with unhandled exceptions that occur
     # in the threadpool. 
     # 
@@ -95,14 +167,27 @@ module ZK
     # the default timeout is 2 seconds per thread
     # 
     def shutdown(timeout=2)
-      @mutex.synchronize do
-        return unless @running
-        @running = false
-        @threadqueue.clear
-        @size.times { @threadqueue << KILL_TOKEN }
+      threads = nil
 
+      @mutex.lock
+      begin
+        return false if @state == :shutdown
+        @state = :shutdown
+
+        @queue.clear
         threads, @threadpool = @threadpool, []
+        @cond.broadcast
+      ensure
+        @mutex.unlock rescue nil
+      end
+
+      join_all(threads)
+
+      nil
+    end
 
+    private
+      def join_all(threads, timeout=nil)
         while th = threads.shift
           begin
             th.join(timeout)
@@ -111,14 +196,8 @@ module ZK
             logger.error { e.to_std_format }
           end
         end
-
-        @threadqueue = ::Queue.new
       end
-
-      nil
-    end
-
-    private
+      
       def dispatch_to_error_handler(e)
         # make a copy that will be free from thread manipulation
         # and doesn't require holding the lock
@@ -150,7 +229,8 @@ module ZK
       end
 
       def prune_dead_threads
-        @mutex.synchronize do
+        @mutex.lock
+        begin
           threads, @threadpool = @threadpool, []
           return if threads.empty?
 
@@ -164,25 +244,41 @@ module ZK
               logger.error { e.to_std_format }
             end
           end
+        ensure
+          @mutex.unlock rescue nil
         end
       end
 
-      def spawn_threadpool #:nodoc:
-        @mutex.synchronize do
-          until @threadpool.size >= @size.to_i
-            thread = Thread.new do
-              while @running
-                begin
-                  op = @threadqueue.pop
-                  break if op == KILL_TOKEN
-                  op.call
-                rescue Exception => e
-                  dispatch_to_error_handler(e)
-                end
-              end
+      def spawn_threadpool
+        until @threadpool.size >= @size.to_i
+          @threadpool << Thread.new(&method(:worker_thread_body))
+        end
+#         logger.debug { "spawn threadpool complete" }
+      end
+
+      def worker_thread_body
+        while true
+          op = nil
+
+          @mutex.lock
+          begin
+            return if @state != :running
+
+            unless op = @queue.shift
+              @cond.wait(@mutex) if @queue.empty? and (@state == :running)
             end
+          ensure
+            @mutex.unlock rescue nil
+          end
+
+          next unless op
 
-            @threadpool << thread
+#           logger.debug { "got #{op.inspect} in thread" }
+
+          begin
+            op.call if op
+          rescue Exception => e
+            dispatch_to_error_handler(e)
           end
         end
       end