Add AtExit hook manager

executors' auto_termination feature is migrated

Author: pitr-ch

lib/concurrent.rb

@@ -1,6 +1,7 @@
 require 'concurrent/version'
 
 require 'concurrent/synchronization'
+require 'concurrent/at_exit'
 
 require 'concurrent/configuration'
 

lib/concurrent/at_exit.rb

@@ -0,0 +1,92 @@
+require 'concurrent/logging'
+require 'concurrent/synchronization'
+
+module Concurrent
+
+  # Provides ability to add and remove hooks to be run at `Kernel#at_exit`, order is undefined.
+  # Each hook is executed at most once.
+  class AtExitImplementation < Synchronization::Object
+    include Logging
+
+    def initialize(enabled = true)
+      super()
+      synchronize do
+        @hooks   = {}
+        @enabled = enabled
+      end
+    end
+
+    # Add a hook to be run at `Kernel#at_exit`
+    # @param [Object] hook_id optionally provide an id, if allready present, hook is replaced
+    # @yield the hook
+    # @return id of the hook
+    def add(hook_id = nil, &hook)
+      id = hook_id || hook.object_id
+      synchronize { @hooks[id] = hook }
+      id
+    end
+
+    # Delete a hook by hook_id
+    # @return [true, false]
+    def delete(hook_id)
+      !!synchronize { @hooks.delete hook_id }
+    end
+
+    # Is hook with hook_id rpesent?
+    # @return [true, false]
+    def hook?(hook_id)
+      synchronize { @hooks.key? hook_id }
+    end
+
+    # @return copy of the hooks
+    def hooks
+      synchronize { @hooks }.clone
+    end
+
+    # install `Kernel#at_exit` callback to execute added hooks
+    def install
+      synchronize do
+        @installed ||= begin
+          at_exit { runner }
+          true
+        end
+        self
+      end
+    end
+
+    # Will it run during `Kernel#at_exit`
+    def enabled?
+      synchronize { @enabled }
+    end
+
+    # Configure if it runs during `Kernel#at_exit`
+    def enabled=(value)
+      synchronize { @enabled = value }
+    end
+
+    # run the hooks manually
+    # @return ids of the hooks
+    def run
+      hooks, _ = synchronize { hooks, @hooks = @hooks, {} }
+      hooks.each do |_, hook|
+        begin
+          hook.call
+        rescue => error
+          log ERROR, error
+        end
+      end
+      hooks.keys
+    end
+
+    private
+
+    def runner
+      run if synchronize { @enabled }
+    end
+  end
+
+  private_constant :AtExitImplementation
+
+  # @see AtExitImplementation
+  AtExit = AtExitImplementation.new.install
+end

lib/concurrent/configuration.rb

@@ -1,6 +1,7 @@
 require 'thread'
 require 'concurrent/atomics'
 require 'concurrent/errors'
+require 'concurrent/at_exit'
 require 'concurrent/executors'
 require 'concurrent/utility/processor_count'
 
@@ -39,11 +40,11 @@ def self.global_logger=(value)
     GLOBAL_LOGGER.value = value
   end
 
-  # Disables pool auto-termination which is called on `at_exit` callback.
+  # Disables AtExit hooks including pool auto-termination hooks.
   # When disabled it will be the application
-  # programmer's responsibility to ensure that the thread pools
+  # programmer's responsibility to ensure that the hooks
   # are shutdown properly prior to application exit
-  # by calling {.terminate_pools!} method.
+  # by calling {AtExit.run} method.
   #
   # @note this option should be needed only because of `at_exit` ordering
   #   issues which may arise when running some of the testing frameworks.
@@ -53,20 +54,28 @@ def self.global_logger=(value)
   # @note This method should *never* be called
   #   from within a gem. It should *only* be used from within the main
   #   application and even then it should be used only when necessary.
+  # @see AtExit
+  def self.disable_at_exit_hooks!
+    AtExit.enabled = false
+  end
+
   def self.disable_executor_auto_termination!
-    Executor::AT_EXIT_AUTO_TERMINATION.enabled = false
+    warn '[DEPRECATED] Use Concurrent.disable_at_exit_hooks! instead'
+    disable_at_exit_hooks!
   end
 
   # @return [true,false]
   # @see .disable_executor_auto_termination!
   def self.disable_executor_auto_termination?
