Accelerate change stream handling (#132)

This optimizes change stream handling thus:
1. `fullDocument` is replaced, in v4.4+, with just the size of that BSON document. (v4.2 & prior lack the required `$bsonSize` aggregation operator.)
2. Both change streams can now enqueue rechecks concurrently.

In local testing these changes allowed a verification that consistently lagged the change stream (and eventually fell off the oplog) to keep up.

Author: FGasper

internal/verifier/change_stream.go

@@ -8,6 +8,7 @@ import (
 	"github.com/10gen/migration-verifier/internal/keystring"
 	"github.com/10gen/migration-verifier/internal/logger"
 	"github.com/10gen/migration-verifier/internal/retry"
+	"github.com/10gen/migration-verifier/internal/types"
 	"github.com/10gen/migration-verifier/internal/util"
 	"github.com/10gen/migration-verifier/msync"
 	"github.com/10gen/migration-verifier/option"
@@ -40,12 +41,13 @@ var supportedEventOpTypes = mapset.NewSet(
 
 // ParsedEvent contains the fields of an event that we have parsed from 'bson.Raw'.
 type ParsedEvent struct {
-	ID           any                  `bson:"_id"`
-	OpType       string               `bson:"operationType"`
-	Ns           *Namespace           `bson:"ns,omitempty"`
-	DocKey       DocKey               `bson:"documentKey,omitempty"`
-	FullDocument bson.Raw             `bson:"fullDocument,omitempty"`
-	ClusterTime  *primitive.Timestamp `bson:"clusterTime,omitEmpty"`
+	ID           any                            `bson:"_id"`
+	OpType       string                         `bson:"operationType"`
+	Ns           *Namespace                     `bson:"ns,omitempty"`
+	DocKey       DocKey                         `bson:"documentKey,omitempty"`
+	FullDocument bson.Raw                       `bson:"fullDocument,omitempty"`
+	FullDocLen   option.Option[types.ByteCount] `bson:"_fullDocLen"`
+	ClusterTime  *primitive.Timestamp           `bson:"clusterTime,omitEmpty"`
 }
 
 func (pe *ParsedEvent) String() string {
@@ -244,7 +246,9 @@ func (verifier *Verifier) HandleChangeStreamEvents(ctx context.Context, batch ch
 		collNames[i] = srcCollName
 		docIDs[i] = changeEvent.DocKey.ID
 
-		if changeEvent.FullDocument == nil {
+		if changeEvent.FullDocLen.OrZero() > 0 {
+			dataSizes[i] = int(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
@@ -316,14 +320,38 @@ func (csr *ChangeStreamReader) GetChangeStreamFilter() (pipeline mongo.Pipeline)
 		}
 	}
 
-	return append(
+	pipeline = append(
 		pipeline,
 		bson.D{
 			{"$unset", []string{
 				"updateDescription",
 			}},
 		},
 	)
+
+	if csr.hasBsonSize() {
+		pipeline = append(
+			pipeline,
+			bson.D{
+				{"$addFields", bson.D{
+					{"_fullDocLen", bson.D{{"$bsonSize", "$fullDocument"}}},
+					{"fullDocument", "$$REMOVE"},
+				}},
+			},
+		)
+	}
+
+	return pipeline
+}
+
+func (csr *ChangeStreamReader) hasBsonSize() bool {
+	major := csr.clusterInfo.VersionArray[0]
+
+	if major == 4 {
+		return csr.clusterInfo.VersionArray[1] >= 4
+	}
+
+	return major > 4
 }
 
 // This function reads a single `getMore` response into a slice.
@@ -345,6 +373,7 @@ func (csr *ChangeStreamReader) readAndHandleOneChangeEventBatch(
 
 	latestEvent := option.None[ParsedEvent]()
 
+	var batchTotalBytes int
 	for hasEventInBatch := true; hasEventInBatch; hasEventInBatch = cs.RemainingBatchLength() > 0 {
 		gotEvent := cs.TryNext(ctx)
 
@@ -364,6 +393,8 @@ func (csr *ChangeStreamReader) readAndHandleOneChangeEventBatch(
 			changeEvents = make([]ParsedEvent, batchSize)
 		}
 
+		batchTotalBytes += len(cs.Current)
+
 		if err := cs.Decode(&changeEvents[eventsRead]); err != nil {
 			return errors.Wrapf(err, "failed to decode change event to %T", changeEvents[eventsRead])
 		}
@@ -373,7 +404,8 @@ func (csr *ChangeStreamReader) readAndHandleOneChangeEventBatch(
 			Stringer("changeStream", csr).
 			Any("event", changeEvents[eventsRead]).
 			Int("eventsPreviouslyReadInBatch", eventsRead).
-			Int("batchSize", len(changeEvents)).
+			Int("batchEvents", len(changeEvents)).
+			Int("batchBytes", batchTotalBytes).
 			Msg("Received a change event.")
 
 		opType := changeEvents[eventsRead].OpType

internal/verifier/change_stream_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/10gen/migration-verifier/contextplus"
 	"github.com/10gen/migration-verifier/internal/logger"
 	"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/mslices"
 	"github.com/10gen/migration-verifier/mstrings"
@@ -85,31 +86,67 @@ func (suite *IntegrationTestSuite) TestChangeStreamFilter_NoNamespaces() {
 		"updateDescription",
 		"update description should be filtered out",
 	)
+}
+
+func (suite *IntegrationTestSuite) TestChangeStreamFilter_BsonSize() {
+	ctx := suite.Context()
+
+	verifier := suite.BuildVerifier()
+	if !verifier.srcChangeStreamReader.hasBsonSize() {
+		suite.T().Skip("Need a source version that has $bsonSize")
+	}
 
-	/*
-	   suite.Assert().Contains(
+	srcColl := verifier.srcClient.Database(suite.DBNameForTest()).Collection("coll")
 
-	   	,
-	   	bson.D{
-	   		{"$match", bson.D{{"ns.db", bson.D{{"$ne", metaDBName}}}}},
-	   	},
+	_, err := srcColl.InsertOne(ctx, bson.D{{"_id", 123}})
+	suite.Require().NoError(err)
+
+	verifier.srcChangeStreamReader.namespaces = mslices.Of(FullName(srcColl))
 
-	   )
-	   verifier.srcChangeStreamReader.namespaces = []string{"foo.bar", "foo.baz", "test.car", "test.chaz"}
-	   suite.Assert().Contains(
+	filter := verifier.srcChangeStreamReader.GetChangeStreamFilter()
 
-	   	verifier.srcChangeStreamReader.GetChangeStreamFilter(),
-	   	bson.D{{"$match", bson.D{
-	   		{"$or", []bson.D{
-	   			{{"ns", bson.D{{"db", "foo"}, {"coll", "bar"}}}},
-	   			{{"ns", bson.D{{"db", "foo"}, {"coll", "baz"}}}},
-	   			{{"ns", bson.D{{"db", "test"}, {"coll", "car"}}}},
-	   			{{"ns", bson.D{{"db", "test"}, {"coll", "chaz"}}}},
-	   		}},
-	   	}}},
+	cs, err := suite.srcMongoClient.Watch(
+		ctx,
+		filter,
+		options.ChangeStream().SetFullDocument("updateLookup"),
+	)
+	suite.Require().NoError(err)
+	defer cs.Close(ctx)
+
+	// Now create a large document and verify that the corresponding event
+	// does NOT contain that full document but _does_ contain the document’s
+	// length.
+
+	_, err = srcColl.InsertOne(ctx, bson.D{
+		{"_id", "abc"},
+		{"bigstring", strings.Repeat("x", 10_000)},
+	})
+	suite.Require().NoError(err)
+
+	suite.Require().True(cs.Next(ctx), "should get event")
+
+	suite.Require().Equal(
+		"abc",
+		cs.Current.Lookup("documentKey", "_id").StringValue(),
+		"event should reference expected document",
+	)
+	suite.Assert().Less(len(cs.Current), 10_000, "event should not be large")
+
+	parsed := ParsedEvent{}
+	suite.Require().NoError(cs.Decode(&parsed))
+	suite.Require().Equal("insert", parsed.OpType)
+
+	suite.Require().True(parsed.FullDocLen.IsSome(), "full doc len should be in event")
+	suite.Assert().Greater(parsed.FullDocLen.MustGet(), types.ByteCount(10_000))
+
+	_, err = srcColl.DeleteOne(ctx, bson.D{{"_id", "abc"}})
+	suite.Require().NoError(err)
 
-	   )
-	*/
+	suite.Require().True(cs.Next(ctx), "should get event")
+	parsed = ParsedEvent{}
+	suite.Require().NoError(cs.Decode(&parsed))
+	suite.Require().Equal("delete", parsed.OpType)
+	suite.Require().True(parsed.FullDocLen.IsNone(), "full doc len not in delete")
 }
 
 func (suite *IntegrationTestSuite) TestChangeStreamFilter_WithNamespaces() {

internal/verifier/migration_verifier.go

@@ -446,9 +446,9 @@ func (verifier *Verifier) getGenerationWhileLocked() (int, bool) {
 
 	// As long as no other goroutine has locked the mux this will
 	// usefully panic if the caller neglected the lock.
-	wasUnlocked := verifier.mux.TryRLock()
+	wasUnlocked := verifier.mux.TryLock()
 	if wasUnlocked {
-		verifier.mux.RUnlock()
+		verifier.mux.Unlock()
 		panic("getGenerationWhileLocked() while unlocked")
 	}
 

internal/verifier/recheck.go

@@ -92,8 +92,8 @@ func (verifier *Verifier) insertRecheckDocs(
 	documentIDs []any,
 	dataSizes []int,
 ) error {
-	verifier.mux.Lock()
-	defer verifier.mux.Unlock()
+	verifier.mux.RLock()
+	defer verifier.mux.RUnlock()
 
 	generation, _ := verifier.getGenerationWhileLocked()
 
@@ -105,8 +105,6 @@ func (verifier *Verifier) insertRecheckDocs(
 	genCollection := verifier.getRecheckQueueCollection(generation)
 
 	for _, curThreadIndexes := range indexesPerThread {
-		curThreadIndexes := curThreadIndexes
-
 		eg.Go(func() error {
 			models := make([]mongo.WriteModel, len(curThreadIndexes))
 			for m, i := range curThreadIndexes {
@@ -162,9 +160,7 @@ func (verifier *Verifier) insertRecheckDocs(
 		})
 	}
 
-	err := eg.Wait()
-
-	if err != nil {
+	if err := eg.Wait(); err != nil {
 		return errors.Wrapf(
 			err,
 			"failed to persist %d recheck(s) for generation %d",