GODRIVER-4018 Reduce allocations when building write batches

Reserve the output buffer once before appending documents in
AppendBatchSequence and AppendBatchArray, instead of letting the
per-document append grow it by repeated reallocation. A first pass
computes how many documents fit within the maxCount/totalSize limits and
their combined size; dst is grown a single time, then the documents are
appended. When dst already has capacity the reservation is a no-op.

AppendBatchArray additionally wrote each element key with strconv.Itoa,
allocating a string per document; it now writes the numeric key in place
with strconv.AppendInt via a helper that composes the same bsoncore
primitives, producing byte-identical output.

Behavior is unchanged: same batch count, byte-identical output, same
maxCount/totalSize handling, and the same empty-result early return.

Benchmark (1000 x 256-byte docs, fresh buffer):
  AppendBatchSequence: 24 -> 3 allocs/op, 1.19MB -> 262KB, ~130us -> ~29us
  AppendBatchArray:    925 -> 3 allocs/op, 1.19MB -> 262KB, ~166us -> ~38us

Author: paveon

x/mongo/driver/batches.go

@@ -28,7 +28,7 @@ var _ OperationBatches = &Batches{}
 
 // AppendBatchSequence appends dst with document sequence of batches as long as the limits of max count, max
 // document size, or total size allows. It returns the number of batches appended, the new appended slice, and
