diff --git a/lib/async/task.rb b/lib/async/task.rb index 6629c0e1..0ecc8df0 100644 --- a/lib/async/task.rb +++ b/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. diff --git a/releases.md b/releases.md index ac220bda..ddb8c65e 100644 --- a/releases.md +++ b/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`. diff --git a/test/async/task.rb b/test/async/task.rb index 0a0c3ebe..dbeb48c6 100644 --- a/test/async/task.rb +++ b/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