diff --git a/lib/async/io/stream.rb b/lib/async/io/stream.rb index 914e709..d1f27a1 100644 --- 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 index 812f319..4dd042c 100644 --- 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