Backport Task#defer_stop.

Author: ioquatix

lib/async/task.rb

@@ -84,6 +84,8 @@ def initialize(reactor, parent = Task.current?, logger: nil, finished: nil, **op
 			@logger = logger || @parent.logger
 
 			@fiber = make_fiber(&block)
+			
+			@defer_stop = nil
 		end
 
 		attr :logger
@@ -162,6 +164,13 @@ def stop(later = false)
 				return
 			end
 
+			# If we are deferring stop...
+			if @defer_stop == false
+				# Don't stop now... but update the state so we know we need to stop later.
+				@defer_stop = true
+				return false
+			end
+			
 			if self.running?
 				if self.current?
 					if later
@@ -182,6 +191,41 @@ def stop(later = false)
 			end
 		end
 
+		# Defer the handling of stop. During the execution of the given block, if a stop is requested, it will be deferred until the block exits. This is useful for ensuring graceful shutdown of servers and other long-running tasks. You should wrap the response handling code in a defer_stop block to ensure that the task is stopped when the response is complete but not before.
+		#
+		# You can nest calls to defer_stop, but the stop will only be deferred until the outermost block exits.
+		#
+		# If stop is invoked a second time, it will be immediately executed.
+		#
+		# @yields {} The block of code to execute.
+		def defer_stop
+			# Tri-state variable for controlling stop:
+			# - nil: defer_stop has not been called.
+			# - false: defer_stop has been called and we are not stopping.
+			# - true: defer_stop has been called and we will stop when exiting the block.
+			if @defer_stop.nil?
+				# If we are not deferring stop already, we can defer it now:
+				@defer_stop = false
+				
+				begin
+					yield
+				rescue Stop
+					# If we are exiting due to a stop, we shouldn't try to invoke stop again:
+					@defer_stop = nil
+					raise
+				ensure
+					# If we were asked to stop, we should do so now:
+					if @defer_stop
+						@defer_stop = nil
+						self.stop
+					end
+				end
+			else
+				# If we are deferring stop already, entering it again is a no-op.
+				yield
+			end
+		end
+		
 		# Lookup the {Task} for the current fiber. Raise `RuntimeError` if none is available.
 		# @return [Async::Task]
 		# @raise [RuntimeError] if task was not {set!} for the current fiber.

spec/async/task_spec.rb

@@ -22,6 +22,7 @@
 
 require 'async'
 require 'async/clock'
+require 'async/notification'
 
 RSpec.describe Async::Task do
 	let(:reactor) {Async::Reactor.new}
@@ -496,4 +497,62 @@ def sleep_forever
 			expect(apples_task.to_s).to include "complete"
 		end
 	end
+	
+	describe '#defer_stop' do
+		it "can defer stopping a task" do
+			child_task = reactor.async do |task|
+				task.defer_stop do
+					task.sleep(10)
+				end
+			end
+			
+			reactor.run_once(0)
+			
+			child_task.stop
+			expect(child_task).to be_running
+			
+			child_task.stop
+			expect(child_task).to be_stopped
+		end
+		
+		it "will stop the task if it was deferred" do
+			condition = Async::Notification.new
+			
+			child_task = reactor.async do |task|
+				task.defer_stop do
+					condition.wait
+				end
+			end
+			
+			reactor.run_once(0)
+			
+			child_task.stop(true)
+			expect(child_task).to be_running
+			
+			reactor.async do
+				condition.signal
+			end
+			
+			reactor.run_once(0)
+			expect(child_task).to be_stopped
+		end
+		
+		it "can defer stop in a deferred stop" do
+			child_task = reactor.async do |task|
+				task.defer_stop do
+					task.defer_stop do
+						task.sleep(10)
+					end
+				end
+			end
+			
+			reactor.run_once(0)
+			
+			child_task.stop
+			expect(child_task).to be_running
+			
+			child_task.stop
+			expect(child_task).to be_stopped
+		end
+	end
 end