Unify cancellation handling between Promise and Task.

Author: ioquatix

lib/async/promise.rb

@@ -6,6 +6,7 @@
 
 require_relative "error"
 require_relative "deadline"
+require_relative "cancel"
 
 module Async
 	# A promise represents a value that will be available in the future.
@@ -77,42 +78,6 @@ def value
 			@mutex.synchronize{@resolved ? @value : nil}
 		end
 
-		# Wait for the promise to be resolved and return the value.
-		# If already resolved, returns immediately. If rejected, raises the stored exception.
-		#
-		# @parameter timeout [Numeric | Nil] Maximum time to wait. If nil, waits indefinitely. If 0, raises immediately if not resolved.
-		# @returns [Object] The resolved value.
-		# @raises [Exception] The rejected or cancelled exception.
-		# @raises [Async::TimeoutError] If timeout expires before the promise is resolved.
-		def wait(timeout: nil)
-			@mutex.synchronize do
-				# Increment waiting count:
-				@waiting += 1
-				
-				begin
-					# Wait for resolution if not already resolved:
-					unless @resolved
-						if timeout.nil?
-							wait_indefinitely
-						else
-							wait_with_timeout(timeout)
-						end
-					end
-					
-					# Return value or raise exception based on resolution type:
-					if @resolved == :completed
-						return @value
-					else
-						# Both :failed and :cancelled store exceptions in @value
-						raise @value
-					end
-				ensure
-					# Decrement waiting count when done:
-					@waiting -= 1
-				end
-			end
-		end
-		
 		# Wait indefinitely for the promise to be resolved.
 		private def wait_indefinitely
 			until @resolved
@@ -122,14 +87,14 @@ def wait(timeout: nil)
 
 		# Wait for the promise to be resolved, respecting the deadline timeout.
 		# @parameter timeout [Numeric] The timeout duration.
-		# @raises [Async::TimeoutError] If the timeout expires before resolution.
+		# @returns [Boolean] True if resolved, false if timeout expires.
 		private def wait_with_timeout(timeout)
 			# Create deadline for timeout tracking:
 			deadline = Deadline.start(timeout)
 
 			# Handle immediate timeout (non-blocking):
 			if deadline == Deadline::Zero && !@resolved
-				raise Async::TimeoutError, "Promise wait not resolved!"
+				return false
 			end
 
 			# Wait with deadline tracking:
@@ -139,11 +104,67 @@ def wait(timeout: nil)
 
 				# Check if deadline has expired before waiting:
 				if remaining <= 0
-					raise Async::TimeoutError, "Promise wait timed out!"
+					return false
 				end
 
 				@condition.wait(@mutex, remaining)
 			end
+			
+			return true
+		end
+		
+		# Wait for the promise to be resolved (without raising exceptions).
+		#
+		# If already resolved, returns immediately. Otherwise, waits until resolution or timeout.
+		#
+		# @parameter timeout [Numeric | Nil] Maximum time to wait. If nil, waits indefinitely. If 0, returns immediately if not resolved.
+		# @returns [Boolean] True if the promise is resolved, false if timeout expires
+		def wait?(timeout: nil)
+			unless @resolved
+				@mutex.synchronize do
+					# Increment waiting count:
+					@waiting += 1
+					
+					begin
+						# Wait for resolution if not already resolved:
+						unless @resolved
+							if timeout.nil?
+								wait_indefinitely
+							else
+								unless wait_with_timeout(timeout)
+									# We don't want to race on @resolved after exiting the mutex:
+									return nil
+								end
+							end
+						end
+					ensure
+						# Decrement waiting count when done:
+						@waiting -= 1
+					end
+				end
+			end
+			
+			return @resolved
+		end
+		
+		# Wait for the promise to be resolved and return the value.
+		#
+		# If already resolved, returns immediately. If rejected, raises the stored exception.
+		#
+		# @returns [Object] The resolved value.
+		# @raises [Exception] The rejected or cancelled exception.
+		# @raises [Async::TimeoutError] If timeout expires before the promise is resolved.
+		def wait(...)
+			resolved = wait?(...)
+			
+			if resolved.nil?
+				raise TimeoutError, "Timeout while waiting for promise!"
+			elsif resolved == :completed
+				return @value
+			elsif @value
+				# If we aren't completed, we should have an exception or cancel reason stored:
+				raise @value
+			end
 		end
 
 		# Resolve the promise with a value.
@@ -155,8 +176,8 @@ def resolve(value)
 			@mutex.synchronize do
 				return if @resolved
 
-				@value = value
 				@resolved = :completed
+				@value = value
 
 				# Wake up all waiting fibers:
 				@condition.broadcast
@@ -174,8 +195,8 @@ def reject(exception)
 			@mutex.synchronize do
 				return if @resolved
 
-				@value = exception
 				@resolved = :failed
+				@value = exception
 
 				# Wake up all waiting fibers:
 				@condition.broadcast
@@ -184,20 +205,16 @@ def reject(exception)
 			return nil
 		end
 
-		# Exception used to indicate cancellation.
-		class Cancel < Exception
-		end
-		
 		# Cancel the promise, indicating cancellation.
 		# All current and future waiters will receive nil.
 		# Can only be called on pending promises - no-op if already resolved.
-		def cancel(exception = Cancel.new("Promise was cancelled!"))
+		def cancel(exception = Cancel.new("Promise cancelled!"))
 			@mutex.synchronize do
 				# No-op if already in any final state
 				return if @resolved
 
-				@value = exception
 				@resolved = :cancelled
+				@value = exception
 
 				# Wake up all waiting fibers:
 				@condition.broadcast

lib/async/task.rb

