Add fetchBytes

Author: prestonvasquez

bson/value_reader.go

@@ -259,26 +259,7 @@ func (vr *valueReader) readBytes(length int32) ([]byte, error) {
 		return nil, fmt.Errorf("invalid length: %d", length)
 	}
 
-	// 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
-	}
-
-	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 data, nil
+	return fetchBytes(vr.src, int(length))
 }
 
 func (vr *valueReader) typeError(t Type) error {
@@ -900,32 +881,25 @@ func (vr *valueReader) readString() (string, error) {
 		return "", fmt.Errorf("invalid string length: %d", length)
 	}
 
-	// Peek the next length bytes without advancing.
-	buf, err := vr.src.peek(int(length))
+	raw, err := fetchBytes(vr.src, int(length))
 	if err != nil {
 		return "", err
 	}
 
 	// Check that the last byte is the NUL terminator.
-	if buf[len(buf)-1] != 0x00 {
-		return "", fmt.Errorf("string does not end with null byte, but with %v", buf[len(buf)-1])
-	}
-
-	// Advance past the length bytes.
-	if _, err := vr.src.discard(int(length)); err != nil {
-		return "", err
+	if raw[len(raw)-1] != 0x00 {
+		return "", fmt.Errorf("string does not end with null byte, but with %v", raw[len(raw)-1])
 	}
 
 	// Convert and strip the trailing NUL.
-	return string(buf[:len(buf)-1]), nil
+	return string(raw[:len(raw)-1]), nil
 }
 
 func (vr *valueReader) peekLength() (int32, error) {
 	buf, err := vr.src.peek(4)
 	if err != nil {
 		return 0, err
 	}
-
 	return int32(binary.LittleEndian.Uint32(buf)), nil
 }
 
@@ -934,57 +908,62 @@ func (vr *valueReader) readLength() (int32, error) {
 }
 
 func (vr *valueReader) readi32() (int32, error) {
-	buf, err := vr.src.peek(4)
+	raw, err := fetchBytes(vr.src, 4)
 	if err != nil {
 		return 0, err
 	}
 
-	// Advance the cursor past the 4 bytes.
-	if _, err := vr.src.discard(4); err != nil {
-		return 0, err
-	}
-
-	return int32(binary.LittleEndian.Uint32(buf)), nil
+	return int32(binary.LittleEndian.Uint32(raw)), nil
 }
 
 func (vr *valueReader) readu32() (uint32, error) {
-	buf, err := vr.src.peek(4)
+	raw, err := fetchBytes(vr.src, 4)
 	if err != nil {
 		return 0, err
 	}
 
-	// Advance the cursor past the 4 bytes.
-	if _, err := vr.src.discard(4); err != nil {
-		return 0, err
-	}
-
-	return binary.LittleEndian.Uint32(buf), nil
+	return binary.LittleEndian.Uint32(raw), nil
 }
 
 func (vr *valueReader) readi64() (int64, error) {
-	buf, err := vr.src.peek(8)
+	raw, err := fetchBytes(vr.src, 8)
 	if err != nil {
 		return 0, err
 	}
 
-	// Advance the cursor past the 8 bytes.
-	if _, err := vr.src.discard(8); err != nil {
-		return 0, err
-	}
-
-	return int64(binary.LittleEndian.Uint64(buf)), nil
+	return int64(binary.LittleEndian.Uint64(raw)), nil
 }
 
 func (vr *valueReader) readu64() (uint64, error) {
-	buf, err := vr.src.peek(8)
+	raw, err := fetchBytes(vr.src, 8)
 	if err != nil {
 		return 0, err
 	}
 
-	// Advance the cursor past the 8 bytes.
-	if _, err := vr.src.discard(8); err != nil {
-		return 0, err
+	return binary.LittleEndian.Uint64(raw), nil
+}
+
+// fetchBytes tries to grab the next n bytes zero-allocation using peek+discard.
+// If peek fails (e.g. bufio buffer full), it falls back to io.ReadFull.
+func fetchBytes(src valueReaderByteSrc, n int) ([]byte, error) {
+	if src.streamable() {
+		data := make([]byte, n)
+		if _, err := io.ReadFull(src, data); err != nil {
+			if errors.Is(err, io.ErrUnexpectedEOF) {
+				err = io.EOF // Convert io.ErrUnexpectedEOF to io.EOF for consistency.
+			}
+			return nil, err
+		}
+
+		return data, nil
+	}
+
+	// Zero-allocation path.
+	buf, err := src.peek(n)
+	if err != nil {
+		return nil, err
 	}
 
-	return binary.LittleEndian.Uint64(buf), nil
+	_, _ = src.discard(n) // Discard the bytes from the source.
+	return buf, nil
 }