-// any error raised. It returns the origenal input slice if nothing can be appends within the limits.
+// any error raised. It returns the original input slice if nothing can be appended within the limits.
 func (b *Batches) AppendBatchSequence(dst []byte, maxCount, totalSize int) (int, []byte, error) {
 	if b.Size() == 0 {
 		return 0, dst, io.EOF
@@ -39,53 +39,79 @@ func (b *Batches) AppendBatchSequence(dst []byte, maxCount, totalSize int) (int,
 	idx, dst = bsoncore.ReserveLength(dst)
 	dst = append(dst, b.Identifier...)
 	dst = append(dst, 0x00)
+
+	// First pass: total the documents that fit within the limits so dst can be
+	// grown once instead of reallocating on each append.
 	size := len(dst)
-	var n int
-	for i := b.offset; i < len(b.Documents); i++ {
-		if n == maxCount {
+	n, docsSize := 0, 0
+	for n < maxCount && b.offset+n < len(b.Documents) {
+		doc := b.Documents[b.offset+n]
+		if size+len(doc) > totalSize {
 			break
 		}
-		doc := b.Documents[i]
 		size += len(doc)
-		if size > totalSize {
-			break
-		}
-		dst = append(dst, doc...)
+		docsSize += len(doc)
 		n++
 	}
 	if n == 0 {
 		return 0, dst[:l], nil
 	}
+
+	// Reserve once; no-op when dst already has room.
+	if cap(dst) < len(dst)+docsSize {
+		dst = append(make([]byte, 0, len(dst)+docsSize), dst...)
+	}
+	for _, doc := range b.Documents[b.offset : b.offset+n] {
+		dst = append(dst, doc...)
+	}
+
 	dst = bsoncore.UpdateLength(dst, idx, int32(len(dst[idx:])))
 	return n, dst, nil
 }
 
 // AppendBatchArray appends dst with array of batches as long as the limits of max count, max document size, or
 // total size allows. It returns the number of batches appended, the new appended slice, and any error raised. It
-// returns the origenal input slice if nothing can be appends within the limits.
+// returns the original input slice if nothing can be appended within the limits.
 func (b *Batches) AppendBatchArray(dst []byte, maxCount, totalSize int) (int, []byte, error) {
 	if b.Size() == 0 {
 		return 0, dst, io.EOF
 	}
 	l := len(dst)
 	aidx, dst := bsoncore.AppendArrayElementStart(dst, b.Identifier)
+
+	// First pass: total the bytes the elements that fit will need so dst can be
+	// grown once. Each element is a type byte, the decimal index key, a null
+	// terminator and the document; keyLen is the current index's digit count,
+	// bumped as n crosses each power of ten.
 	size := len(dst)
-	var n int
-	for i := b.offset; i < len(b.Documents); i++ {
-		if n == maxCount {
+	n, appendSize := 0, 0
+	keyLen, nextPow10 := 1, 10
+	for n < maxCount && b.offset+n < len(b.Documents) {
+		doc := b.Documents[b.offset+n]
+		if size+len(doc) > totalSize {
 			break
 		}
-		doc := b.Documents[i]
 		size += len(doc)
-		if size > totalSize {
-			break
+		if n == nextPow10 {
+			keyLen++
+			nextPow10 *= 10
 		}
-		dst = bsoncore.AppendDocumentElement(dst, strconv.Itoa(n), doc)
+		appendSize += 2 + keyLen + len(doc)
 		n++
 	}
 	if n == 0 {
 		return 0, dst[:l], nil
 	}
+	appendSize++ // closing byte written by AppendArrayEnd
+
+	// Reserve once; no-op when dst already has room.
+	if cap(dst) < len(dst)+appendSize {
+		dst = append(make([]byte, 0, len(dst)+appendSize), dst...)
+	}
+	for j := 0; j < n; j++ {
+		dst = appendDocumentElementInt(dst, j, b.Documents[b.offset+j])
+	}
+
 	var err error
 	dst, err = bsoncore.AppendArrayEnd(dst, aidx)
 	if err != nil {
@@ -94,6 +120,15 @@ func (b *Batches) AppendBatchArray(dst []byte, maxCount, totalSize int) (int, []
 	return n, dst, nil
 }
 
+// appendDocumentElementInt is bsoncore.AppendDocumentElement with an integer
+// key, formatting the key in place to avoid allocating a string per element.
+func appendDocumentElementInt(dst []byte, key int, doc []byte) []byte {
+	dst = bsoncore.AppendType(dst, bsoncore.TypeEmbeddedDocument)
+	dst = strconv.AppendInt(dst, int64(key), 10)
+	dst = append(dst, 0x00)
+	return bsoncore.AppendDocument(dst, doc)
+}
+
 // IsOrdered indicates if the batches are ordered.
 func (b *Batches) IsOrdered() *bool {
 	return b.Ordered

x/mongo/driver/batches_test.go

@@ -7,69 +7,196 @@
 package driver
 
 import (
+	"strconv"
 	"testing"
 
 	"go.mongodb.org/mongo-driver/v2/internal/assert"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/wiremessage"
 )
 
-func newTestBatches(t *testing.T) *Batches {
-	t.Helper()
-	return &Batches{
-		Identifier: "foobar",
-		Documents: []bsoncore.Document{
-			[]byte("Lorem ipsum dolor sit amet"),
-			[]byte("consectetur adipiscing elit"),
-		},
+const testIdentifier = "foobar"
+
+var testDocs = [][]byte{
+	[]byte("Lorem ipsum dolor sit amet"),
+	[]byte("consectetur adipiscing elit"),
+}
+
+// newBatches builds a Batches with the test identifier from docs.
+func newBatches(docs ...[]byte) *Batches {
+	b := &Batches{Identifier: testIdentifier}
+	for _, doc := range docs {
+		b.Documents = append(b.Documents, bsoncore.Document(doc))
+	}
+	return b
+}
+
+// wantSequence builds the expected AppendBatchSequence output: prefix followed by
+// a document sequence section holding docs.
+func wantSequence(prefix []byte, docs [][]byte) []byte {
+	dst := append([]byte(nil), prefix...)
+	dst = wiremessage.AppendMsgSectionType(dst, wiremessage.DocumentSequence)
+	idx, dst := bsoncore.ReserveLength(dst)
+	dst = append(dst, testIdentifier...)
+	dst = append(dst, 0x00)
+	for _, doc := range docs {
+		dst = append(dst, doc...)
 	}
+	return bsoncore.UpdateLength(dst, idx, int32(len(dst[idx:])))
+}
+
+// wantArray builds the expected AppendBatchArray output using the straightforward
+// bsoncore.AppendDocumentElement / strconv.Itoa construction. Comparing against it
+// also verifies AppendBatchArray's in-place integer-key encoding.
+func wantArray(prefix []byte, docs [][]byte) []byte {
+	dst := append([]byte(nil), prefix...)
+	idx, dst := bsoncore.AppendArrayElementStart(dst, testIdentifier)
+	for i, doc := range docs {
+		dst = bsoncore.AppendDocumentElement(dst, strconv.Itoa(i), doc)
+	}
+	dst, _ = bsoncore.AppendArrayEnd(dst, idx)
+	return dst
+}
+
+// benchBatches builds a Batches of numDocs identical documents of docSize bytes.
+func benchBatches(numDocs, docSize int) *Batches {
+	b := &Batches{Identifier: "documents"}
+	doc := bsoncore.Document(make([]byte, docSize))
+	for i := 0; i < numDocs; i++ {
+		b.Documents = append(b.Documents, doc)
+	}
+	return b
 }
 
 func TestAdvancing(t *testing.T) {
-	batches := newTestBatches(t)
+	batches := newBatches(testDocs...)
 	batches.AdvanceBatches(3)
 	size := batches.Size()
-	assert.Equal(t, 0, size, "expected Size(): %d, got: %d", 1, size)
+	assert.Equal(t, 0, size, "expected Size(): %d, got: %d", 0, size)
 }
 
 func TestAppendBatchSequence(t *testing.T) {
-	batches := newTestBatches(t)
-
-	got := []byte{42}
-	sizeLimit := len(batches.Documents[0]) + len(batches.Documents[1])
-	var n int
-	var err error
-	n, got, err = batches.AppendBatchSequence(got, 2, sizeLimit)
-	assert.NoError(t, err)
-	assert.Equal(t, 1, n)
-
-	var idx int32
-	dst := []byte{42}
-	dst = wiremessage.AppendMsgSectionType(dst, wiremessage.DocumentSequence)
-	idx, dst = bsoncore.ReserveLength(dst)
-	dst = append(dst, "foobar"...)
-	dst = append(dst, 0x00)
-	dst = append(dst, "Lorem ipsum dolor sit amet"...)
-	dst = bsoncore.UpdateLength(dst, idx, int32(len(dst[idx:])))
-	assert.Equal(t, dst, got)
+	tests := []struct {
+		name      string
+		docs      [][]byte
+		advance   int
+		maxCount  int
+		sizeLimit int
+		want      [][]byte
+	}{
+		{
+			name:      "size limit fits only the first document",
+			docs:      testDocs,
+			maxCount:  2,
+			sizeLimit: len(testDocs[0]) + len(testDocs[1]),
+			want:      testDocs[:1],
+		},
+		{
+			name:      "all documents fit",
+			docs:      testDocs,
+			maxCount:  10,
+			sizeLimit: 16 * 1024 * 1024,
+			want:      testDocs,
+		},
+		{
+			name:      "offset skips leading documents",
+			docs:      testDocs,
+			advance:   1,
+			maxCount:  10,
+			sizeLimit: 16 * 1024 * 1024,
+			want:      testDocs[1:],
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			batches := newBatches(tt.docs...)
+			batches.AdvanceBatches(tt.advance)
+
+			n, got, err := batches.AppendBatchSequence([]byte{42}, tt.maxCount, tt.sizeLimit)
+			assert.NoError(t, err)
+			assert.Equal(t, len(tt.want), n)
+			assert.Equal(t, wantSequence([]byte{42}, tt.want), got)
+		})
+	}
 }
 
 func TestAppendBatchArray(t *testing.T) {
-	batches := newTestBatches(t)
-
-	got := []byte{42}
-	sizeLimit := len(batches.Documents[0]) + len(batches.Documents[1])
-	var n int
-	var err error
-	n, got, err = batches.AppendBatchArray(got, 2, sizeLimit)
-	assert.NoError(t, err)
-	assert.Equal(t, 1, n)
-
-	var idx int32
-	dst := []byte{42}
-	idx, dst = bsoncore.AppendArrayElementStart(dst, "foobar")
-	dst = bsoncore.AppendDocumentElement(dst, "0", []byte("Lorem ipsum dolor sit amet"))
-	dst, err = bsoncore.AppendArrayEnd(dst, idx)
-	assert.NoError(t, err)
-	assert.Equal(t, dst, got)
+	manyDocs := make([][]byte, 150) // exercises multi-digit array keys
+	for i := range manyDocs {
+		manyDocs[i] = testDocs[0]
+	}
+
+	tests := []struct {
+		name      string
+		docs      [][]byte
+		advance   int
+		maxCount  int
+		sizeLimit int
+		want      [][]byte
+	}{
+		{
+			name:      "size limit fits only the first document",
+			docs:      testDocs,
+			maxCount:  2,
+			sizeLimit: len(testDocs[0]) + len(testDocs[1]),
+			want:      testDocs[:1],
+		},
+		{
+			name:      "all documents fit",
+			docs:      testDocs,
+			maxCount:  10,
+			sizeLimit: 16 * 1024 * 1024,
+			want:      testDocs,
+		},
+		{
+			name:      "offset skips leading documents",
+			docs:      testDocs,
+			advance:   1,
+			maxCount:  10,
+			sizeLimit: 16 * 1024 * 1024,
+			want:      testDocs[1:],
+		},
+		{
+			name:      "many elements with multi-digit keys",
+			docs:      manyDocs,
+			maxCount:  len(manyDocs),
+			sizeLimit: 16 * 1024 * 1024,
+			want:      manyDocs,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			batches := newBatches(tt.docs...)
+			batches.AdvanceBatches(tt.advance)
+
+			n, got, err := batches.AppendBatchArray([]byte{42}, tt.maxCount, tt.sizeLimit)
+			assert.NoError(t, err)
+			assert.Equal(t, len(tt.want), n)
+			assert.Equal(t, wantArray([]byte{42}, tt.want), got)
+		})
+	}
+}
+
+func BenchmarkAppendBatch(b *testing.B) {
+	batches := benchBatches(1000, 256)
+
+	benchmarks := []struct {
+		name string
+		fn   func(*Batches, []byte, int, int) (int, []byte, error)
+	}{
+		{"Sequence", (*Batches).AppendBatchSequence},
+		{"Array", (*Batches).AppendBatchArray},
+	}
+	for _, bm := range benchmarks {
+		b.Run(bm.name, func(b *testing.B) {
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				batches.offset = 0
+				// A nil buffer measures the growth a BulkWrite pays per operation.
+				if _, _, err := bm.fn(batches, nil, len(batches.Documents), 16*1024*1024); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
 }