Tidy up implementation.

Author: samuel-williams-shopify

lib/async/promise.rb

@@ -14,7 +14,7 @@ module Async
 	class Promise
 		# Create a new promise.
 		def initialize
-			# nil = pending, true = success, :error = failure:
+			# nil = pending, :completed = success, :failed = failure, :cancelled = cancelled:
 			@resolved = nil
 
 			# Stores either the result value or the exception:
@@ -29,33 +29,55 @@ def initialize
 
 		# @returns [Boolean] Whether the promise has been resolved or rejected.
 		def resolved?
-			@mutex.synchronize { !!@resolved }
+			@mutex.synchronize {!!@resolved}
+		end
+		
+		# @returns [Symbol | Nil] The internal resolved state (:completed, :failed, :cancelled, or nil if pending).
+		# @private For internal use by Task.
+		def resolved
+			@mutex.synchronize {@resolved}
+		end
+		
+		# @returns [Boolean] Whether the promise has been cancelled.
+		def cancelled?
+			@mutex.synchronize {@resolved == :cancelled}
+		end
+		
+		# @returns [Boolean] Whether the promise failed with an exception.
+		def failed?
+			@mutex.synchronize {@resolved == :failed}
+		end
+		
+		# @returns [Boolean] Whether the promise has completed successfully.
+		def completed?
+			@mutex.synchronize {@resolved == :completed}
 		end
 
 		# @returns [Boolean] Whether any fibers are currently waiting for this promise.
 		def waiting?
-			@mutex.synchronize { @waiting > 0 }
+			@mutex.synchronize {@waiting > 0}
 		end
 
 		# Artificially mark that someone is waiting (useful for suppressing warnings).
 		# @private Internal use only.
 		def suppress_warnings!
-			@mutex.synchronize { @waiting += 1 }
+			@mutex.synchronize {@waiting += 1}
 		end
 
 		# Non-blocking access to the current value. Returns nil if not yet resolved.
-		# Does not raise exceptions even if the promise was rejected.
+		# Does not raise exceptions even if the promise was rejected or cancelled.
+		# For resolved promises, returns the raw stored value (result, exception, or cancel exception).
 		#
-		# @returns [Object | Nil] The resolved value, rejected exception, or nil if pending.
+		# @returns [Object | Nil] The stored value, or nil if pending.
 		def value
-			@mutex.synchronize { @resolved ? @value : nil }
+			@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.
 		#
 		# @returns [Object] The resolved value.
-		# @raises [Exception] The rejected exception.
+		# @raises [Exception] The rejected or cancelled exception.
 		def wait
 			@mutex.synchronize do
 				# Increment waiting count:
@@ -66,10 +88,11 @@ def wait
 					@condition.wait(@mutex) unless @resolved
 
 					# Return value or raise exception based on resolution type:
-					if @resolved == :error
-						raise @value
-					else
+					if @resolved == :completed
 						return @value
+					else
+						# Both :failed and :cancelled store exceptions in @value
+						raise @value
 					end
 				ensure
 					# Decrement waiting count when done:
@@ -88,11 +111,13 @@ def resolve(value)
 				return if @resolved
 
 				@value = value
-				@resolved = true
+				@resolved = :completed
 
 				# Wake up all waiting fibers:
 				@condition.broadcast
 			end
+			
+			return value
 		end
 
 		# Reject the promise with an exception.
@@ -105,11 +130,35 @@ def reject(exception)
 				return if @resolved
 
 				@value = exception
-				@resolved = :error
+				@resolved = :failed
 
 				# Wake up all waiting fibers:
 				@condition.broadcast
 			end
+			
+			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!"))
+			@mutex.synchronize do
+				# No-op if already in any final state
+				return if @resolved
+				
+				@value = exception
+				@resolved = :cancelled
+				
+				# Wake up all waiting fibers:
+				@condition.broadcast
+			end
+			
+			return nil
 		end
 
 		# Resolve the promise with the result of the block.
@@ -121,12 +170,17 @@ def fulfill(&block)
 			raise "Promise already resolved!" if @resolved
 
 			begin
-				result = yield
-				resolve(result)
-				return result
+				return self.resolve(yield)
+			rescue Cancel => exception
+				return self.cancel(exception)
 			rescue => error
-				reject(error)
-				raise  # Re-raise so caller knows it failed
+				return self.reject(error)
+			rescue Exception => exception
+				self.reject(exception)
+				raise
+			ensure
+				# Handle non-local exits (throw, etc.) that bypass normal flow:
+				self.resolve(nil) unless @resolved
 			end
 		end
 

lib/async/task.rb

@@ -64,19 +64,18 @@ def initialize(parent = Task.current?, finished: nil, **options, &block)
 
 			# These instance variables are critical to the state of the task.
 			# In the initialized state, the @block should be set, but the @fiber should be nil.
