Add functionality to be able to peek n bytes off the wire (#72)

maruth-stripe · web-flow · commit 81c6728dd4c2 · 2023-11-10T01:08:16.000+13:00
diff --git a/lib/async/io/stream.rb b/lib/async/io/stream.rb
@@ -125,10 +125,21 @@ def read_until(pattern, offset = 0, chomp: true)
 				return matched
 			end
 			
-			def peek
-				until yield(@read_buffer) or @eof
+			def peek(size = nil)
+				if size
+					until @eof or @read_buffer.bytesize >= size
+						# Compute the amount of data we need to read from the underlying stream:
+						read_size = size - @read_buffer.bytesize
+						
+						# Don't read less than @block_size to avoid lots of small reads:
+						fill_read_buffer(read_size > @block_size ? read_size : @block_size)
+					end
+					return @read_buffer[..([size, @read_buffer.size].min - 1)]
+				end
+				until (block_given? && yield(@read_buffer)) or @eof
 					fill_read_buffer
 				end
+				return @read_buffer
 			end
 			
 			def gets(separator = $/, **options)
diff --git a/spec/async/io/stream_spec.rb b/spec/async/io/stream_spec.rb
@@ -184,6 +184,43 @@
 				expect(subject.read_partial(20)).to be == "o World"
 				expect(subject).to be_eof
 			end
+			
+			it "should peek everything" do
+				io.write "Hello World"
+				io.seek(0)
+				
+				expect(subject.io).to receive(:read_nonblock).and_call_original.twice
+				
+				expect(subject.peek).to be == "Hello World"
+				expect(subject.read).to be == "Hello World"
+				expect(subject).to be_eof
+			end
+		
+			it "should peek only the amount requested" do
+				io.write "Hello World"
+				io.seek(0)
+				
+				expect(subject.io).to receive(:read_nonblock).and_call_original.twice
+				
+				expect(subject.peek(4)).to be == "Hell"
+				expect(subject.read_partial(4)).to be == "Hell"
+				expect(subject).to_not be_eof
+				
+				expect(subject.peek(20)).to be == "o World"
+				expect(subject.read_partial(20)).to be == "o World"
+				expect(subject).to be_eof
+			end
+
+			it "peeks everything when requested bytes is too large" do
+				io.write "Hello World"
+				io.seek(0)
+				
+				expect(subject.io).to receive(:read_nonblock).and_call_original.twice
+				
+				expect(subject.peek(400)).to be == "Hello World"
+				expect(subject.read_partial(400)).to be == "Hello World"
+				expect(subject).to be_eof
+			end
 		
 			context "with large content", if: !Async::IO.buffer? do
 				it "allocates expected amount of bytes" do
