GODRIVER-3533 Optimize value reader and writer

Author: prestonvasquez

bson/marshal.go

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

bson/unmarshal.go

@@ -43,6 +43,8 @@ type ValueUnmarshaler interface {
 // pointer, the pointer is set to nil without calling UnmarshalBSONValue.
 func Unmarshal(data []byte, val interface{}) error {
 	vr := newDocumentReader(bytes.NewReader(data))
+	defer releaseDocumentReader(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,13 +30,30 @@ 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
-	offset int64
+	r  *bufio.Reader
+	ra io.ReaderAt // The underlying reader
 
-	stack []vrState
-	frame int64
+	poolReader bool // Set if r is from bufioReaderPool
+
+	offset int64
+	stack  []vrState
+	frame  int64
 }
 
 // NewDocumentReader returns a ValueReader using b for the underlying BSON
@@ -57,14 +75,33 @@ func newValueReader(t Type, r io.Reader) ValueReader {
 }
 
 func newDocumentReader(r io.Reader) *valueReader {
-	stack := make([]vrState, 1, 5)
-	stack[0] = vrState{
-		mode: mTopLevel,
-	}
-	return &valueReader{
-		r:     bufio.NewReader(r),
-		stack: stack,
+	vr := vrPool.Get().(*valueReader)
+
+	vr.offset = 0
+	vr.frame = 0
+	vr.poolReader = true
+
+	vr.stack = vr.stack[:1]
+	vr.stack[0].mode = mTopLevel
+
+	br := bufioReaderPool.Get().(*bufio.Reader)
+	br.Reset(r)
+	vr.r = br
+
+	return vr
+}
+
+func releaseDocumentReader(vr *valueReader) {
+	if !vr.poolReader {
+		return
 	}
+
+	bufioReaderPool.Put(vr.r)
+	vr.r = nil
+	vr.poolReader = false
+
+	vr.ra = nil
+	vrPool.Put(vr)
 }
 
 func (vr *valueReader) advanceFrame() {
@@ -253,14 +290,18 @@ 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 {
 		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

@@ -26,6 +26,16 @@ var vwPool = sync.Pool{
 	},
 }
 
+func getValueWriter(w io.Writer) *valueWriter {
+	vw := vwPool.Get().(*valueWriter)
+
+	vw.reset(vw.buf)
+	vw.buf = vw.buf[:0]
+	vw.w = w
+
+	return vw
+}
+
 func putValueWriter(vw *valueWriter) {
 	if vw != nil {
 		vw.w = nil // don't leak the writer