-			# In the running state, the @fiber should be set.
+			# In the running state, the @fiber should be set, and @block should be nil.
 			# In a finished state, the @block should be nil, and the @fiber should be nil.
 			@block = block
 			@fiber = nil
 
-			@status = :initialized
-			@finished = Promise.new
+			@promise = Promise.new
 
 			# Handle finished: parameter for backward compatibility:
 			case finished
 			when false
 				# `finished: false` suppresses warnings for expected task failures:
-				@finished.suppress_warnings!
+				@promise.suppress_warnings!
 			when nil
 				# `finished: nil` is the default, no special handling:
 			else
@@ -121,7 +120,7 @@ def annotation
 
 		# @returns [String] A description of the task and it's current status.
 		def to_s
-			"\#<#{self.description} (#{@status})>"
+			"\#<#{self.description} (#{self.status})>"
 		end
 
 		# @deprecated Prefer {Kernel#sleep} except when compatibility with `stable-v1` is required.
@@ -158,22 +157,22 @@ def finished?
 
 		# @returns [Boolean] Whether the task is running.
 		def running?
-			@status == :running
+			self.alive?
 		end
 
 		# @returns [Boolean] Whether the task failed with an exception.
 		def failed?
-			@status == :failed
+			@promise.failed?
 		end
 
 		# @returns [Boolean] Whether the task has been stopped.
 		def stopped?
-			@status == :stopped
+			@promise.cancelled?
 		end
 
 		# @returns [Boolean] Whether the task has completed execution and generated a result.
 		def completed?
-			@status == :completed
+			@promise.completed?
 		end
 
 		# Alias for {#completed?}.
@@ -182,20 +181,32 @@ def complete?
 		end
 
 		# @attribute [Symbol] The status of the execution of the task, one of `:initialized`, `:running`, `:complete`, `:stopped` or `:failed`.
-		attr :status
+		def status
+			case @promise.resolved
+			when :cancelled
+				:stopped
+			when :failed
+				:failed
+			when :completed
+				:completed
+			when nil
+				self.running? ? :running : :initialized
+			end
+		end
 
 		# Begin the execution of the task.
 		#
 		# @raises [RuntimeError] If the task is already running.
 		def run(*arguments)
-			if @status == :initialized
-				@status = :running
+			# Move from initialized to running by clearing @block
+			if block = @block
+				@block = nil
 
 				schedule do
-					@block.call(self, *arguments)
+					block.call(self, *arguments)
 				rescue => error
 					# I'm not completely happy with this overhead, but the alternative is to not log anything which makes debugging extremely difficult. Maybe we can introduce a debug wrapper which adds extra logging.
-					if !@finished.waiting?
+					unless @promise.waiting?
 						warn(self, "Task may have ended with unhandled exception.", exception: error)
 					end
 
@@ -242,13 +253,23 @@ def wait
 			raise "Cannot wait on own fiber!" if Fiber.current.equal?(@fiber)
 
 			# Wait for the task to complete - Promise handles all the complexity:
-			# It will either return the result or raise the exception automatically
-			@finished.wait
+			begin
+				@promise.wait
+			rescue Promise::Cancel
+				# For backward compatibility, stopped tasks return nil
+				return nil
+			end
 		end
 
 		# Access the result of the task without waiting. May be nil if the task is not completed. Does not raise exceptions.
 		def result
-			@finished.value
+			value = @promise.value
+			# For backward compatibility, return nil for stopped tasks
+			if @promise.cancelled?
+				nil
+			else
+				value
+			end
 		end
 
 		# Stop the task and all of its children.
@@ -386,26 +407,21 @@ def finish!
 
 		# State transition into the completed state.
 		def completed!(result)
-			@status = :completed
-			
 			# Resolve the promise with the result:
-			@finished&.resolve(result)
+			@promise&.resolve(result)
 		end
 
 		# State transition into the failed state.
 		def failed!(exception = false)
-			@status = :failed
-			
 			# Reject the promise with the exception:
-			@finished&.reject(exception)
+			@promise&.reject(exception)
 		end
 
 		def stopped!
 			# Console.info(self, status:) {"Task #{self} was stopped with #{@children&.size.inspect} children!"}
-			@status = :stopped
 
-			# Resolve the promise with nil for stopped tasks:
-			@finished&.resolve(nil)
+			# Cancel the promise:
+			@promise&.cancel
 
 			stopped = false
 
@@ -450,7 +466,7 @@ def schedule(&block)
 
 			@fiber.async_task = self
 
-			self.root.resume(@fiber)
+			(Fiber.scheduler || self.reactor).resume(@fiber)
 		end
 	end
 end