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,90 @@
+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`
+    # @yield the hook
+    # @return id of the hook
+    def add(&hook)
+      synchronize { @hooks[hook.object_id] = hook }
+      self
+    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,24 @@ 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
+        @auto_terminate = AtExit.add { terminate_at_exit }
+      when false
+        AtExit.delete(@auto_terminate)
+        @auto_terminate = nil
+      else
+        raise ArgumentError
+      end
+    end
+
+    def terminate_at_exit
+      kill # TODO be gentle first
+      wait_for_termination
     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