Skip to content

Commit 68fa286

Browse files
authored
Reduce memory usage of the migration verifier (#27)
Generation 0 heap pprof file has shown high memory usage by the `ImportFromCursor` function itself. This PR attempts to improve the memory usage of the migration verifier by: - removing an extra BSON decode step before storing a document to a doc map in `ImportFromCursor` - using a byte buffer when constructing map keys in `getMapKey`
1 parent bf64c12 commit 68fa286

File tree

1 file changed

+11
-13
lines changed

1 file changed

+11
-13
lines changed

internal/documentmap/documentmap.go

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@ package documentmap
2828
// to try to duplicate that here.)
2929

3030
import (
31+
"bytes"
3132
"context"
33+
"errors"
3234
"fmt"
3335

3436
"github.com/10gen/migration-verifier/internal/logger"
@@ -103,30 +105,24 @@ func (m *Map) ImportFromCursor(ctx context.Context, cursor *mongo.Cursor) error
103105
for cursor.Next(ctx) {
104106
err := cursor.Err()
105107
if err != nil {
106-
if err == mongo.ErrNoDocuments {
108+
if errors.Is(err, mongo.ErrNoDocuments) {
107109
break
108110
}
109111

110112
return err
111113
}
112114

113115
nDocumentsReturned++
116+
bytesReturned += (int64)(len(cursor.Current))
114117

115-
var rawDoc bson.Raw
116-
err = cursor.Decode(&rawDoc)
117-
if err != nil {
118-
return err
119-
}
120-
bytesReturned += (int64)(len(rawDoc))
121-
122-
m.addDocument(rawDoc)
118+
m.copyAndAddDocument(cursor.Current)
123119
}
124120
m.logger.Debug().Msgf("Find returned %d documents containing %d bytes", nDocumentsReturned, bytesReturned)
125121

126122
return nil
127123
}
128124

129-
func (m *Map) addDocument(rawDoc bson.Raw) {
125+
func (m *Map) copyAndAddDocument(rawDoc bson.Raw) {
130126
rawDocCopy := make(bson.Raw, len(rawDoc))
131127
copy(rawDocCopy, rawDoc)
132128

@@ -200,11 +196,13 @@ func (m *Map) TotalDocsBytes() types.ByteCount {
200196

201197
// called in tests as well as internally
202198
func (m *Map) getMapKey(doc *bson.Raw) MapKey {
203-
key := MapKey("")
199+
var keyBuffer bytes.Buffer
204200
for _, keyName := range m.fieldNames {
205201
value := doc.Lookup(keyName)
206-
key += MapKey(value.Type) + MapKey(value.Value)
202+
keyBuffer.Grow(1 + len(value.Value))
203+
keyBuffer.WriteByte(byte(value.Type))
204+
keyBuffer.Write(value.Value)
207205
}
208206

209-
return key
207+
return MapKey(keyBuffer.String())
210208
}

0 commit comments

Comments
 (0)