Asynchronous pruning for RubyThreadPoolExecutor

joshuay03 · joshuay03 · commit fd02056db0c8 · 2025-03-17T08:10:47.000+10:00
diff --git a/lib/concurrent-ruby/concurrent/collection/ruby_timeout_queue.rb b/lib/concurrent-ruby/concurrent/collection/ruby_timeout_queue.rb
@@ -0,0 +1,73 @@
+module Concurrent
+  module Collection
+    # @!visibility private
+    # @!macro ruby_timeout_queue
+    class RubyTimeoutQueue < ::Queue
+      def initialize(*args)
+        if RUBY_VERSION >= '3.2'
+          raise "#{self.class.name} is not needed on Ruby 3.2 or later, use ::Queue instead"
+        end
+
+        super(*args)
+
+        @mutex = Mutex.new
+        @cond_var = ConditionVariable.new
+      end
+
+      def push(obj)
+        @mutex.synchronize do
+          super(obj)
+          @cond_var.signal
+        end
+      end
+      alias_method :enq, :push
+      alias_method :<<, :push
+
+      def pop(non_block = false, timeout: nil)
+        if non_block && timeout
+          raise ArgumentError, "can't set a timeout if non_block is enabled"
+        end
+
+        if non_block
+          super(true)
+        elsif timeout
+          @mutex.synchronize do
+            while empty? && timeout.positive? # handle spurious wakeups
+              timed_out, timeout = ns_timed_out?(timeout) { @cond_var.wait(@mutex, timeout) }
+              return nil if timed_out
+            end
+
+            super(true) rescue nil
+          end
+        else
+          super(false)
+        end
+      end
+      alias_method :deq, :pop
+      alias_method :shift, :pop
+
+      private
+
+      def ns_timed_out?(timeout, &cond_var_wait)
+        # https://github.com/ruby/ruby/pull/4256
+        timed_out = if RUBY_VERSION >= '3.1'
+          start_time = Concurrent.monotonic_time
+          result = cond_var_wait.call.nil?
+          end_time = Concurrent.monotonic_time
+          result
+        else
+          start_time = Concurrent.monotonic_time
+          deadline = start_time + timeout
+          cond_var_wait.call
+          end_time = Concurrent.monotonic_time
+          end_time >= deadline
+        end
+
+        elapsed = end_time - start_time
+        new_timeout = [timeout - elapsed, 0].max
+
+        [timed_out, new_timeout]
+      end
+    end
+  end
+end
diff --git a/lib/concurrent-ruby/concurrent/collection/timeout_queue.rb b/lib/concurrent-ruby/concurrent/collection/timeout_queue.rb
@@ -0,0 +1,18 @@
+module Concurrent
+  module Collection
+    # @!visibility private
+    # @!macro internal_implementation_note
+    TimeoutQueueImplementation = if RUBY_VERSION >= '3.2'
+                                   ::Queue
+                                 else
+                                   require 'concurrent/collection/ruby_timeout_queue'
+                                   RubyTimeoutQueue
+                                 end
+    private_constant :TimeoutQueueImplementation
+
+    # @!visibility private
+    # @!macro timeout_queue
+    class TimeoutQueue < TimeoutQueueImplementation
+    end
+  end
+end
diff --git a/lib/concurrent-ruby/concurrent/executor/fixed_thread_pool.rb b/lib/concurrent-ruby/concurrent/executor/fixed_thread_pool.rb
@@ -75,17 +75,6 @@ module Concurrent
   #   @return [Integer] Number of tasks that may be enqueued before reaching `max_queue` and rejecting
   #     new tasks. A value of -1 indicates that the queue may grow without bound.
 
-  # @!macro thread_pool_executor_method_prune_pool
-  #   Prune the thread pool of unneeded threads
-  #
-  #   What is being pruned is controlled by the min_threads and idletime
-  #   parameters passed at pool creation time
-  #
-  #   This is a no-op on some pool implementation (e.g. the Java one).  The Ruby
-  #   pool will auto-prune each time a new job is posted. You will need to call
-  #   this method explicitly in case your application post jobs in bursts (a
-  #   lot of jobs and then nothing for long periods)
-
   # @!macro thread_pool_executor_public_api
   #
   #   @!macro abstract_executor_service_public_api
