Skip to content

GODRIVER-4018 Reduce allocations when building write batches #2476

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
paveon wants to merge 1 commit into
base: master
Choose a base branch
from
+225 −63
Open

GODRIVER-4018 Reduce allocations when building write batches #2476

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 52 additions & 17 deletions x/mongo/driver/batches.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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
Expand Down
219 changes: 173 additions & 46 deletions x/mongo/driver/batches_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
})
}
}