-    Executor::AT_EXIT_AUTO_TERMINATION.enabled?
+    warn '[DEPRECATED] Use Concurrent::AtExit.enabled? instead'
+    AtExit.enabled?
   end
 
   # terminates all pools and blocks until they are terminated
   # @see .disable_executor_auto_termination!
   def self.terminate_pools!
-    Executor::AT_EXIT_AUTO_TERMINATION.terminate
+    warn '[DEPRECATED] Use Concurrent::AtExit.run instead'
+    AtExit.run
   end
 
   # Global thread pool optimized for short, fast *operations*.

lib/concurrent/executor/executor.rb

@@ -1,61 +1,11 @@
 require 'concurrent/errors'
 require 'concurrent/logging'
+require 'concurrent/at_exit'
 require 'concurrent/atomic/event'
 
 module Concurrent
 
   module Executor
-    class AtExitAutoTermination
-      def initialize(enabled = true)
-        @mutex     = Mutex.new
-        @executors = Set.new
-        @enabled   = enabled
-        @installed = false
-      end
-
-      def add(executor)
-        @mutex.synchronize { @executors.add executor }
-      end
-
-      def remove(executor)
-        @mutex.synchronize { @executors.delete executor }
-      end
-
-      def enabled?
-        @mutex.synchronize { @enabled }
-      end
-
-      def enabled=(value)
-        @mutex.synchronize { @enabled = value }
-      end
-
-      def terminate
-        executors = @mutex.synchronize { @executors.to_a if @enabled }
-        if executors
-          executors.each(&:kill) # TODO be gentle first
-          executors.each(&:wait_for_termination)
-          executors.clear
-          true
-        else
-          false
-        end
-      end
-
-      def install
-        @installed ||= begin
-          at_exit { terminate }
-          true
-        end
-        self
-      end
-
-      def executors
-        @mutex.synchronize { @executors.to_a }
-      end
-    end
-
-    AT_EXIT_AUTO_TERMINATION = AtExitAutoTermination.new.install
-
     # The policy defining how rejected tasks (tasks received once the
     # queue size reaches the configured `max_queue`, or after the
     # executor has shut down) are handled. Must be one of the values
@@ -126,13 +76,25 @@ def serialized?
     private
 
     def ns_auto_terminate?
-      @auto_terminate
+      !!@auto_terminate
     end
 
     def ns_auto_terminate=(value)
-      AT_EXIT_AUTO_TERMINATION.add self if value == true
-      AT_EXIT_AUTO_TERMINATION.remove self if value == false
-      @auto_terminate = value
+      case value
+      when true
+        AtExit.add(self) { terminate_at_exit }
+        @auto_terminate = true
+      when false
+        AtExit.delete(self)
+        @auto_terminate = false
+      else
+        raise ArgumentError
+      end
+    end
+
+    def terminate_at_exit
+      kill # TODO be gentle first
+      wait_for_termination(10)
     end
   end
 

lib/concurrent/executor/serialized_execution.rb

@@ -98,6 +98,8 @@ def work(job)
         job = @stash.shift || (@being_executed = false)
       end
 
+      # TODO maybe be able to tell caching pool to just enqueue this job, because the current one end at the end
+      # of this block
       call_job job if job
     end
   end

spec/concurrent/executor/ruby_cached_thread_pool_spec.rb

@@ -35,16 +35,16 @@ module Concurrent
       it 'deals with dead threads' do
         expect(subject).to receive(:ns_worker_died).exactly(5).times.and_call_original
 
-        5.times { subject.post { sleep 0.1 } }
-        worker_ids = subject.instance_variable_get(:@pool).map(&:object_id)
-        5.times { subject.post { sleep 0.1; raise Exception } }
-        sleep(0.1)
+        dead_threads_queue = Queue.new
+        5.times { subject.post { sleep 0.1; dead_threads_queue.push Thread.current; raise Exception } }
+        sleep(0.2)
         latch = Concurrent::CountDownLatch.new(5)
         5.times { subject.post { sleep 0.1; latch.count_down } }
-        # processes
         expect(latch.wait(1)).to be true
 
-        expect(worker_ids - subject.instance_variable_get(:@pool).map(&:object_id)).not_to be_empty # some were replaced
+        dead_threads = []
+        dead_threads << dead_threads_queue.pop until dead_threads_queue.empty?
+        expect(dead_threads.all? { |t| !t.alive? }).to be true
       end
     end