Enqueue rechecks more efficiently (#154)

- Send a raw `insert` command rather than InsertMany. This prevents the driver from “sneakily” ensuring that each inserted document has an `_id`. (This also requires making TolerateSimpleDuplicateKeyInBulk smarter.)
- Avoid bson.Marshal() and bson.Unmarshal in hot paths. Even when the struct implements the Marshaler & Unmarshaler interfaces, the driver still does some extra work that, while not as problematic as reflection, is still better avoided.
- Avoid reallocations when marshaling recheck structs.

Author: FGasper

internal/util/error.go

@@ -10,6 +10,7 @@ import (
 	"github.com/10gen/migration-verifier/mmongo"
 	mapset "github.com/deckarep/golang-set/v2"
 	"github.com/pkg/errors"
+	"github.com/rs/zerolog"
 	"github.com/samber/lo"
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
@@ -465,26 +466,45 @@ func TolerateSimpleDuplicateKeyInBulk(
 	docsCount int,
 	err error,
 ) error {
-	bwe := mongo.BulkWriteException{}
-	if errors.As(err, &bwe) && bwe.WriteConcernError == nil {
-		writeCodes := lo.Map(
-			bwe.WriteErrors,
-			func(we mongo.BulkWriteError, _ int) int {
-				return we.Code
-			},
-		)
-
-		codesSet := mapset.NewSet(writeCodes...)
-		if codesSet.Cardinality() == 1 && writeCodes[0] == DuplicateKeyErrCode {
-			// This will be fairly common since we now listen for change events
-			// on both source & destination, so use trace level here.
+	if err == nil {
+		return nil
+	}
+
+	var writeCodes []int
+
+	if bwe := (mongo.BulkWriteException{}); errors.As(err, &bwe) {
+		if bwe.WriteConcernError == nil {
+			writeCodes = lo.Map(
+				bwe.WriteErrors,
+				func(we mongo.BulkWriteError, _ int) int {
+					return we.Code
+				},
+			)
+		}
+	} else if we := (mongo.WriteException{}); errors.As(err, &we) {
+		if we.WriteConcernError == nil {
+			writeCodes = lo.Map(
+				we.WriteErrors,
+				func(we mongo.WriteError, _ int) int {
+					return we.Code
+				},
+			)
+		}
+	}
+
+	codesSet := mapset.NewSet(writeCodes...)
+	if codesSet.Cardinality() == 1 && writeCodes[0] == DuplicateKeyErrCode {
+		// This will be fairly common since we now listen for change events
+		// on both source & destination, so use trace level here, and don’t
+		// do the logger calls at all unless the level warrants.
+		if logger.GetLevel() <= zerolog.TraceLevel {
 			logger.Trace().
 				Int("documentsSubmitted", docsCount).
 				Int("duplicates", len(writeCodes)).
 				Msg("Ignoring duplicate key error on recheck inserts.")
-
-			err = nil
 		}
+
+		err = nil
 	}
 
 	return err

internal/verifier/change_stream.go

@@ -177,7 +177,7 @@ func (verifier *Verifier) HandleChangeStreamEvents(ctx context.Context, batch ch
 	dbNames := make([]string, len(batch.events))
 	collNames := make([]string, len(batch.events))
 	docIDs := make([]bson.RawValue, len(batch.events))
-	dataSizes := make([]int, len(batch.events))
+	dataSizes := make([]int32, len(batch.events))
 
 	latestTimestamp := bson.Timestamp{}
 
@@ -230,14 +230,14 @@ func (verifier *Verifier) HandleChangeStreamEvents(ctx context.Context, batch ch
 		docIDs[i] = changeEvent.DocID
 
 		if changeEvent.FullDocLen.OrZero() > 0 {
-			dataSizes[i] = int(changeEvent.FullDocLen.OrZero())
+			dataSizes[i] = int32(changeEvent.FullDocLen.OrZero())
 		} else if changeEvent.FullDocument == nil {
 			// This happens for deletes and for some updates.
 			// The document is probably, but not necessarily, deleted.
 			dataSizes[i] = fauxDocSizeForDeleteEvents
 		} else {
 			// This happens for inserts, replaces, and most updates.
-			dataSizes[i] = len(changeEvent.FullDocument)
+			dataSizes[i] = int32(len(changeEvent.FullDocument))
 		}
 
 		if err := eventRecorder.AddEvent(&changeEvent); err != nil {
@@ -382,7 +382,7 @@ func (csr *ChangeStreamReader) readAndHandleOneChangeEventBatch(
 
 		batchTotalBytes += len(cs.Current)
 
-		if err := cs.Decode(&changeEvents[eventsRead]); err != nil {
+		if err := (&changeEvents[eventsRead]).UnmarshalFromBSON(cs.Current); err != nil {
 			return errors.Wrapf(err, "failed to decode change event to %T", changeEvents[eventsRead])
 		}
 

internal/verifier/change_stream_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/10gen/migration-verifier/internal/testutil"
 	"github.com/10gen/migration-verifier/internal/types"
 	"github.com/10gen/migration-verifier/internal/util"
+	"github.com/10gen/migration-verifier/internal/verifier/recheck"
 	"github.com/10gen/migration-verifier/mbson"
 	"github.com/10gen/migration-verifier/mslices"
 	"github.com/10gen/migration-verifier/mstrings"
@@ -127,7 +128,7 @@ func (suite *IntegrationTestSuite) TestChangeStreamFilter_BsonSize() {
 	suite.Assert().Less(len(cs.Current), 10_000, "event should not be large")
 
 	parsed := ParsedEvent{}
-	suite.Require().NoError(cs.Decode(&parsed))
+	suite.Require().NoError((&parsed).UnmarshalFromBSON(cs.Current))
 	suite.Require().Equal("insert", parsed.OpType)
 
 	suite.Require().Equal(
@@ -144,7 +145,7 @@ func (suite *IntegrationTestSuite) TestChangeStreamFilter_BsonSize() {
 
 	suite.Require().True(cs.Next(ctx), "should get event")
 	parsed = ParsedEvent{}
-	suite.Require().NoError(cs.Decode(&parsed))
+	suite.Require().NoError((&parsed).UnmarshalFromBSON(cs.Current))
 	suite.Require().Equal("delete", parsed.OpType)
 	suite.Require().True(parsed.FullDocLen.IsNone(), "full doc len not in delete")
 }
@@ -916,7 +917,7 @@ func (suite *IntegrationTestSuite) TestRecheckDocsWithDstChangeEvents() {
 	_, err = coll2.InsertOne(ctx, bson.D{{"_id", 1}})
 	suite.Require().NoError(err)
 
-	var rechecks []RecheckDoc
+	var rechecks []recheck.Doc
 	require.Eventually(
 		suite.T(),
 		func() bool {

internal/verifier/compare.go

@@ -322,7 +322,7 @@ func (verifier *Verifier) compareDocsFromChannels(
 				Details:      Missing,
 				Cluster:      ClusterTarget,
 				NameSpace:    namespace,
-				dataSize:     len(docWithTs.doc),
+				dataSize:     int32(len(docWithTs.doc)),
 				SrcTimestamp: option.Some(docWithTs.ts),
 			},
 		)
@@ -339,7 +339,7 @@ func (verifier *Verifier) compareDocsFromChannels(
 				Details:      Missing,
 				Cluster:      ClusterSource,
 				NameSpace:    namespace,
-				dataSize:     len(docWithTs.doc),
+				dataSize:     int32(len(docWithTs.doc)),
 				DstTimestamp: option.Some(docWithTs.ts),
 			},
 		)
@@ -732,7 +732,7 @@ func (verifier *Verifier) compareOneDocument(srcClientDoc, dstClientDoc bson.Raw
 			Details:   Mismatch + fmt.Sprintf(" : Document %s has fields in different order", srcClientDoc.Lookup("_id")),
 			Cluster:   ClusterTarget,
 			NameSpace: namespace,
-			dataSize:  dataSize,
+			dataSize:  int32(dataSize),
 		}}, nil
 	}
 	results := mismatchResultsToVerificationResults(mismatch, srcClientDoc, dstClientDoc, namespace, srcClientDoc.Lookup("_id"), "" /* fieldPrefix */)

internal/verifier/migration_verifier.go

@@ -565,7 +565,7 @@ func (verifier *Verifier) ProcessVerifyTask(ctx context.Context, workerNum int,
 				Int("mismatchesCount", len(problems)).
 				Msg("Discrepancies found. Will recheck in the next generation.")
 
-			dataSizes := make([]int, 0, len(problems))
+			dataSizes := make([]int32, 0, len(problems))
 
 			// This stores all IDs for the next generation to check.
 			// Its length should equal len(mismatches) + len(missingIds).

internal/verifier/migration_verifier_test.go

@@ -941,21 +941,21 @@ func (suite *IntegrationTestSuite) TestFailedVerificationTaskInsertions() {
 		ctx,
 		"foo.bar",
 		mslices.Of(mbson.ToRawValue(42)),
-		[]int{100},
+		[]int32{100},
 	)
 	suite.Require().NoError(err)
 	err = verifier.InsertFailedCompareRecheckDocs(
 		ctx,
 		"foo.bar",
 		mslices.Of(mbson.ToRawValue(43), mbson.ToRawValue(44)),
-		[]int{100, 100},
+		[]int32{100, 100},
 	)
 	suite.Require().NoError(err)
 	err = verifier.InsertFailedCompareRecheckDocs(
 		ctx,
 		"foo.bar2",
 		mslices.Of(mbson.ToRawValue(42)),
-		[]int{100},
+		[]int32{100},
 	)
 	suite.Require().NoError(err)
 	event := ParsedEvent{

internal/verifier/parsed_event.go

@@ -26,16 +26,20 @@ func (pe *ParsedEvent) String() string {
 
 var _ bson.Unmarshaler = &ParsedEvent{}
 
-// UnmarshalBSON implements bson.Unmarshaler. We define this manually to
-// avoid reflection, which can substantially impede performance in “hot”
-// code paths like this.
+// UnmarshalBSON implements bson.Unmarshaler but panics because
+// it’s faster to use UnmarshalFromBSON than bson.Unmarshal().
 func (pe *ParsedEvent) UnmarshalBSON(in []byte) error {
-	els, err := bson.Raw(in).Elements()
-	if err != nil {
-		return errors.Wrapf(err, "parsing elements")
-	}
+	panic("Use UnmarshalFromBSON instead.")
+}
+
+// UnmarshalFromBSON unmarshals from BSON without the overhead of
+// bson.Unmarshal.
+func (pe *ParsedEvent) UnmarshalFromBSON(in []byte) error {
+	for el, err := range mbson.RawElements(in) {
+		if err != nil {
+			return errors.Wrapf(err, "parsing elements")
+		}
 
-	for _, el := range els {
 		key, err := el.KeyErr()
 		if err != nil {
 			return errors.Wrapf(err, "parsing field name")
@@ -59,7 +63,7 @@ func (pe *ParsedEvent) UnmarshalBSON(in []byte) error {
 
 			ns := Namespace{}
 
-			err = bson.Unmarshal(rvDoc, &ns)
+			err = (&ns).UnmarshalFromBSON(rvDoc)
 			if err != nil {
 				return errors.Wrapf(err, "unmarshaling %#q value", key)
 			}