Make Async::Queue thread safe.

Author: samuel-williams-shopify

lib/async/queue.rb

@@ -10,7 +10,10 @@
 require_relative "notification"
 
 module Async
-	# A queue which allows items to be processed in order.
+	# A thread-safe queue which allows items to be processed in order.
+	#
+	# This implementation uses Thread::Queue internally for thread safety while
+	# maintaining compatibility with the fiber scheduler.
 	#
 	# It has a compatible interface with {Notification} and {Condition}, except that it's multi-value.
 	#
@@ -21,53 +24,42 @@ class Queue
 		class ClosedError < RuntimeError
 		end
 
-		# Create a new queue.
+		# Create a new thread-safe queue.
 		#
 		# @parameter parent [Interface(:async) | Nil] The parent task to use for async operations.
-		# @parameter available [Notification] The notification to use for signaling when items are available.
-		def initialize(parent: nil, available: Notification.new)
-			@items = []
-			@closed = false
+		# @parameter available [Notification] The notification to use for signaling when items are available. (ignored, for compatibility)
+		def initialize(parent: nil, available: nil, queue: Thread::Queue.new)
+			@queue = queue
 			@parent = parent
-			@available = available
 		end
 
 		# @returns [Boolean] Whether the queue is closed.
 		def closed?
-			@closed
+			@queue.closed?
 		end
 
 		# Close the queue, causing all waiting tasks to return `nil`. Any subsequent calls to {enqueue} will raise an exception.
 		def close
-			@closed = true
-			
-			while @available.waiting?
-				@available.signal(nil)
-			end
+			@queue.close
 		end
 
-		# @attribute [Array] The items in the queue.
-		attr :items
-		
 		# @returns [Integer] The number of items in the queue.
 		def size
-			@items.size
+			@queue.size
 		end
 
 		# @returns [Boolean] Whether the queue is empty.
 		def empty?
-			@items.empty?
+			@queue.empty?
 		end
 
 		# Add an item to the queue.
 		def push(item)
-			if @closed
+			if @queue.closed?
 				raise ClosedError, "Cannot push items to a closed queue."
 			end
 
-			@items << item
-			
-			@available.signal unless self.empty?
+			@queue.push(item)
 		end
 
 		# Compatibility with {::Queue#push}.
@@ -77,26 +69,30 @@ def <<(item)
 
 		# Add multiple items to the queue.
 		def enqueue(*items)
-			if @closed
+			if @queue.closed?
 				raise ClosedError, "Cannot enqueue items to a closed queue."
 			end
 
-			@items.concat(items)
-			
-			@available.signal unless self.empty?
+			items.each { |item| @queue.push(item) }
 		end
 
 		# Remove and return the next item from the queue.
 		def dequeue
-			while @items.empty?
-				if @closed
-					return nil
-				end
+			return nil if @queue.closed? && @queue.empty?
+			
+			begin
+				# Try non-blocking first
+				@queue.pop(true)
+			rescue ThreadError
+				# Queue is empty, check if closed
+				return nil if @queue.closed?
 
-				@available.wait
+				# Use blocking pop - the fiber scheduler will handle this properly
+				# in Ruby's fiber scheduler implementation
+				@queue.pop(false)
 			end
-			
-			@items.shift
+		rescue ClosedQueueError
+			nil
 		end
 
 		# Compatibility with {::Queue#pop}.
@@ -136,7 +132,7 @@ def wait
 		end
 	end
 
-	# A queue which limits the number of items that can be enqueued.
+	# A thread-safe queue which limits the number of items that can be enqueued.
 	# @public Since *Async v1*.
 	class LimitedQueue < Queue
 		# @private This exists purely for emitting a warning.
@@ -149,30 +145,19 @@ def self.new(...)
 		# Create a new limited queue.
 		#
 		# @parameter limit [Integer] The maximum number of items that can be enqueued.
-		# @parameter full [Notification] The notification to use for signaling when the queue is full.
-		def initialize(limit = 1, full: Notification.new, **options)
-			super(**options)
-			
-			@limit = limit
-			@full = full
+		# @parameter full [Notification] The notification to use for signaling when the queue is full. (ignored, for compatibility)
+		def initialize(limit = 1, full: nil, **options)
+			super(**options, queue: Thread::SizedQueue.new(limit))
 		end
 
 		# @attribute [Integer] The maximum number of items that can be enqueued.
-		attr :limit
-		
-		# Close the queue, causing all waiting tasks to return `nil`. Any subsequent calls to {enqueue} will raise an exception.
-		# Also signals all tasks waiting for the queue to be full.
-		def close
-			super
-			
-			while @full.waiting?
-				@full.signal(nil)
-			end
+		def limit
+			@queue.max
 		end
 
 		# @returns [Boolean] Whether trying to enqueue an item would block.
 		def limited?
-			!@closed && @items.size >= @limit
+			!@queue.closed? && @queue.size >= @queue.max
 		end
 
 		# Add an item to the queue.
@@ -181,46 +166,24 @@ def limited?
 		#
 		# @parameter item [Object] The item to add to the queue.
 		def push(item)
-			while limited?
-				@full.wait
+			if @queue.closed?
+				raise ClosedError, "Cannot push items to a closed queue."
 			end
 
-			super
+			begin
+				@queue.push(item) # This will block if queue is full
+			rescue ClosedQueueError
+				raise ClosedError, "Cannot push items to a closed queue."
+			end
 		end
 
 		# Add multiple items to the queue.
 		#
-		# If the queue is full, this method will block until there is space available. 
+		# If the queue is full, this method will block until there is space available.
 		#
 		# @parameter items [Array] The items to add to the queue.
 		def enqueue(*items)
-			while !items.empty?
-				while limited?
-					@full.wait
-				end
-				
-				if @closed
-					raise ClosedError, "Cannot enqueue items to a closed queue."
-				end
-				
-				available = @limit - @items.size
-				@items.concat(items.shift(available))
-				
-				@available.signal unless self.empty?
-			end
-		end
-		
-		# Remove and return the next item from the queue.
-		#
-		# If the queue is empty, this method will block until an item is available.
-		#
-		# @returns [Object] The next item in the queue.
-		def dequeue
-			item = super
-			
-			@full.signal
-			
-			return item
+			items.each { |item| push(item) }
 		end
 	end
 end