Add streaming support

Author: prestonvasquez

bson/buffered_value_reader.go

@@ -21,13 +21,6 @@ type bufferedValueReader struct {
 
 var _ valueReaderByteSrc = (*bufferedValueReader)(nil)
 
-// Read reads up to len(p) bytes from buf[offset:] into p, advancing offset.
-//
-// TODO: Depending on streaming requirements, we might be able to remove this.
-func (b *bufferedValueReader) Read([]byte) (n int, err error) {
-	panic("Read not implemented for bufferedValueReader")
-}
-
 // 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)) {
@@ -109,3 +102,11 @@ func (b *bufferedValueReader) regexLength() (int32, error) {
 	// (options) + 2 null terminators
 	return int32(i + j + 2), nil
 }
+
+func (*bufferedValueReader) streamable() bool {
+	return false
+}
+
+func (*bufferedValueReader) reset() {
+	// No resources to release for bufferedValueReader.
+}

bson/copier.go

@@ -180,7 +180,7 @@ func copyValueFromBytes(dst ValueWriter, t Type, src []byte) error {
 		return wvb.writeValueBytes(t, src)
 	}
 
-	vr := newDocumentReader(bytes.NewReader(src))
+	vr := newBufferedDocumentReader(bytes.NewReader(src))
 	vr.pushElement(t)
 
 	return copyValue(dst, vr)

bson/copier_test.go

@@ -40,7 +40,7 @@ func TestCopier(t *testing.T) {
 			doc = bsoncore.AppendStringElement(doc, "Hello", "world")
 			doc, err := bsoncore.AppendDocumentEnd(doc, idx)
 			noerr(t, err)
-			src := newDocumentReader(bytes.NewReader(doc))
+			src := newBufferedDocumentReader(bytes.NewReader(doc))
 			dst := newValueWriterFromSlice(make([]byte, 0))
 			want := doc
 			err = copyDocument(dst, src)
@@ -77,7 +77,7 @@ func TestCopier(t *testing.T) {
 			noerr(t, err)
 			doc, err = bsoncore.AppendDocumentEnd(doc, idx)
 			noerr(t, err)
-			src := newDocumentReader(bytes.NewReader(doc))
+			src := newBufferedDocumentReader(bytes.NewReader(doc))
 
 			_, err = src.ReadDocument()
 			noerr(t, err)
@@ -450,7 +450,7 @@ func TestCopier(t *testing.T) {
 				idx,
 			)
 			noerr(t, err)
-			vr := newDocumentReader(bytes.NewReader(b))
+			vr := newBufferedDocumentReader(bytes.NewReader(b))
 			_, err = vr.ReadDocument()
 			noerr(t, err)
 			_, _, err = vr.ReadElement()
@@ -489,7 +489,7 @@ func TestCopier(t *testing.T) {
 				idx,
 			)
 			noerr(t, err)
-			vr := newDocumentReader(bytes.NewReader(b))
+			vr := newBufferedDocumentReader(bytes.NewReader(b))
 			_, err = vr.ReadDocument()
 			noerr(t, err)
 			_, _, err = vr.ReadElement()

bson/streaming_value_reader.go

@@ -6,7 +6,9 @@
 
 package bson
 
-import "bufio"
+import (
+	"bufio"
+)
 
 // streamingValueReader reads from an ioReader wrapped in a bufio.Reader. It
 // first reads the BSON length header, then ensures it only ever reads exactly
@@ -15,43 +17,76 @@ import "bufio"
 // Note: this approach trades memory usage for extra buffering and reader calls,
 // so it is less performanted than the in-memory bufferedValueReader.
 type streamingValueReader struct {
-	br *bufio.Reader
+	br     *bufio.Reader
+	offset int64 // offset is the current read position in the buffer
 }
 
 var _ valueReaderByteSrc = (*streamingValueReader)(nil)
 
-// Read reads up to len(p) bytes from buf[offset:] into p, advancing offset.
-func (b *streamingValueReader) Read([]byte) (n int, err error) {
-	panic("Read not implemented for streamingValueReader")
-}
-
 // ReadByte returns the single byte at buf[offset] and advances offset by 1.
-func (b *streamingValueReader) ReadByte() (byte, error) {
-	panic("ReadByte not implemented for streamingValueReader")
+func (s *streamingValueReader) ReadByte() (byte, error) {
+	c, err := s.br.ReadByte()
+	if err == nil {
+		s.offset++
+	}
+	return c, err
 }
 
 // peek returns buf[offset:offset+n] without advancing offset.
-func (b *streamingValueReader) peek(int) ([]byte, error) {
-	panic("peek not implemented for streamingValueReader")
+func (s *streamingValueReader) peek(n int) ([]byte, error) {
+	return s.br.Peek(n)
 }
 
 // discard advances offset by n bytes, returning the number of bytes discarded.
-func (b *streamingValueReader) discard(int) (int, error) {
-	panic("discard not implemented for streamingValueReader")
+func (s *streamingValueReader) discard(n int) (int, error) {
+	m, err := s.br.Discard(n)
+	s.offset += int64(m)
+	return m, err
 }
 
 // readSlice scans buf[offset:] for the first occurrence of delim, returns
 // buf[offset:idx+1], and advances offset past it; errors if delim not found.
-func (b *streamingValueReader) readSlice(byte) ([]byte, error) {
-	panic("readSlice not implemented for streamingValueReader")
+func (s *streamingValueReader) readSlice(delim byte) ([]byte, error) {
+	data, err := s.br.ReadSlice(delim)
+	if err != nil {
+		return nil, err
+	}
+	s.offset += int64(len(data))
+	return data, nil
 }
 
 // pos returns the current read position in the buffer.
-func (b *streamingValueReader) pos() int64 {
-	panic("pos not implemented for streamingValueReader")
+func (s *streamingValueReader) pos() int64 {
+	return s.offset
 }
 
 // regexLength will return the total byte length of a BSON regex value.
-func (b *streamingValueReader) regexLength() (int32, error) {
-	panic("regexLength not implemented for streamingValueReader")
+func (s *streamingValueReader) regexLength() (int32, error) {
+	var (
+		count    int32
+		nulCount int
+	)
+
+	for nulCount < 2 {
+		buf, err := s.br.Peek(int(count) + 1)
+		if err != nil {
+			return 0, err
+		}
+
+		b := buf[count]
+		count++
+		if b == 0x00 {
+			nulCount++
+		}
+	}
+
+	return count, nil
+}
+
+func (*streamingValueReader) streamable() bool {
+	return true
+}
+
+func (s *streamingValueReader) reset() {
+	s.offset = 0
 }