GODRIVER-3533 Optimize value reader and writer (#2022)

Author: prestonvasquez

bson/marshal.go

@@ -74,7 +74,10 @@ func Marshal(val interface{}) ([]byte, error) {
 		}
 	}()
 	sw.Reset()
-	vw := NewDocumentWriter(sw)
+
+	vw := getDocumentWriter(sw)
+	defer putDocumentWriter(vw)
+
 	enc := encPool.Get().(*Encoder)
 	defer encPool.Put(enc)
 	enc.Reset(vw)

bson/unmarshal.go

@@ -42,7 +42,9 @@ type ValueUnmarshaler interface {
 // When unmarshaling BSON, if the BSON value is null and the Go value is a
 // pointer, the pointer is set to nil without calling UnmarshalBSONValue.
 func Unmarshal(data []byte, val interface{}) error {
-	vr := newDocumentReader(bytes.NewReader(data))
+	vr := getDocumentReader(bytes.NewReader(data))
+	defer putDocumentReader(vr)
+
 	if l, err := vr.peekLength(); err != nil {
 		return err
 	} else if int(l) != len(data) {

bson/value_reader.go

@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"sync"
 )
 
 var _ ValueReader = &valueReader{}
@@ -29,6 +30,20 @@ type vrState struct {
 	end   int64
 }
 
+var bufioReaderPool = sync.Pool{
+	New: func() interface{} {
+		return bufio.NewReader(nil)
+	},
+}
+
+var vrPool = sync.Pool{
+	New: func() interface{} {
+		return &valueReader{
+			stack: make([]vrState, 1, 5),
+		}
+	},
+}
+
 // valueReader is for reading BSON values.
 type valueReader struct {
 	r      *bufio.Reader
@@ -38,6 +53,33 @@ type valueReader struct {
 	frame int64
 }
 
+func getDocumentReader(r io.Reader) *valueReader {
+	vr := vrPool.Get().(*valueReader)
+
+	vr.offset = 0
+	vr.frame = 0
+
+	vr.stack = vr.stack[:1]
+	vr.stack[0] = vrState{mode: mTopLevel}
+
+	br := bufioReaderPool.Get().(*bufio.Reader)
+	br.Reset(r)
+	vr.r = br
+
+	return vr
+}
+
+func putDocumentReader(vr *valueReader) {
+	if vr == nil {
+		return
+	}
+
+	bufioReaderPool.Put(vr.r)
+	vr.r = nil
+
+	vrPool.Put(vr)
+}
+
 // NewDocumentReader returns a ValueReader using b for the underlying BSON
 // representation.
 func NewDocumentReader(r io.Reader) ValueReader {
@@ -253,14 +295,28 @@ func (vr *valueReader) appendNextElement(dst []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	buf := make([]byte, length)
-	_, err = io.ReadFull(vr.r, buf)
+	buf, err := vr.r.Peek(int(length))
 	if err != nil {
+		if err == bufio.ErrBufferFull {
+			temp := make([]byte, length)
+			if _, err = io.ReadFull(vr.r, temp); err != nil {
+				return nil, err
+			}
+			dst = append(dst, temp...)
+			vr.offset += int64(len(temp))
+			return dst, nil
+		}
+
 		return nil, err
 	}
+
 	dst = append(dst, buf...)
-	vr.offset += int64(len(buf))
-	return dst, err
+	if _, err = vr.r.Discard(int(length)); err != nil {
+		return nil, err
+	}
+
+	vr.offset += int64(length)
+	return dst, nil
 }
 
 func (vr *valueReader) readValueBytes(dst []byte) (Type, []byte, error) {

bson/value_writer.go

@@ -33,6 +33,29 @@ func putValueWriter(vw *valueWriter) {
 	}
 }
 
+var documentWriterPool = sync.Pool{
+	New: func() interface{} {
+		return newDocumentWriter(nil)
+	},
+}
+
+func getDocumentWriter(w io.Writer) *valueWriter {
+	vw := documentWriterPool.Get().(*valueWriter)
+
+	vw.reset(vw.buf)
+	vw.buf = vw.buf[:0]
+	vw.w = w
+
+	return vw
+}
+
+func putDocumentWriter(vw *valueWriter) {
+	if vw != nil {
+		vw.w = nil // don't leak the writer
+		documentWriterPool.Put(vw)
+	}
+}
+
 // This is here so that during testing we can change it and not require
 // allocating a 4GB slice.
 var maxSize = math.MaxInt32