@@ -122,9 +111,6 @@ module Concurrent
   #
   #   @!method can_overflow?
   #     @!macro executor_service_method_can_overflow_question
-  #
-  #   @!method prune_pool
-  #     @!macro thread_pool_executor_method_prune_pool
 
 
 
diff --git a/lib/concurrent-ruby/concurrent/executor/java_thread_pool_executor.rb b/lib/concurrent-ruby/concurrent/executor/java_thread_pool_executor.rb
@@ -98,10 +98,6 @@ def running?
         super && !@executor.isTerminating
       end
 
-      # @!macro thread_pool_executor_method_prune_pool
-      def prune_pool
-      end
-
       private
 
       def ns_initialize(opts)
diff --git a/lib/concurrent-ruby/concurrent/executor/ruby_thread_pool_executor.rb b/lib/concurrent-ruby/concurrent/executor/ruby_thread_pool_executor.rb
@@ -3,6 +3,7 @@
 require 'concurrent/concern/logging'
 require 'concurrent/executor/ruby_executor_service'
 require 'concurrent/utility/monotonic_time'
+require 'concurrent/collection/timeout_queue'
 
 module Concurrent
 
@@ -95,8 +96,16 @@ def remaining_capacity
     end
 
     # @!visibility private
-    def remove_busy_worker(worker)
-      synchronize { ns_remove_busy_worker worker }
+    def prunable_capacity
+      synchronize { ns_prunable_capacity }
+    end
+
+    # @!visibility private
+    def remove_worker(worker)
+      synchronize do
+        ns_remove_ready_worker worker
+        ns_remove_busy_worker worker
+      end
     end
 
     # @!visibility private
@@ -114,11 +123,6 @@ def worker_task_completed
       synchronize { @completed_task_count += 1 }
     end
 
-    # @!macro thread_pool_executor_method_prune_pool
-    def prune_pool
-      synchronize { ns_prune_pool }
-    end
-
     private
 
     # @!visibility private
@@ -146,9 +150,6 @@ def ns_initialize(opts)
       @largest_length       = 0
       @workers_counter      = 0
       @ruby_pid             = $$ # detects if Ruby has forked
-
-      @gc_interval  = opts.fetch(:gc_interval, @idletime / 2.0).to_i # undocumented
-      @next_gc_time = Concurrent.monotonic_time + @gc_interval
     end
 
     # @!visibility private
@@ -162,12 +163,10 @@ def ns_execute(*args, &task)
 
       if ns_assign_worker(*args, &task) || ns_enqueue(*args, &task)
         @scheduled_task_count += 1
+        nil
       else
-        return fallback_action(*args, &task)
+        fallback_action(*args, &task)
       end
-
-      ns_prune_pool if @next_gc_time < Concurrent.monotonic_time
-      nil
     end
 
     # @!visibility private
@@ -218,7 +217,7 @@ def ns_assign_worker(*args, &task)
     # @!visibility private
     def ns_enqueue(*args, &task)
       return false if @synchronous
-      
+
       if !ns_limited_queue? || @queue.size < @max_queue
         @queue << [task, args]
         true
@@ -265,7 +264,7 @@ def ns_ready_worker(worker, last_message, success = true)
       end
     end
 
-    # removes a worker which is not in not tracked in @ready
+    # removes a worker which is not tracked in @ready
     #
     # @!visibility private
     def ns_remove_busy_worker(worker)
@@ -274,23 +273,19 @@ def ns_remove_busy_worker(worker)
       true
     end
 
-    # try oldest worker if it is idle for enough time, it's returned back at the start
-    #
-    # @!visibility private
-    def ns_prune_pool
-      now = Concurrent.monotonic_time
-      stopped_workers = 0
-      while !@ready.empty? && (@pool.size - stopped_workers > @min_length)
-        worker, last_message = @ready.first
-        if now - last_message > self.idletime
-          stopped_workers += 1
-          @ready.shift
-          worker << :stop
-        else break
-        end
+    def ns_remove_ready_worker(worker)
+      if index = @ready.index { |rw, _| rw == worker }
+        @ready.delete_at(index)
       end
+      true
+    end
 
