Optimize insertion of rechecks.

This minimizes GC overhead, especially from BSON marshaling.

Author: FGasper

internal/verifier/recheck.go

@@ -10,10 +10,12 @@ import (
 	"github.com/10gen/migration-verifier/internal/types"
 	"github.com/10gen/migration-verifier/internal/util"
 	"github.com/pkg/errors"
+	"github.com/samber/lo"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/bsontype"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
 const (
@@ -31,9 +33,22 @@ const (
 // sorting by _id will guarantee that all rechecks for a given
 // namespace appear consecutively.
 type RecheckPrimaryKey struct {
-	SrcDatabaseName   string `bson:"db"`
-	SrcCollectionName string `bson:"coll"`
-	DocumentID        any    `bson:"docID"`
+	SrcDatabaseName   string        `bson:"db"`
+	SrcCollectionName string        `bson:"coll"`
+	DocumentID        bson.RawValue `bson:"docID"`
+}
+
+var _ bson.Marshaler = &RecheckPrimaryKey{}
+
+func (rk *RecheckPrimaryKey) MarshalBSON() ([]byte, error) {
+	return bsoncore.NewDocumentBuilder().
+		AppendString("db", rk.SrcDatabaseName).
+		AppendString("coll", rk.SrcCollectionName).
+		AppendValue("docID", bsoncore.Value{
+			Type: rk.DocumentID.Type,
+			Data: rk.DocumentID.Value,
+		}).
+		Build(), nil
 }
 
 // RecheckDoc stores the necessary information to know which documents must be rechecked.
@@ -46,6 +61,15 @@ type RecheckDoc struct {
 	DataSize int `bson:"dataSize"`
 }
 
+var _ bson.Marshaler = &RecheckDoc{}
+
+func (rd *RecheckDoc) MarshalBSON() ([]byte, error) {
+	return bsoncore.NewDocumentBuilder().
+		AppendDocument("_id", lo.Must(bson.Marshal(rd.PrimaryKey))).
+		AppendInt64("dataSize", int64(rd.DataSize)).
+		Build(), nil
+}
+
 // InsertFailedCompareRecheckDocs is for inserting RecheckDocs based on failures during Check.
 func (verifier *Verifier) InsertFailedCompareRecheckDocs(
 	ctx context.Context,
@@ -95,50 +119,16 @@ func (verifier *Verifier) insertRecheckDocs(
 
 	genCollection := verifier.getRecheckQueueCollection(generation)
 
-	var recheckBatches [][]mongo.WriteModel
-	var curRechecks []mongo.WriteModel
-	curBatchSize := 0
-	for i, dbName := range dbNames {
-		recheckDoc := RecheckDoc{
-			PrimaryKey: RecheckPrimaryKey{
-				SrcDatabaseName:   dbName,
-				SrcCollectionName: collNames[i],
-				DocumentID:        rawDocIDs[i],
-			},
-			DataSize: dataSizes[i],
-		}
-
-		recheckRaw, err := bson.Marshal(recheckDoc)
-		if err != nil {
-			return errors.Wrapf(err, "marshaling recheck for %#q", dbName+"."+collNames[i])
-		}
-
-		curRechecks = append(
-			curRechecks,
-			mongo.NewInsertOneModel().SetDocument(recheckDoc),
-		)
-		curBatchSize += len(recheckRaw)
-		if curBatchSize > recheckBatchByteLimit || len(curRechecks) >= recheckBatchCountLimit {
-			recheckBatches = append(recheckBatches, curRechecks)
-			curRechecks = nil
-			curBatchSize = 0
-		}
-	}
-
-	if len(curRechecks) > 0 {
-		recheckBatches = append(recheckBatches, curRechecks)
-	}
-
-	for _, models := range recheckBatches {
+	sendRechecks := func(rechecks []bson.Raw) {
 		eg.Go(func() error {
 
 			retryer := retry.New()
 			err := retryer.WithCallback(
 				func(retryCtx context.Context, _ *retry.FuncInfo) error {
-					_, err := genCollection.BulkWrite(
+					_, err := genCollection.InsertMany(
 						retryCtx,
-						models,
-						options.BulkWrite().SetOrdered(false),
+						lo.ToAnySlice(rechecks),
+						options.InsertMany().SetOrdered(false),
 					)
 
 					// We expect duplicate-key errors from the above because:
@@ -157,20 +147,54 @@ func (verifier *Verifier) insertRecheckDocs(
 					// and document sizes probably remain stable(-ish) across updates.
 					err = util.TolerateSimpleDuplicateKeyInBulk(
 						verifier.logger,
-						len(models),
+						len(rechecks),
 						err,
 					)
 
 					return err
 				},
 				"persisting %d recheck(s)",
-				len(models),
+				len(rechecks),
 			).Run(groupCtx, verifier.logger)
 
-			return errors.Wrapf(err, "batch of %d rechecks", len(models))
+			return errors.Wrapf(err, "batch of %d rechecks", len(rechecks))
 		})
 	}
 
+	curRechecks := make([]bson.Raw, 0, recheckBatchCountLimit)
+	curBatchBytes := 0
+	for i, dbName := range dbNames {
+		recheckDoc := RecheckDoc{
+			PrimaryKey: RecheckPrimaryKey{
+				SrcDatabaseName:   dbName,
+				SrcCollectionName: collNames[i],
+				DocumentID:        rawDocIDs[i],
+			},
+			DataSize: dataSizes[i],
+		}
+
+		recheckRaw, err := bson.Marshal(recheckDoc)
+		if err != nil {
+			return errors.Wrapf(err, "marshaling recheck for %#q", dbName+"."+collNames[i])
+		}
+
+		curRechecks = append(
+			curRechecks,
+			bson.Raw(recheckRaw),
+		)
+
+		curBatchBytes += len(recheckRaw)
+		if curBatchBytes > recheckBatchByteLimit || len(curRechecks) >= recheckBatchCountLimit {
+			sendRechecks(curRechecks)
+			curRechecks = make([]bson.Raw, 0, recheckBatchCountLimit)
+			curBatchBytes = 0
+		}
+	}
+
+	if len(curRechecks) > 0 {
+		sendRechecks(curRechecks)
+	}
+
 	if err := eg.Wait(); err != nil {
 		return errors.Wrapf(
 			err,

internal/verifier/recheck_test.go

@@ -39,7 +39,7 @@ func (suite *IntegrationTestSuite) TestFailedCompareThenReplace() {
 				PrimaryKey: RecheckPrimaryKey{
 					SrcDatabaseName:   "the",
 					SrcCollectionName: "namespace",
-					DocumentID:        "theDocID",
+					DocumentID:        mbson.ToRawValue("theDocID"),
 				},
 			},
 		},
@@ -74,7 +74,7 @@ func (suite *IntegrationTestSuite) TestFailedCompareThenReplace() {
 				PrimaryKey: RecheckPrimaryKey{
 					SrcDatabaseName:   "the",
 					SrcCollectionName: "namespace",
-					DocumentID:        "theDocID",
+					DocumentID:        mbson.ToRawValue("theDocID"),
 				},
 			},
 		},
@@ -282,13 +282,13 @@ func (suite *IntegrationTestSuite) TestLargeIDInsertions() {
 		PrimaryKey: RecheckPrimaryKey{
 			SrcDatabaseName:   "testDB",
 			SrcCollectionName: "testColl",
-			DocumentID:        id1,
+			DocumentID:        mbson.ToRawValue(id1),
 		},
 	}
 	d2 := d1
-	d2.PrimaryKey.DocumentID = id2
+	d2.PrimaryKey.DocumentID = mbson.ToRawValue(id2)
 	d3 := d1
-	d3.PrimaryKey.DocumentID = id3
+	d3.PrimaryKey.DocumentID = mbson.ToRawValue(id3)
 
 	results := suite.fetchRecheckDocs(ctx, verifier)
 	suite.ElementsMatch([]any{d1, d2, d3}, results)
@@ -342,13 +342,13 @@ func (suite *IntegrationTestSuite) TestLargeDataInsertions() {
 		PrimaryKey: RecheckPrimaryKey{
 			SrcDatabaseName:   "testDB",
 			SrcCollectionName: "testColl",
-			DocumentID:        id1,
+			DocumentID:        mbson.ToRawValue(id1),
 		},
 	}
 	d2 := d1
-	d2.PrimaryKey.DocumentID = id2
+	d2.PrimaryKey.DocumentID = mbson.ToRawValue(id2)
 	d3 := d1
-	d3.PrimaryKey.DocumentID = id3
+	d3.PrimaryKey.DocumentID = mbson.ToRawValue(id3)
 
 	results := suite.fetchRecheckDocs(ctx, verifier)
 	suite.ElementsMatch([]any{d1, d2, d3}, results)
@@ -451,11 +451,11 @@ func (suite *IntegrationTestSuite) TestGenerationalClear() {
 		PrimaryKey: RecheckPrimaryKey{
 			SrcDatabaseName:   "testDB",
 			SrcCollectionName: "testColl",
-			DocumentID:        id1,
+			DocumentID:        mbson.ToRawValue(id1),
 		},
 	}
 	d2 := d1
-	d2.PrimaryKey.DocumentID = id2
+	d2.PrimaryKey.DocumentID = mbson.ToRawValue(id2)
 
 	results := suite.fetchRecheckDocs(ctx, verifier)
 	suite.Assert().ElementsMatch([]any{d1, d2}, results)