Avoid peeking on unbuffered stream

Author: prestonvasquez

bson/buffered_value_reader.go

@@ -21,6 +21,17 @@ type bufferedValueReader struct {
 
 var _ valueReaderByteSrc = (*bufferedValueReader)(nil)
 
+// Read reads up to len(p) bytes from the in-memory buffer, advancing the offset
+// by the number of bytes read.
+func (b *bufferedValueReader) Read(p []byte) (int, error) {
+	if b.offset >= int64(len(b.buf)) {
+		return 0, io.EOF
+	}
+	n := copy(p, b.buf[b.offset:])
+	b.offset += int64(n)
+	return n, nil
+}
+
 // ReadByte returns the single byte at buf[offset] and advances offset by 1.
 func (b *bufferedValueReader) ReadByte() (byte, error) {
 	if b.offset >= int64(len(b.buf)) {

bson/streaming_value_reader.go

@@ -23,6 +23,14 @@ type streamingValueReader struct {
 
 var _ valueReaderByteSrc = (*streamingValueReader)(nil)
 
+// Read reads up to len(p) bytes from the underlying bufio.Reader, advancing
+// the offset by the number of bytes read.
+func (s *streamingValueReader) Read(p []byte) (int, error) {
+	n, err := s.br.Read(p)
+	s.offset += int64(n)
+	return n, err
+}
+
 // ReadByte returns the single byte at buf[offset] and advances offset by 1.
 func (s *streamingValueReader) ReadByte() (byte, error) {
 	c, err := s.br.ReadByte()

bson/value_reader.go

@@ -35,6 +35,7 @@ var vrPool = sync.Pool{
 }
 
 type valueReaderByteSrc interface {
+	io.Reader
 	io.ByteReader
 
 	// Peek returns the next n bytes without advancing the cursor. It must return
@@ -258,18 +259,26 @@ func (vr *valueReader) readBytes(length int32) ([]byte, error) {
 		return nil, fmt.Errorf("invalid length: %d", length)
 	}
 
-	// Peek the next length bytes.
-	buf, err := vr.src.peek(int(length))
-	if err != nil {
-		return nil, err
+	// If we can peek and discard the bytes, we can avoid an allocation.
+	if buf, err := vr.src.peek(int(length)); err == nil {
+		_, _ = vr.src.discard(int(length)) // Discard the bytes from the source.
+		return buf, nil
 	}
 
-	// Advance the cursor past those bytes.
-	if _, err := vr.src.discard(int(length)); err != nil {
+	data := make([]byte, length)
+	reader, ok := vr.src.(io.Reader)
+	if !ok {
+		return nil, fmt.Errorf("source does not implement io.Reader: %T", vr.src)
+	}
+
+	if _, err := io.ReadFull(reader, data); err != nil {
+		if errors.Is(err, io.ErrUnexpectedEOF) {
+			err = io.EOF // Convert io.ErrUnexpectedEOF to io.EOF for consistency.
+		}
 		return nil, err
 	}
 
-	return buf, nil
+	return data, nil
 }
 
 func (vr *valueReader) typeError(t Type) error {