Barrier waits in order of completion/failure.

Author: ioquatix

lib/async/barrier.rb

@@ -5,6 +5,7 @@
 
 require_relative "list"
 require_relative "task"
+require_relative "queue"
 
 module Async
 	# A general purpose synchronisation primitive, which allows one task to wait for a number of other tasks to complete. It can be used in conjunction with {Semaphore}.
@@ -16,6 +17,7 @@ class Barrier
 		# @public Since *Async v1*.
 		def initialize(parent: nil)
 			@tasks = List.new
+			@finished = Queue.new
 
 			@parent = parent
 		end
@@ -41,11 +43,15 @@ def size
 		# Execute a child task and add it to the barrier.
 		# @asynchronous Executes the given block concurrently.
 		def async(*arguments, parent: (@parent or Task.current), **options, &block)
-			task = parent.async(*arguments, **options, &block)
+			waiting = nil
 
-			@tasks.append(TaskNode.new(task))
-			
-			return task
+			parent.async(*arguments, **options) do |task, *arguments|
+				waiting = TaskNode.new(task)
+				@tasks.append(waiting)
+				block.call(task, *arguments)
+			ensure
+				@finished.signal(waiting)
+			end
 		end
 
 		# Whether there are any tasks being held by the barrier.
@@ -55,14 +61,27 @@ def empty?
 		end
 
 		# Wait for all tasks to complete by invoking {Task#wait} on each waiting task, which may raise an error. As long as the task has completed, it will be removed from the barrier.
+		#
+		# @yields {|task| ...} If a block is given, the unwaited task is yielded. You must invoke {Task#wait} yourself. In addition, you may `break` if you have captured enough results.
+		#
 		# @asynchronous Will wait for tasks to finish executing.
 		def wait
-			@tasks.each do |waiting|
+			while !@tasks.empty?
+				# Wait for a task to finish (we get the task node):
+				return unless waiting = @finished.wait
+				
+				# Remove the task as it is now finishing:
+				@tasks.remove?(waiting)
+				
+				# Get the task:
 				task = waiting.task
-				begin
+				
+				# If a block is given, the user can implement their own behaviour:
+				if block_given?
+					yield task
+				else
+					# Wait for it to either complete or raise an error:
 					task.wait
-				ensure
-					@tasks.remove?(waiting) unless task.alive?
 				end
 			end
 		end
@@ -73,6 +92,8 @@ def stop
 			@tasks.each do |waiting|
 				waiting.task.stop
 			end
+			
+			@finished.close
 		end
 	end
 end

lib/async/list.rb

@@ -134,7 +134,7 @@ def remove(node)
 			return removed(node)
 		end
 
-		# @returns [Boolean] Returns true if the list is empty.
+		# @returns [Boolean] True if the list is empty.
 		def empty?
 			@size == 0
 		end

lib/async/notification.rb

@@ -10,12 +10,14 @@ module Async
 	# @public Since *Async v1*.
 	class Notification < Condition
 		# Signal to a given task that it should resume operations.
+		#
+		# @returns [Boolean] if a task was signalled.
 		def signal(value = nil, task: Task.current)
-			return if @waiting.empty?
+			return false if @waiting.empty?
 
 			Fiber.scheduler.push Signal.new(self.exchange, value)
 
-			return nil
+			return true
 		end
 
 		Signal = Struct.new(:waiting, :value) do

lib/async/queue.rb

@@ -20,10 +20,19 @@ class Queue
 		# @parameter available [Notification] The notification to use for signaling when items are available.
 		def initialize(parent: nil, available: Notification.new)
 			@items = []
+			@closed = false
 			@parent = parent
 			@available = available
 		end
 
+		def close
+			@closed = true
+			
+			while @available.waiting?
+				@available.signal(nil)
+			end
+		end
+		
 		# @attribute [Array] The items in the queue.
 		attr :items
 
@@ -59,6 +68,10 @@ def enqueue(*items)
 		# Remove and return the next item from the queue.
 		def dequeue
 			while @items.empty?
+				if @closed
+					return nil
+				end
+				
 				@available.wait
 			end