Improved streamable error handling.

Author: ioquatix

lib/protocol/http/body/buffered.rb

@@ -56,6 +56,7 @@ def finish
 					self
 				end
 
+				# Ensure that future reads return nil, but allow for rewinding.
 				def close(error = nil)
 					@index = @chunks.length
 				end
@@ -112,8 +113,6 @@ def rewind
 				def inspect
 					if @chunks
 						"\#<#{self.class} #{@chunks.size} chunks, #{self.length} bytes>"
-					else
-						"\#<#{self.class} closed>"
 					end
 				end
 			end

lib/protocol/http/body/streamable.rb

@@ -17,9 +17,14 @@ module Body
 			#
 			# When invoking `call(stream)`, the stream can be read from and written to, and closed. However, the stream is only guaranteed to be open for the duration of the `call(stream)` call. Once the method returns, the stream **should** be closed by the server.
 			module Streamable
+				# Raised when an operation is attempted on a closed stream.
 				class ClosedError < StandardError
 				end
 
+				# Raised when a streaming body is consumed more than once.
+				class ConsumedError < StandardError
+				end
+				
 				def self.new(*arguments)
 					if arguments.size == 1
 						DeferredBody.new(*arguments)
@@ -63,20 +68,36 @@ def write(chunk)
 						end
 					end
 
+					# Indicates that no further output will be generated.
+					def close_write(error = nil)
+						# We might want to specialize the implementation later...
+						close(error)
+					end
+					
 					# Can be invoked by the block to close the stream. Closing the output means that no more chunks will be generated.
 					def close(error = nil)
 						if from = @from
 							# We are closing from within the output fiber, so we need to transfer back to `@from`:
 							@from = nil
-							from.transfer(nil)
+							if error
+								from.raise(error)
+							else
+								from.transfer(nil)
+							end
 						elsif @fiber
 							# We are closing from outside the output fiber, so we need to resume the fiber appropriately:
 							@from = Fiber.current
 
 							if error
+								# The fiber will be resumed from where it last called write, and we will raise the error there:
 								@fiber.raise(error)
 							else
-								@fiber.transfer(nil)
+								begin
+									# If we get here, it means we are closing the fiber from the outside, so we need to transfer control back to the fiber:
+									@fiber.transfer(nil)
+								rescue Protocol::HTTP::Body::Streamable::ClosedError
+									# If the fiber then tries to write to the stream, it will raise a ClosedError, and we will end up here. We can ignore it, as we are already closing the stream and don't care about further writes.
+								end
 							end
 						end
 					end
@@ -105,7 +126,7 @@ def stream?
 					def read
 						if @output.nil?
 							if @block.nil?
-								raise "Streaming body has already been consumed!"
+								raise ConsumedError, "Streaming body has already been consumed!"
 							end
 
 							@output = Output.new(@input, @block)
@@ -120,7 +141,7 @@ def read
 					# The block can read and write to the stream, and must close the stream when finishing.
 					def call(stream)
 						if @block.nil?
-							raise "Streaming block has already been consumed!"
+							raise ConsumedError, "Streaming block has already been consumed!"
 						end
 
 						block = @block
@@ -153,25 +174,17 @@ def close(error = nil)
 				class DeferredBody < Body
 					def initialize(block)
 						super(block, Writable.new)
-						@finishing = false
-					end
-					
-					def close(error = nil)
-						return unless @finishing
-						super
 					end
 
 					# Stream the response body into the block's input.
 					def stream(input)
 						input&.each do |chunk|
 							@input&.write(chunk)
 						end
+						@input&.close_write
 					rescue => error
 						raise
 					ensure
-						@finishing = true
-						@input&.close
-						
 						self.close(error)
 					end
 				end

lib/protocol/http/body/writable.rb

@@ -16,15 +16,9 @@ class Closed < StandardError
 				# @param [Integer] length The length of the response body if known.
 				# @param [Async::Queue] queue Specify a different queue implementation, e.g. `Async::LimitedQueue.new(8)` to enable back-pressure streaming.
 				def initialize(length = nil, queue: Thread::Queue.new)
-					@queue = queue
-					
 					@length = length
-					
+					@queue = queue
 					@count = 0
-					
-					@finished = false
-					
-					@closed = false
 					@error = nil
 				end
 
@@ -34,18 +28,16 @@ def length
 
 				# Stop generating output; cause the next call to write to fail with the given error. Does not prevent existing chunks from being read. In other words, this indicates both that no more data will be or should be written to the body.
 				def close(error = nil)
