Fix policy-driven container shutdown

Assisted-By: devx/3236e566-7538-432e-a30a-2bdf37265ed4

Author: samuel-williams-shopify

lib/async/container/generic.rb

@@ -173,7 +173,7 @@ def stop(timeout = true)
 				@stopping = true
 				@group.stop(timeout)
 
-				if @group.running?
+				if @group.running?(except: current_task)
 					Console.warn(self, "Group is still running after stopping it!")
 				else
 					Console.info(self, "Group has stopped.")
@@ -183,6 +183,12 @@ def stop(timeout = true)
 				raise
 			end
 
+			private def current_task
+				Async::Task.current
+			rescue RuntimeError
+				nil
+			end
+			
 			protected def health_check_failed(child, age_clock, health_check_timeout)
 				begin
 					@policy.health_check_failed(

lib/async/container/group.rb

@@ -38,8 +38,9 @@ def initialize(health_check_interval: 1.0)
 
 				@mutex = Mutex.new
 				@children = {}
-				@supervisors = 0
-				@events = ::Thread::Queue.new
+				@supervisors = {}
+				@pending_events = 0
+				@waiters = []
 			end
 
 			# @attribute [Numeric | Nil] The interval used to wake waiters for periodic health checks.
@@ -58,9 +59,10 @@ def size
 			end
 
 			# Check whether any supervisor tasks are still running.
+			# @parameter except [Async::Task | Nil] The supervisor task to ignore.
 			# @returns [Boolean] Whether any supervisor tasks are running.
-			def running?
-				@mutex.synchronize{@supervisors.positive?}
+			def running?(except: nil)
+				running_supervisors?(except: except)
 			end
 
 			# Check whether any supervisor tasks are still running.
@@ -85,13 +87,16 @@ def running
 			# @yields {...} The supervisor block to execute.
 			# @returns [Async::Task] The supervisor task.
 			def supervise(&block)
-				@mutex.synchronize{@supervisors += 1}
+				parent = Async::Task.current
 
-				Async::Task.current.async(transient: true) do
+				parent.async(transient: true) do
+					task = Async::Task.current
+					@mutex.synchronize{@supervisors[task] = true}
+					
 					begin
 						block.call
 					ensure
-						@mutex.synchronize{@supervisors -= 1}
+						@mutex.synchronize{@supervisors.delete(task)}
 						signal!
 					end
 				end
@@ -116,15 +121,29 @@ def delete(child)
 			# @parameter duration [Numeric | Nil] The maximum duration to sleep.
 			# @returns [Object | Nil] The queued signal value, or `nil` if the sleep timed out.
 			def sleep(duration = nil)
+				events = ::Thread::Queue.new
+				
+				@mutex.synchronize do
+					if @pending_events.positive?
+						@pending_events -= 1
+						return true
+					end
+					
+					@waiters << events
+				end
+				
 				::Thread.handle_interrupt(SignalException => :immediate) do
-					@events.pop(timeout: duration)
+					events.pop(timeout: duration)
 				end
+			ensure
+				@mutex.synchronize{@waiters.delete(events)} if events
 			end
 
-			# Wait until all supervisor tasks have stopped.
+			# Wait until all other supervisor tasks have stopped.
+			# @parameter except [Async::Task | Nil] The supervisor task to ignore while waiting.
 			# @returns [Nil]
-			def wait
-				sleep while running?
+			def wait(except: current_supervisor)
+				sleep while running_supervisors?(except: except)
 			end
 
 			# Wake any waiters so they can re-check child health or state.
@@ -159,21 +178,26 @@ def kill
 			# @returns [Nil]
 			def stop(graceful = GRACEFUL_TIMEOUT)
 				Console.debug(self, "Stopping all children...", graceful: graceful)
+				except = current_supervisor
 
 				if graceful
 					interrupt
 
 					graceful = DEFAULT_GRACEFUL_TIMEOUT if graceful == true
-					wait_for_exit(Clock.start, graceful)
+					wait_for_children(Clock.start, graceful)
 				end
 			ensure
-				if running?
+				if size.positive?
 					if graceful
 						Console.warn(self, "Killing children after graceful shutdown failed...", size: size)
 					end
 
 					kill
-					wait
+					sleep while size.positive?
+				end
+				
+				if running_supervisors?(except: except)
+					wait(except: except)
 				end
 			end
 
@@ -189,8 +213,8 @@ def each_child
 				end
 			end
 
-			def wait_for_exit(clock, timeout)
-				while running?
+			def wait_for_children(clock, timeout)
+				while size.positive?
 					duration = timeout - clock.total
 
 					break if duration.negative?
@@ -199,8 +223,40 @@ def wait_for_exit(clock, timeout)
 				end
 			end
 
+			def current_supervisor
+				task = Async::Task.current
+				
+				@mutex.synchronize do
+					@supervisors.key?(task) ? task : nil
+				end
+			rescue RuntimeError
+				nil
+			end
+			
+			def running_supervisors?(except: nil)
+				@mutex.synchronize do
+					if except
+						@supervisors.any?{|task, _| task != except}
+					else
+						@supervisors.any?
+					end
+				end
+			end
+			
 			def signal!
-				@events << true
+				waiters = @mutex.synchronize do
+					if @waiters.empty?
+						@pending_events += 1
+					end
+					
+					@waiters.dup
+				end
+				
+				waiters.each do |events|
+					events << true
+				end
+				
+				return true
 			end
 		end
 	end

test/async/container/policy.rb

@@ -237,6 +237,37 @@ def child_exit(container, child, status, name:, key:, **options)
 			expect(container.statistics.failures).to be == 1
 		end
 
+		it "can stop the container from child_exit" do
+			stop_policy = Class.new(Async::Container::Policy) do
+				def initialize
+					@stop_count = 0
+				end
+				
+				attr :stop_count
+				
+				def child_exit(container, child, status, name:, key:, **options)
+					unless container.stopping?
+						@stop_count += 1
+						container.stop
+					end
+				end
+			end.new
+			
+			container = Async::Container.best_container_class.new(policy: stop_policy)
+			
+			3.times do |i|
+				container.spawn(name: "worker-#{i}") do |instance|
+					instance.ready!
+					exit(1)
+				end
+			end
+			
+			container.wait
+			
+			expect(stop_policy.stop_count).to be == 1
+			expect(container).not.to be(:running?)
+		end
+		
 		it "invokes callbacks for multiple children" do
 			container = Async::Container.best_container_class.new(policy: tracking_policy)