Merge branch 'release/2.0' into DEVPROD-17319-cherry-pick-2.0

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

internal/integration/client_side_encryption_prose_test.go

@@ -795,7 +795,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 				"endpoint":   "doesnotexist.invalid:443",
 			},
 			"kmip": {
-				"endpoint": "doesnotexist.local:5698",
+				"endpoint": "doesnotexist.invalid:5698",
 			},
 		}
 
@@ -819,9 +819,8 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			"endpoint": "kms.us-east-1.amazonaws.com:443",
 		}
 		awsFailureConnectionError := map[string]interface{}{
-			"region":   "us-east-1",
-			"key":      "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0",
-			"endpoint": "kms.us-east-1.amazonaws.com:12345",
+			"keyId":    "1",
+			"endpoint": "localhost:12345",
 		}
 		awsFailureInvalidEndpoint := map[string]interface{}{
 			"region":   "us-east-1",
@@ -860,7 +859,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		}
 		kmipFailureInvalidEndpoint := map[string]interface{}{
 			"keyId":    "1",
-			"endpoint": "doesnotexist.local:5698",
+			"endpoint": "doesnotexist.invalid:5698",
 		}
 
 		const (
@@ -906,7 +905,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			},
 			{
 				name:                                  "Case 4: aws failure with connection error",
-				provider:                              "aws",
+				provider:                              "kmip",
 				masterKey:                             awsFailureConnectionError,
 				errorSubstring:                        []string{errConnectionRefused, errWindowsTLSConnectionRefused},
 				testInvalidClientEncryption:           false,
@@ -1583,7 +1582,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			"endpoint": "127.0.0.1:9001",
 		}
 		azureMasterKey := map[string]interface{}{
-			"keyVaultEndpoint": "doesnotexist.local",
+			"keyVaultEndpoint": "doesnotexist.invalid",
 			"keyName":          "foo",
 		}
 		gcpMasterKey := map[string]interface{}{