add test

Author: tdq45gj

internal/verifier/change_stream.go

@@ -15,12 +15,6 @@ import (
 	"golang.org/x/exp/constraints"
 )
 
-// ChangeEventRecheckBuffer contains rechecks from change events as a map namespace -> document _ids, and
-type ChangeEventRecheckBuffer struct {
-	buf     map[string][]interface{}
-	bufSize map[string]uint64
-}
-
 // ParsedEvent contains the fields of an event that we have parsed from 'bson.Raw'.
 type ParsedEvent struct {
 	ID          interface{}          `bson:"_id"`
@@ -40,8 +34,16 @@ type DocKey struct {
 	ID interface{} `bson:"_id"`
 }
 
+// ChangeEventRecheckBuffer contains rechecks from change events as a map namespace -> document _ids, and
+type ChangeEventRecheckBuffer struct {
+	buf     map[string][]interface{}
+	bufSize map[string]uint64
+}
+
+// minChangeStreamCheckpointInterval is a var instead of a const for tests to overwrite.
+var minChangeStreamCheckpointInterval = time.Second * 10
+
 const (
-	minChangeStreamCheckpointInterval  = time.Second * 10
 	metadataChangeStreamCollectionName = "changeStream"
 )
 
@@ -112,6 +114,14 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 		return nil
 	}
 
