Optimize BSON for rechecks & mismatches. (#165)

This optimizes a few BSON marshal/unmarshal operations:

- Unmarshaling document rechecks
  - This optimizes generation of recheck tasks, which entails reading & unmarshaling rechecks in series.
- Marshaling VerificationResult and MismatchInfo
  - This optimizes behavior when there are lots of mismatches.

This also includes a new UnmarshalCursor function that works like mongo.Cursor.All but for structs that implement UnmarshalFromBSON (and forbid bson.Unmarshal).

It also fixes 2 intermittently-failing tests:
- TestChangeStream_Resume_NoSkip
- TestRecheckDocsWithDstChangeEvents

Author: FGasper

internal/util/cluster_time.go

@@ -1,23 +1,19 @@
 package util
 
 import (
+	"github.com/10gen/migration-verifier/mbson"
 	"github.com/pkg/errors"
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 )
 
 func GetClusterTimeFromSession(sess *mongo.Session) (bson.Timestamp, error) {
-	ctStruct := struct {
-		ClusterTime struct {
-			ClusterTime bson.Timestamp `bson:"clusterTime"`
-		} `bson:"$clusterTime"`
-	}{}
-
 	clusterTimeRaw := sess.ClusterTime()
-	err := bson.Unmarshal(sess.ClusterTime(), &ctStruct)
+
+	ctrv, err := clusterTimeRaw.LookupErr("$clusterTime", "clusterTime")
 	if err != nil {
-		return bson.Timestamp{}, errors.Wrapf(err, "failed to find clusterTime in session cluster time document (%v)", clusterTimeRaw)
+		return bson.Timestamp{}, errors.Wrapf(err, "finding clusterTime in session cluster time document (%v)", clusterTimeRaw)
 	}
 
-	return ctStruct.ClusterTime.ClusterTime, nil
+	return mbson.CastRawValue[bson.Timestamp](ctrv)
 }

internal/verifier/change_stream_test.go

@@ -25,7 +25,6 @@ import (
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
-	"go.mongodb.org/mongo-driver/v2/mongo/readconcern"
 )
 
 func (suite *IntegrationTestSuite) TestChangeStreamFilter_NoNamespaces() {
@@ -255,12 +254,6 @@ func (suite *IntegrationTestSuite) TestChangeStream_Resume_NoSkip() {
 
 	verifier1 := suite.BuildVerifier()
 
-	// Use of linearizable read concern below seems to freeze pre-4.4 servers.
-	srcVersion := verifier1.srcClusterInfo.VersionArray
-	if util.CmpMinorVersions([2]int(srcVersion), [2]int{4, 4}) == -1 {
-		suite.T().Skipf("Source version (%v) is too old for this test.", srcVersion)
-	}
-
 	srcDB := verifier1.srcClient.Database(suite.DBNameForTest())
 	srcColl := srcDB.Collection("coll")
 
@@ -289,9 +282,10 @@ func (suite *IntegrationTestSuite) TestChangeStream_Resume_NoSkip() {
 		"should see a change stream resume token persisted",
 	)
 
-	insertCtx, cancelInserts := contextplus.WithCancelCause(ctx)
-	defer cancelInserts(ctx.Err())
+	stopInserts := make(chan struct{})
 	insertsDone := make(chan struct{})
+
+	var insertedIDs []any
 	go func() {
 		defer func() {
 			close(insertsDone)
@@ -302,21 +296,25 @@ func (suite *IntegrationTestSuite) TestChangeStream_Resume_NoSkip() {
 		)
 		require.NoError(suite.T(), err)
 
-		sessCtx := mongo.NewSessionContext(insertCtx, sess)
+		sessCtx := mongo.NewSessionContext(ctx, sess)
 
-		docID := int32(1)
 		for {
-			_, err := srcColl.InsertOne(
+			select {
+			case <-ctx.Done():
+				return
+			case <-stopInserts:
+				return
+			default:
+			}
+
+			inserted, err := srcColl.InsertOne(
 				sessCtx,
-				bson.D{{"_id", docID}},
+				bson.D{},
 			)
 
-			if err != nil {
-				require.ErrorIs(suite.T(), err, context.Canceled)
-				return
-			}
+			require.NoError(suite.T(), err, "should insert")
 
-			docID++
+			insertedIDs = append(insertedIDs, inserted.InsertedID)
 		}
 	}()
 
@@ -340,58 +338,43 @@ func (suite *IntegrationTestSuite) TestChangeStream_Resume_NoSkip() {
 	verifier2 := suite.BuildVerifier()
 	suite.startSrcChangeStreamReaderAndHandler(ctx, verifier2)
 
-	cancelInserts(fmt.Errorf("verifier2 started"))
+	close(stopInserts)
 	<-insertsDone
 
-	lastIDRes := srcColl.Database().Collection(
-		srcColl.Name(),
-		options.Collection().SetReadConcern(readconcern.Linearizable()),
-	).FindOne(
-		ctx,
-		bson.D{},
-		options.FindOne().
-			SetSort(bson.D{{"_id", -1}}),
-	)
-	require.NoError(suite.T(), lastIDRes.Err())
-
-	lastDocID := lo.Must(lo.Must(lastIDRes.Raw()).LookupErr("_id")).Int32()
-
+	var foundIDs []any
 	assert.Eventually(
 		suite.T(),
 		func() bool {
 			rechecks := suite.fetchPendingVerifierRechecks(ctx, verifier2)
 
-			return lo.SomeBy(
+			foundIDs = lo.Map(
 				rechecks,
-				func(cur bson.M) bool {
+				func(cur bson.M, _ int) any {
 					id := cur["_id"].(bson.D)
 
 					for _, el := range id {
-						if el.Key != "docID" {
-							continue
+						if el.Key == "docID" {
+							return el.Value
 						}
-
-						return el.Value.(int32) == lastDocID
 					}
 
 					panic(fmt.Sprintf("no docID in _id: %+v", id))
 				},
 			)
+
+			return len(foundIDs) == len(insertedIDs)
 		},
 		time.Minute,
 		100*time.Millisecond,
 		"last-inserted doc shows as recheck",
 	)
 
-	sess := lo.Must(verifier2.verificationDatabase().Client().StartSession())
-	sctx := mongo.NewSessionContext(ctx, sess)
-
-	rechecks := suite.fetchPendingVerifierRechecks(sctx, verifier2)
-	if !assert.EqualValues(suite.T(), lastDocID, len(rechecks), "all source docs should be rechecked") {
-		for _, recheck := range rechecks {
-			suite.T().Logf("found recheck: %v", recheck)
-		}
-	}
+	assert.ElementsMatch(
+		suite.T(),
+		insertedIDs,
+		foundIDs,
+		"all source docs should be rechecked",
+	)
 }
 
 // TestChangeStreamResumability creates a verifier, starts its change stream,
@@ -1069,15 +1052,28 @@ func (suite *IntegrationTestSuite) TestRecheckDocsWithDstChangeEvents() {
 	require.Eventually(
 		suite.T(),
 		func() bool {
-			recheckColl := verifier.getRecheckQueueCollection(1 + verifier.generation)
-			cursor, err := recheckColl.Find(ctx, bson.D{})
-			if errors.Is(err, mongo.ErrNoDocuments) {
+			rechecksM := suite.fetchPendingVerifierRechecks(ctx, verifier)
+
+			if len(rechecksM) < 3 {
 				return false
 			}
 
-			suite.Require().NoError(err)
-			suite.Require().NoError(cursor.All(ctx, &rechecks))
-			return len(rechecks) == 3
+			require.Len(suite.T(), rechecksM, 3)
+
+			rechecks = lo.Map(
+				rechecksM,
+				func(m bson.M, _ int) recheck.Doc {
+					raw, err := bson.Marshal(m)
+					require.NoError(suite.T(), err)
+
+					doc := recheck.Doc{}
+					require.NoError(suite.T(), (&doc).UnmarshalFromBSON(raw))
+
+					return doc
+				},
+			)
+
+			return true
 		},
 		time.Minute,
 		500*time.Millisecond,

internal/verifier/mismatches.go

@@ -2,13 +2,16 @@ package verifier
 
 import (
 	"context"
+	"encoding/binary"
+	"fmt"
 
 	"github.com/10gen/migration-verifier/option"
 	"github.com/pkg/errors"
 	"github.com/samber/lo"
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
 )
 
 const (
@@ -20,6 +23,36 @@ type MismatchInfo struct {
 	Detail VerificationResult
 }
 
+var _ bson.Marshaler = MismatchInfo{}
+
+func (mi MismatchInfo) MarshalBSON() ([]byte, error) {
+	panic("Use MarshalToBSON().")
+}
+
+func (mi MismatchInfo) MarshalToBSON() []byte {
+	detail := mi.Detail.MarshalToBSON()
+
+	bsonLen := 4 + // header
+		1 + 4 + 1 + len(bson.ObjectID{}) + // Task
+		1 + 6 + 1 + len(detail) + // Detail
+		1 // NUL
+
+	buf := make(bson.Raw, 4, bsonLen)
+
+	binary.LittleEndian.PutUint32(buf, uint32(bsonLen))
+
+	buf = bsoncore.AppendObjectIDElement(buf, "task", mi.Task)
+	buf = bsoncore.AppendDocumentElement(buf, "detail", detail)
+
+	buf = append(buf, 0)
+
+	if len(buf) != bsonLen {
+		panic(fmt.Sprintf("%T BSON length is %d but expected %d", mi, len(buf), bsonLen))
+	}
+
+	return buf
+}
+
 func createMismatchesCollection(ctx context.Context, db *mongo.Database) error {
 	_, err := db.Collection(mismatchesCollectionName).Indexes().CreateMany(
 		ctx,
@@ -108,7 +141,7 @@ func recordMismatches(
 				Document: MismatchInfo{
 					Task:   taskID,
 					Detail: r,
-				},
+				}.MarshalToBSON(),
 			}
 		},
 	)

internal/verifier/mismatches_test.go

@@ -0,0 +1,49 @@
+package verifier
+
+import (
+	"testing"
+
+	"github.com/10gen/migration-verifier/mbson"
+	"github.com/10gen/migration-verifier/option"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.mongodb.org/mongo-driver/v2/bson"
+)
+
+func TestMismatchesInfoMarshal(t *testing.T) {
+	results := []VerificationResult{
+		{},
+		{
+			ID: mbson.ToRawValue("hello"),
+		},
+		{
+			SrcTimestamp: option.Some(bson.Timestamp{12, 23}),
+		},
+		{
+			DstTimestamp: option.Some(bson.Timestamp{12, 23}),
+		},
+		{
+			ID:           mbson.ToRawValue("hello"),
+			Field:        "yeah",
+			Cluster:      "hoho",
+			Details:      "aaa",
+			NameSpace:    "xxxxx",
+			SrcTimestamp: option.Some(bson.Timestamp{12, 23}),
+			DstTimestamp: option.Some(bson.Timestamp{23, 24}),
+		},
+	}
+
+	for _, result := range results {
+		mi := MismatchInfo{
+			Task:   bson.NewObjectID(),
+			Detail: result,
+		}
+
+		raw := mi.MarshalToBSON()
+
+		var rt MismatchInfo
+		require.NoError(t, bson.Unmarshal(raw, &rt))
+
+		assert.Equal(t, mi, rt, "should round-trip")
+	}
+}

internal/verifier/recheck.go

@@ -139,10 +139,7 @@ func (verifier *Verifier) insertRecheckDocs(
 			DataSize: dataSizes[i],
 		}
 
-		recheckRaw, err := recheckDoc.MarshalToBSON()
-		if err != nil {
-			return errors.Wrapf(err, "marshaling recheck for %#q", dbName+"."+collNames[i])
-		}
+		recheckRaw := recheckDoc.MarshalToBSON()
 
 		curRechecks = append(
 			curRechecks,
@@ -339,15 +336,12 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 	// subject to a 16MB limit on group size.
 	for cursor.Next(ctx) {
 		var doc recheck.Doc
-		err = cursor.Decode(&doc)
+		err = (&doc).UnmarshalFromBSON(cursor.Current)
 		if err != nil {
 			return err
 		}
 
-		idRaw, err := cursor.Current.LookupErr("_id", "docID")
-		if err != nil {
-			return errors.Wrapf(err, "failed to find docID in enqueued recheck %v", cursor.Current)
-		}
+		idRaw := doc.PrimaryKey.DocumentID
 
 		// We persist rechecks if any of these happen:
 		// - the namespace has changed
@@ -395,7 +389,7 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 
 	err = cursor.Err()
 	if err != nil {
-		return err
+		return errors.Wrapf(err, "reading persisted rechecks")
 	}
 
 	err = persistBufferedRechecks()