Add Task#wait_cancel_deferred.

lib/async/task.rb

@@ -100,6 +100,7 @@ def initialize(parent = Task.current?, finished: nil, **options, &block)
 			end
 
 			@defer_cancel = nil
+			@cancel_deferred = nil
 
 			# Call this after all state is initialized, as it may call `add_child` which will set the parent and make it visible to the scheduler.
 			super(parent, **options)
@@ -351,6 +352,7 @@ def cancel(later = false, cause: $!)
 				if @defer_cancel == false
 					# Don't cancel now... but update the state so we know we need to cancel later.
 					@defer_cancel = cause
+					@cancel_deferred&.resolve(cause)
 					return false
 				end
 
@@ -407,6 +409,7 @@ def defer_cancel
 
 					# We need to ensure the state is reset before we exit the block:
 					@defer_cancel = nil
+					@cancel_deferred = nil
 
 					# If we were asked to cancel, we should do so now:
 					if defer_cancel
@@ -436,6 +439,34 @@ def stop_deferred?
 			cancel_deferred?
 		end
 
+		# Wait until a cancel has been deferred, either by this task or by the nearest ancestor task deferring cancel. This allows work running within a {defer_cancel} block, including child tasks, to learn that it has been asked to finish.
+		#
+		# If no task is deferring cancel, this waits indefinitely, as any cancel will be delivered directly.
+		#
+		# @returns [Exception] The cause of the deferred cancel.
+		def wait_cancel_deferred
+			node = self
+
+			while node.is_a?(Task)
+				if promise = node.cancel_deferred_promise
+					return promise.wait
+				end
+
+				node = node.parent
+			end
+
+			Promise.new.wait
+		end
+
+		# @returns [Promise | Nil] A promise resolved when cancel is deferred, if this task is deferring cancel.
+		protected def cancel_deferred_promise
+			unless @defer_cancel.nil?
+				@cancel_deferred ||= Promise.new.tap do |promise|
+					promise.resolve(@defer_cancel) if @defer_cancel
+				end
+			end
+		end
+
 		# Lookup the {Task} for the current fiber. Raise `RuntimeError` if none is available.
 		# @returns [Task]
 		# @raises[RuntimeError] If task was not {set!} for the current fiber.

releases.md

@@ -1,5 +1,9 @@
 # Releases
 
+## Unreleased
+
+  - Add `Async::Task#wait_cancel_deferred`, which waits until a cancel has been deferred by the task or its nearest ancestor within `defer_cancel`. This lets work that never finishes on its own, such as a streaming response, learn that it has been asked to finish.
+
 ## v2.46.0
 
   - Fixed `Thread#kill` raising a `TypeError` while shutting down a thread with an active `Async::Reactor`.

test/async/task.rb

@@ -1389,6 +1389,56 @@ def sleep_forever
 		end
 	end
 
+	with "#wait_cancel_deferred" do
+		it "can wait for a deferred cancel from a child task" do
+			cause = nil
+
+			parent_task = reactor.async do |task|
+				task.defer_cancel do
+					task.async do |child_task|
+						cause = child_task.wait_cancel_deferred
+					end.wait
+				end
+			end
+
+			reactor.run_once(0)
+
+			parent_task.cancel
+			expect(parent_task).to be(:running?)
+
+			while parent_task.running?
+				reactor.run_once(0)
+			end
+
+			expect(cause).to be_a(Async::Cancel::Cause)
+			expect(parent_task).to be(:cancelled?)
+		end
+
+		it "returns immediately if cancel was already deferred" do
+			condition = Async::Notification.new
+			cause = nil
+
+			child_task = reactor.async do |task|
+				task.defer_cancel do
+					condition.wait
+					cause = task.wait_cancel_deferred
+				end
+			end
+
+			reactor.run_once(0)
+
+			child_task.cancel
+			condition.signal
+
+			while child_task.running?
+				reactor.run_once(0)
+			end
+
+			expect(cause).to be_a(Async::Cancel::Cause)
+			expect(child_task).to be(:cancelled?)
+		end
+	end
+
 	with "failing task" do
 		include Sus::Fixtures::Console::CapturedLogger