From d6bf08de67b81031380dd1ada6104a08ffe1ab13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20Pavela?= Date: Thu, 9 Jul 2026 21:27:05 +0200 Subject: [PATCH] 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 --- x/mongo/driver/batches.go | 69 ++++++++--- x/mongo/driver/batches_test.go | 219 ++++++++++++++++++++++++++------- 2 files changed, 225 insertions(+), 63 deletions(-) diff --git a/x/mongo/driver/batches.go b/x/mongo/driver/batches.go index ed1b2882b..a12257bbe 100644 --- a/x/mongo/driver/batches.go +++ b/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 diff --git a/x/mongo/driver/batches_test.go b/x/mongo/driver/batches_test.go index 5efc57c7e..f87c60e50 100644 --- a/x/mongo/driver/batches_test.go +++ b/x/mongo/driver/batches_test.go @@ -7,6 +7,7 @@ package driver import ( + "strconv" "testing" "go.mongodb.org/mongo-driver/v2/internal/assert" @@ -14,62 +15,188 @@ import ( "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) + } + } + }) + } }