+	// Do a final flush of the buffered change event rechecks before the function returns.
+	defer func() {
+		err := verifier.flushAllBufferedChangeEventRechecks(ctx)
+		if err != nil {
+			verifier.changeStreamErrChan <- err
+		}
+	}()
+
 	for {
 		var err error
 

internal/verifier/change_stream_test.go

@@ -2,6 +2,8 @@ package verifier
 
 import (
 	"context"
+	"github.com/rs/zerolog"
+	"strings"
 	"testing"
 	"time"
 
@@ -198,3 +200,69 @@ func (suite *MultiSourceVersionTestSuite) TestNoStartAtTime() {
 	suite.Require().NotNil(verifier.srcStartAtTs)
 	suite.Require().LessOrEqual(origStartTs.Compare(*verifier.srcStartAtTs), 0)
 }
+
+func (suite *MultiSourceVersionTestSuite) TestBatchInsertChangeEventRecheckDocs() {
+	zerolog.SetGlobalLevel(zerolog.DebugLevel)
+
+	verifier := buildVerifier(suite.T(), suite.srcMongoInstance, suite.dstMongoInstance, suite.metaMongoInstance)
+
+	ctx := context.Background()
+	vCtx, cancel := context.WithCancel(ctx)
+
+	// Don't do a checkpoint for this test.
+	origInterval := minChangeStreamCheckpointInterval
+	minChangeStreamCheckpointInterval = 10 * time.Hour
+	defer func() {
+		minChangeStreamCheckpointInterval = origInterval
+	}()
+
+	err := verifier.StartChangeStream(vCtx)
+	suite.Require().NoError(err)
+
+	// A large recheck document should be flushed immediately.
+	_, err = suite.srcMongoClient.Database("testDb").Collection("testColl").InsertOne(
+		ctx,
+		bson.D{{"_id", strings.Repeat("a", 4*1024*1024)}},
+	)
+	suite.Require().NoError(err)
+	require.Eventually(
+		suite.T(),
+		func() bool {
+			return len(suite.fetchVerifierRechecks(ctx, verifier)) == 1
+		},
+		time.Minute,
+		500*time.Millisecond,
+		"the verifier should flush a recheck",
+	)
+	suite.Require().Empty(verifier.changeEventRecheckBuf.buf["testDB.testColl"])
+
+	// A small recheck document should be buffered in-memory.
+	_, err = suite.srcMongoClient.Database("testDb").Collection("testColl").InsertOne(
+		ctx,
+		bson.D{{"_id", 0}},
+	)
+	suite.Require().NoError(err)
+	require.Eventually(
+		suite.T(),
+		func() bool {
+			suite.Require().Len(suite.fetchVerifierRechecks(ctx, verifier), 1)
+			return len(verifier.changeEventRecheckBuf.buf["testDB.testColl"]) == 1
+		},
+		time.Minute,
+		500*time.Millisecond,
+		"the verifier should buffer a recheck",
+	)
+
+	// Any recheck docs remaining in the buffer should be flushed before the change stream reader exits.
+	cancel()
+	require.Eventually(
+		suite.T(),
+		func() bool {
+			return len(suite.fetchVerifierRechecks(ctx, verifier)) == 2
+		},
+		time.Minute,
+		500*time.Millisecond,
+		"the verifier should have flushed all recheck docs",
+	)
+	suite.Require().Empty(verifier.changeEventRecheckBuf.buf["testDB.testColl"])
+}

internal/verifier/migration_verifier.go

@@ -139,7 +139,7 @@ type Verifier struct {
 
 	pprofInterval time.Duration
 
-	changeEventRecheckBuf ChangeEventRecheckBuffer
+	changeEventRecheckBuf *ChangeEventRecheckBuffer
 }
 
 // VerificationStatus holds the Verification Status
@@ -199,6 +199,10 @@ func NewVerifier(settings VerifierSettings) *Verifier {
 		changeStreamErrChan:   make(chan error),
 		changeStreamDoneChan:  make(chan struct{}),
 		readConcernSetting:    readConcern,
+		changeEventRecheckBuf: &ChangeEventRecheckBuffer{
+			buf:     make(map[string][]interface{}),
+			bufSize: make(map[string]uint64),
+		},
 	}
 }
 

internal/verifier/migration_verifier_test.go

@@ -232,7 +232,7 @@ func (suite *MultiMetaVersionTestSuite) TestGetNamespaceStatistics_Recheck() {
 	verifier := buildVerifier(suite.T(), suite.srcMongoInstance, suite.dstMongoInstance, suite.metaMongoInstance)
 
 	suite.Require().NoError(
-		verifier.InsertChangeEventRecheckDoc(
+		verifier.AddAndMaybeFlushChangeEventRecheckDoc(
 			ctx,
 			&ParsedEvent{
 				OpType: "insert",
@@ -245,7 +245,7 @@ func (suite *MultiMetaVersionTestSuite) TestGetNamespaceStatistics_Recheck() {
 	)
 
 	suite.Require().NoError(
-		verifier.InsertChangeEventRecheckDoc(
+		verifier.AddAndMaybeFlushChangeEventRecheckDoc(
 			ctx,
 			&ParsedEvent{
 				ID: bson.M{
@@ -260,6 +260,8 @@ func (suite *MultiMetaVersionTestSuite) TestGetNamespaceStatistics_Recheck() {
 		),
 	)
 
+	suite.Require().NoError(verifier.flushAllBufferedChangeEventRechecks(ctx))
+
 	verifier.generation++
 
 	func() {

internal/verifier/recheck.go

@@ -50,8 +50,9 @@ func (verifier *Verifier) AddAndMaybeFlushChangeEventRecheckDoc(ctx context.Cont
 	}
 	verifier.changeEventRecheckBuf.bufSize[namespace] += uint64(len(namespace) + len(bsonID))
 
-	// Flush all recheck documents once a buffer reaches 5 MB.
-	if verifier.changeEventRecheckBuf.bufSize[changeEvent.Ns.String()] > 5*1024*1024 {
+	// Flush all recheck documents once a buffer reaches 4 MB. It is a conservative threshold
+	// to prevent a recheck task document exceeding 16MB size limit.
+	if verifier.changeEventRecheckBuf.bufSize[namespace] > 4*1024*1024 {
 		if err := verifier.flushChangeEventRechecksForNamespace(ctx, namespace); err != nil {
 			return err
 		}