-					unless @closed
-						@queue.close
-						
-						@closed = true
-						@error = error
-					end
+					@error ||= error
+					
+					@queue.clear
+					@queue.close
 
 					super
 				end
 
 				def closed?
-					@closed
+					@queue.closed?
 				end
 
 				def ready?
@@ -59,23 +51,23 @@ def empty?
 
 				# Read the next available chunk.
 				def read
+					if @error
+						raise @error
+					end
+					
 					@queue.pop
 				end
 
 				# Write a single chunk to the body. Signal completion by calling `#finish`.
 				def write(chunk)
-					# If the reader breaks, the writer will break.
-					if @closed
+					if @queue.closed?
 						raise(@error || Closed)
 					end
 
 					@queue.push(chunk)
 					@count += 1
 				end
 
-				# This alias is provided for compatibility with template generation.
-				alias << write
-				
 				def close_write(error = nil)
 					@error ||= error
 					@queue.close
@@ -95,6 +87,8 @@ def write(chunk)
 						@writable.write(chunk)
 					end
 
+					alias << write
+					
 					def close(error = nil)
 						@closed = true
 

test/protocol/http/body/buffered.rb

@@ -161,6 +161,15 @@
 		end
 	end
 
+	with "#clear" do
+		it "clears all chunks and resets length" do
+			body.clear
+			expect(body.chunks).to be(:empty?)
+			expect(body.read).to be == nil
+			expect(body.length).to be == 0
+		end
+	end
+	
 	with '#inspect' do
 		it "can be inspected" do
 			expect(body.inspect).to be =~ /\d+ chunks, \d+ bytes/

test/protocol/http/body/streamable.rb

@@ -96,6 +96,24 @@
 			expect(stream.string).to be == "HelloWorld"
 		end
 
+		it "will fail if invoked twice" do
+			stream = StringIO.new
+			body.call(stream)
+			
+			expect do
+				body.call(stream)
+			end.to raise_exception(Protocol::HTTP::Body::Streamable::ConsumedError)
+		end
+		
+		it "will fail if trying to read after streaming" do
+			stream = StringIO.new
+			body.call(stream)
+			
+			expect do
+				body.read
+			end.to raise_exception(Protocol::HTTP::Body::Streamable::ConsumedError)
+		end
+		
 		with "a block that raises an error" do
 			let(:block) do
 				proc do |stream|
@@ -125,8 +143,14 @@
 	with '#close' do
 		it "can close the body" do
 			expect(body.read).to be == "Hello"
+			
 			body.close
 		end
+		
+		it "can raise an error on the block" do
+			expect(body.read).to be == "Hello"
+			body.close(RuntimeError.new("Oh no!"))
+		end
 	end
 
 	with "nested fiber" do

test/protocol/http/body/writable.rb

@@ -44,14 +44,15 @@
 		it "should be empty if closed with no pending chunks" do
 			expect(body).not.to be(:empty?)
 
-			body.close
+			body.close_write
 
 			expect(body).to be(:empty?)
 		end
 
 		it "should become empty when pending chunks are read" do
 			body.write("Hello")
-			body.close
+			
+			body.close_write
 
 			expect(body).not.to be(:empty?)
 			body.read
@@ -105,7 +106,7 @@
 				body.write("#{i}")
 			end
 
-			body.close
+			body.close_write
 
 			expect(body.join).to be == "012"
 		end
@@ -117,7 +118,7 @@
 				body.write("Hello World #{i}")
 			end
 
-			body.close
+			body.close_write
 
 			3.times do |i|
 				chunk = body.read
@@ -157,7 +158,7 @@
 
 		it "will stop after finishing" do
 			body.write("Hello World!")
-			body.close
+			body.close_write
 
 			expect(body).not.to be(:empty?)
 
@@ -168,4 +169,31 @@
 			expect(body).to be(:empty?)
 		end
 	end
+	
+	with "#output" do
+		it "can be used to write data" do
+			body.output do |output|
+				output.write("Hello World!")
+			end
+			
+			expect(body.output).to be(:closed?)
+			
+			expect(body.read).to be == "Hello World!"
+			expect(body.read).to be_nil
+		end
+		
+		it "can propagate errors" do
+			expect do
+				body.output do |output|
+					raise "Oops!"
+				end
+			end.to raise_exception(RuntimeError, message: be =~ /Oops/)
+			
+			expect(body).to be(:closed?)
+			
+			expect do
+				body.read
+			end.to raise_exception(RuntimeError, message: be =~ /Oops/)
+		end
+	end
 end