@@ -270,13 +270,8 @@ def async(*arguments, **options, &block)
 		def wait
 			raise "Cannot wait on own fiber!" if Fiber.current.equal?(@fiber)
 
-			# Wait for the task to complete - Promise handles all the complexity:
-			begin
-				@promise.wait
-			rescue Promise::Cancel
-				# For backward compatibility, cancelled tasks return nil:
-				return nil
-			end
+			# Wait for the task to complete:
+			@promise.wait
 		end
 
 		# For compatibility with `Thread#join` and similar interfaces.
@@ -483,8 +478,8 @@ def failed!(exception = false)
 		def cancelled!
 			# Console.info(self, status:) {"Task #{self} was cancelled with #{@children&.size.inspect} children!"}
 
-			# Cancel the promise:
-			@promise.cancel
+			# Cancel the promise, specify nil here so that no exception is raised when waiting on the promise:
+			@promise.cancel(nil)
 
 			cancelled = false
 

test/async/promise.rb

@@ -201,7 +201,7 @@
 
 			expect do
 				promise.wait(timeout: 0)
-			end.to raise_exception(Async::TimeoutError, message: be =~ /not resolved/)
+			end.to raise_exception(Async::TimeoutError, message: be =~ /Timeout/i)
 		end
 
 		it "raises TimeoutError after timeout expires" do
@@ -223,7 +223,7 @@
 			elapsed = Time.now - start_time
 
 			expect(error).to be_a(Async::TimeoutError)
-			expect(error.message).to be =~ /timed out/
+			expect(error.message).to be =~ /Timeout/i
 			expect(elapsed).to be >= 0.1
 			expect(elapsed).to be < 0.5
 			expect(promise).not.to be(:waiting?)
@@ -333,11 +333,11 @@
 			expect(promise.cancelled?).to be == true
 			expect(promise.completed?).to be == false
 			expect(promise.failed?).to be == false
-			expect(promise.value).to be_a(Async::Promise::Cancel)
+			expect(promise.value).to be_a(Async::Cancel)
 
 			expect do
 				promise.wait
-			end.to raise_exception(Async::Promise::Cancel, message: be =~ /cancelled/)
+			end.to raise_exception(Async::Cancel, message: be =~ /cancelled/)
 		end
 
 		it "can be cancelled with custom exception" do
@@ -368,7 +368,7 @@
 			expect(promise.cancelled?).to be == true
 			expect do
 				promise.wait
-			end.to raise_exception(Async::Promise::Cancel)
+			end.to raise_exception(Async::Cancel)
 		end
 
 		it "ignores reject after cancel" do
@@ -378,7 +378,7 @@
 			expect(promise.cancelled?).to be == true
 			expect do
 				promise.wait
-			end.to raise_exception(Async::Promise::Cancel)
+			end.to raise_exception(Async::Cancel)
 		end
 
 		it "handles multiple concurrent waiters with cancellation" do
@@ -524,7 +524,7 @@
 		end
 
 		it "handles Cancel exceptions (absorbed, not re-raised)" do
-			cancel_exception = Async::Promise::Cancel.new("user cancelled")
+			cancel_exception = Async::Cancel.new("user cancelled")
 
 			result = promise.fulfill do
 				raise cancel_exception
@@ -538,14 +538,14 @@
 			# But promise.wait will raise the cancel exception:
 			expect do
 				promise.wait
-			end.to raise_exception(Async::Promise::Cancel, message: be =~ /user cancelled/)
+			end.to raise_exception(Async::Cancel, message: be =~ /user cancelled/)
 		end
 
 		it "handles custom Cancel exceptions" do
 			custom_cancel = StandardError.new("custom stop")
 
 			result = promise.fulfill do
-				raise Async::Promise::Cancel.new("wrapper").tap{|c| c.instance_variable_set(:@cause, custom_cancel)}
+				raise Async::Cancel.new("wrapper").tap{|c| c.instance_variable_set(:@cause, custom_cancel)}
 			end
 
 			expect(result).to be_nil
@@ -646,7 +646,7 @@
 			promise3 = Async::Promise.new
 
 			# Cancel exception should be caught specifically:
-			promise1.fulfill{raise Async::Promise::Cancel.new}
+			promise1.fulfill{raise Async::Cancel.new}
 			expect(promise1.cancelled?).to be == true
 
 			# StandardError should be caught by rescue =>:
@@ -736,7 +736,7 @@
 			expect(promise.value).to be == original_value
 			expect do
 				promise.wait
-			end.to raise_exception(Async::Promise::Cancel)
+			end.to raise_exception(Async::Cancel)
 		end
 
 		it "reject after cancel is ignored" do
@@ -749,7 +749,7 @@
 			expect(promise.value).to be == original_value
 			expect do
 				promise.wait
-			end.to raise_exception(Async::Promise::Cancel)
+			end.to raise_exception(Async::Cancel)
 		end
 
 		it "fulfill after cancel raises already resolved error" do

test/async/task.rb

@@ -358,9 +358,9 @@
 		end
 
 		it "can stop nested tasks with exception handling" do
+			parent = child = nil
+			
 			reactor.run do
-				child = nil
-				
 				parent = reactor.async do |task|
 					child = task.async do |subtask|
 						# Never comes back:
@@ -370,20 +370,20 @@
 					begin
 						child.wait
 					ensure
-						child.stop
+						child.cancel
 					end
 				end
 
 				# Ensure the parent has a chance to run:
 				Fiber.scheduler.yield
 
-				parent.stop
+				parent.cancel
 				parent.wait
-				expect(parent).to be(:cancelled?)
-				
 				child.wait
-				expect(child).to be(:cancelled?)
 			end
+			
+			expect(parent).to be(:cancelled?)
+			expect(child).to be(:cancelled?)
 		end
 
 		it "can stop current task" do