-      @next_gc_time = Concurrent.monotonic_time + @gc_interval
+    def ns_prunable_capacity
+      if running?
+        [@pool.size - @min_length, @ready.size].min
+      else
+        @pool.size
+      end
     end
 
     def ns_reset_if_forked
@@ -312,7 +307,7 @@ class Worker
 
       def initialize(pool, id)
         # instance variables accessed only under pool's lock so no need to sync here again
-        @queue  = Queue.new
+        @queue  = Collection::TimeoutQueue.new
         @pool   = pool
         @thread = create_worker @queue, pool, pool.idletime
 
@@ -338,17 +333,26 @@ def kill
       def create_worker(queue, pool, idletime)
         Thread.new(queue, pool, idletime) do |my_queue, my_pool, my_idletime|
           catch(:stop) do
-            loop do
+            prunable = true
 
-              case message = my_queue.pop
+            loop do
+              timeout = prunable && my_pool.running? ? my_idletime : nil
+              case message = my_queue.pop(timeout: timeout)
+              when nil
+                if my_pool.prunable_capacity.positive?
+                  my_pool.remove_worker(self)
+                  throw :stop
+                end
+
+                prunable = false
               when :stop
-                my_pool.remove_busy_worker(self)
+                my_pool.remove_worker(self)
                 throw :stop
-
               else
                 task, args = message
                 run_task my_pool, task, args
                 my_pool.ready_worker(self, Concurrent.monotonic_time)
+                prunable = true
               end
             end
           end
diff --git a/spec/concurrent/executor/cached_thread_pool_spec.rb b/spec/concurrent/executor/cached_thread_pool_spec.rb
@@ -152,15 +152,13 @@ module Concurrent
 
         context 'garbage collection' do
 
-          subject { described_class.new(idletime: 0.1, max_threads: 2, gc_interval: 0) }
+          subject { described_class.new(idletime: 0.1, max_threads: 2) }
 
           it 'removes from pool any thread that has been idle too long' do
             latch = Concurrent::CountDownLatch.new(4)
             4.times { subject.post { sleep 0.1; latch.count_down } }
+            sleep 0.4
             expect(latch.wait(1)).to be true
-            sleep 0.2
-            subject.post {}
-            sleep 0.2
             expect(subject.length).to be < 4
           end
 
@@ -197,25 +195,27 @@ module Concurrent
             expect(subject.length).to be >= 5
             3.times { subject << proc { sleep(1) } }
             sleep(0.1)
-            expect(subject.length).to be >= 5
+            expect(subject.length).to be >= 3
           end
         end
       end
 
       context 'stress' do
         configurations = [
-            { min_threads:    2,
-              max_threads:    ThreadPoolExecutor::DEFAULT_MAX_POOL_SIZE,
-              idletime:       0.1, # 1 minute
-              max_queue: 0, # unlimited
+            {
+              min_threads:     2,
+              max_threads:     ThreadPoolExecutor::DEFAULT_MAX_POOL_SIZE,
+              idletime:        60, # 1 minute
+              max_queue:       0, # unlimited
               fallback_policy: :caller_runs, # shouldn't matter -- 0 max queue
-              gc_interval: 0.1 },
-            { min_threads:    2,
-              max_threads:    4,
-              idletime:       0.1, # 1 minute
-              max_queue: 0, # unlimited
+            },
+            {
+              min_threads:     2,
+              max_threads:     4,
+              idletime:        60, # 1 minute
+              max_queue:       0, # unlimited
               fallback_policy: :caller_runs, # shouldn't matter -- 0 max queue
-              gc_interval: 0.1 }
+            }
         ]
 
         configurations.each do |config|
diff --git a/spec/concurrent/executor/java_thread_pool_executor_spec.rb b/spec/concurrent/executor/java_thread_pool_executor_spec.rb
@@ -26,13 +26,6 @@ module Concurrent
 
       it_should_behave_like :thread_pool_executor
 
-      context :prune do
-        it "is a no-op, pruning is handled by the JVM" do
-          executor = JavaThreadPoolExecutor.new
-          executor.prune_pool
-        end
-      end
-
       context '#overload_policy' do
 
         specify ':abort maps to AbortPolicy' do
diff --git a/spec/concurrent/executor/ruby_thread_pool_executor_spec.rb b/spec/concurrent/executor/ruby_thread_pool_executor_spec.rb
