stream.rb in protocol-http-0.22.9

- old
+ new

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
-
+#
 # Copyright, 2019, by Samuel G. D. Williams. <http://www.codeotaku.com>
 # 
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
@@ -18,19 +18,23 @@
 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
+require_relative 'buffered'
+
 module Protocol
 	module HTTP
 		module Body
 			# The input stream is an IO-like object which contains the raw HTTP POST data. When applicable, its external encoding must be “ASCII-8BIT” and it must be opened in binary mode, for Ruby 1.9 compatibility. The input stream must respond to gets, each, read and rewind.
 			class Stream
-				def initialize(input, output)
+				def initialize(input, output = Buffered.new)
 					@input = input
 					@output = output
 					
+					raise ArgumentError, "Non-writable output!" unless output.respond_to?(:write)
+					
 					# Will hold remaining data in `#read`.
 					@buffer = nil
 					@closed = false
 				end
 				
@@ -43,33 +47,66 @@
 				# read behaves like IO#read. Its signature is read([length, [buffer]]). If given, length must be a non-negative Integer (>= 0) or nil, and buffer must be a String and may not be nil. If length is given and not nil, then this method reads at most length bytes from the input stream. If length is not given or nil, then this method reads all data until EOF. When EOF is reached, this method returns nil if length is given and not nil, or “” if length is not given or is nil. If buffer is given, then the read data will be placed into buffer instead of a newly created String object.
 				# @param length [Integer] the amount of data to read
 				# @param buffer [String] the buffer which will receive the data
 				# @return a buffer containing the data
 				def read(length = nil, buffer = nil)
+					return '' if length == 0
+					
 					buffer ||= Async::IO::Buffer.new
-					buffer.clear
+
+					# Take any previously buffered data and replace it into the given buffer.
+					if @buffer
+						buffer.replace(@buffer)
+						@buffer = nil
+					end
 					
-					until buffer.bytesize == length
-						@buffer = read_next if @buffer.nil?
-						break if @buffer.nil?
+					if length
+						while buffer.bytesize < length and chunk = read_next
+							buffer << chunk
+						end
 						
-						remaining_length = length - buffer.bytesize if length
+						# This ensures the subsequent `slice!` works correctly.
+						buffer.force_encoding(Encoding::BINARY)
+
+						# This will be at least one copy:
+						@buffer = buffer.byteslice(length, buffer.bytesize)
 						
-						if remaining_length && remaining_length < @buffer.bytesize
-							# We know that we are not going to reuse the original buffer.
-							# But byteslice will generate a hidden copy. So let's freeze it first:
-							@buffer.freeze
-							
-							buffer << @buffer.byteslice(0, remaining_length)
-							@buffer = @buffer.byteslice(remaining_length, @buffer.bytesize)
+						# This should be zero-copy:
+						buffer.slice!(length)
+						
+						if buffer.empty?
+							return nil
 						else
-							buffer << @buffer
-							@buffer = nil
+							return buffer
 						end
+					else
+						while chunk = read_next
+							buffer << chunk
+						end
+						
+						return buffer
 					end
+				end
+				
+				# Read at most `length` bytes from the stream. Will avoid reading from the underlying stream if possible.
+				def read_partial(length = nil)
+					if @buffer
+						buffer = @buffer
+						@buffer = nil
+					else
+						buffer = read_next
+					end
 					
-					return nil if buffer.empty? && length && length > 0
+					if buffer and length
+						if buffer.bytesize > length
+							# This ensures the subsequent `slice!` works correctly.
+							buffer.force_encoding(Encoding::BINARY)
+
+							@buffer = buffer.byteslice(length, buffer.bytesize)
+							buffer.slice!(length)
+						end
+					end
 					
 					return buffer
 				end
 				
 				def read_nonblock(length, buffer = nil)
@@ -78,11 +115,11 @@
 					
 					unless @buffer
 						buffer&.clear
 						return
 					end
-
+					
 					if @buffer.bytesize > length
 						chunk = @buffer.byteslice(0, length)
 						@buffer = @buffer.byteslice(length, @buffer.bytesize)
 					else
 						chunk = @buffer
@@ -122,13 +159,13 @@
 					@output&.close
 					@output = nil
 				end
 				
 				# Close the input and output bodies.
-				def close
+				def close(error = nil)
 					self.close_read
 					self.close_write
-					
+
 					return nil
 				ensure
 					@closed = true
 				end
 				