Felipe's review

Author: tdq45gj

internal/verifier/change_stream.go

@@ -59,11 +59,10 @@ type ChangeStreamReader struct {
 	lastChangeEventTime *primitive.Timestamp
 	logger              *logger.Logger
 	namespaces          []string
-	metaDBName          string
 
-	metaClient    *mongo.Client
+	metaDB        *mongo.Database
 	watcherClient *mongo.Client
-	buildInfo     *util.BuildInfo
+	buildInfo     util.BuildInfo
 
 	changeStreamRunning         bool
 	ChangeEventBatchChan        chan []ParsedEvent
@@ -77,35 +76,29 @@ type ChangeStreamReader struct {
 func (verifier *Verifier) initializeChangeStreamReaders() {
 	verifier.srcChangeStreamReader = &ChangeStreamReader{
 		readerType:                  srcReaderType,
-		lastChangeEventTime:         nil,
 		logger:                      verifier.logger,
 		namespaces:                  verifier.srcNamespaces,
-		metaDBName:                  verifier.metaDBName,
-		metaClient:                  verifier.metaClient,
+		metaDB:                      verifier.metaClient.Database(verifier.metaDBName),
 		watcherClient:               verifier.srcClient,
-		buildInfo:                   verifier.srcBuildInfo,
+		buildInfo:                   *verifier.srcBuildInfo,
 		changeStreamRunning:         false,
 		ChangeEventBatchChan:        make(chan []ParsedEvent),
 		ChangeStreamWritesOffTsChan: make(chan primitive.Timestamp),
 		ChangeStreamErrChan:         make(chan error),
 		ChangeStreamDoneChan:        make(chan struct{}),
-		startAtTs:                   nil,
 	}
 	verifier.dstChangeStreamReader = &ChangeStreamReader{
 		readerType:                  dstReaderType,
-		lastChangeEventTime:         nil,
 		logger:                      verifier.logger,
 		namespaces:                  verifier.dstNamespaces,
-		metaDBName:                  verifier.metaDBName,
-		metaClient:                  verifier.metaClient,
+		metaDB:                      verifier.metaClient.Database(verifier.metaDBName),
 		watcherClient:               verifier.dstClient,
-		buildInfo:                   verifier.dstBuildInfo,
+		buildInfo:                   *verifier.dstBuildInfo,
 		changeStreamRunning:         false,
 		ChangeEventBatchChan:        make(chan []ParsedEvent),
 		ChangeStreamWritesOffTsChan: make(chan primitive.Timestamp),
 		ChangeStreamErrChan:         make(chan error),
 		ChangeStreamDoneChan:        make(chan struct{}),
-		startAtTs:                   nil,
 	}
 }
 
@@ -156,12 +149,13 @@ func (verifier *Verifier) HandleChangeStreamEvents(ctx context.Context, batch []
 			// We need to retrieve the source namespaces if change events are from the destination.
 			switch eventOrigin {
 			case dstReaderType:
-				if len(verifier.dstSrcNsMap) == 0 {
+				if verifier.nsMap.Len() == 0 {
+					// Namespace is not remapped. Source namespace is the same as the destination.
 					srcDBName = changeEvent.Ns.DB
 					srcCollName = changeEvent.Ns.Coll
 				} else {
 					dstNs := fmt.Sprintf("%s.%s", changeEvent.Ns.DB, changeEvent.Ns.Coll)
-					srcNs, exist := verifier.dstSrcNsMap[dstNs]
+					srcNs, exist := verifier.nsMap.GetSrcNamespace(dstNs)
 					if !exist {
 						return errors.Errorf("no source namespace corresponding to the destination namepsace %s", dstNs)
 					}
@@ -214,7 +208,7 @@ func (csr *ChangeStreamReader) GetChangeStreamFilter() (pipeline mongo.Pipeline)
 	if len(csr.namespaces) == 0 {
 		pipeline = mongo.Pipeline{
 			{{"$match", bson.D{
-				{"ns.db", bson.D{{"$ne", csr.metaDBName}}},
+				{"ns.db", bson.D{{"$ne", csr.metaDB.Name()}}},
 			}}},
 		}
 	} else {
@@ -328,16 +322,17 @@ func (csr *ChangeStreamReader) iterateChangeStream(
 			return ctx.Err()
 
 		// If the ChangeStreamEnderChan has a message, the user has indicated that
-		// source and destination writes are ended. This means we should exit rather than continue
-		// reading the change stream since there should be no more events.
+		// source writes are ended and the migration tool is finished / committed.
+		// This means we should exit rather than continue reading the change stream
+		// since there should be no more events.
 		case writesOffTs := <-csr.ChangeStreamWritesOffTsChan:
 			csr.logger.Debug().
 				Interface("writesOffTimestamp", writesOffTs).
 				Msg("Change stream thread received writesOff timestamp. Finalizing change stream.")
 
 			gotwritesOffTimestamp = true
 
-			// Read all change events until the source / destination reports no events.
+			// Read change events until the stream reaches the writesOffTs.
 			// (i.e., the `getMore` call returns empty)
 			for {
 				var curTs primitive.Timestamp
@@ -543,7 +538,7 @@ func addTimestampToLogEvent(ts primitive.Timestamp, event *zerolog.Event) *zerol
 }
 
 func (csr *ChangeStreamReader) getChangeStreamMetadataCollection() *mongo.Collection {
-	return csr.metaClient.Database(csr.metaDBName).Collection(metadataChangeStreamCollectionName)
+	return csr.metaDB.Collection(metadataChangeStreamCollectionName)
 }
 
 func (csr *ChangeStreamReader) loadChangeStreamResumeToken(ctx context.Context) (bson.Raw, error) {

internal/verifier/change_stream_test.go

@@ -463,7 +463,10 @@ func (suite *IntegrationTestSuite) TestCreateForbidden() {
 func (suite *IntegrationTestSuite) TestRecheckDocsWithDstChangeEvents() {
 	ctx := suite.Context()
 
-	db := suite.dstMongoClient.Database("dstDB")
+	srcDBName := suite.DBNameForTest() + "_src"
+	dstDBName := suite.DBNameForTest() + "_dst"
+
+	db := suite.dstMongoClient.Database(dstDBName)
 	coll1 := db.Collection("dstColl1")
 	coll2 := db.Collection("dstColl2")
 
@@ -472,8 +475,8 @@ func (suite *IntegrationTestSuite) TestRecheckDocsWithDstChangeEvents() {
 	}
 
 	verifier := suite.BuildVerifier()
-	verifier.SetSrcNamespaces([]string{"srcDB.srcColl1", "srcDB.srcColl2"})
-	verifier.SetDstNamespaces([]string{"dstDB.dstColl1", "dstDB.dstColl2"})
+	verifier.SetSrcNamespaces([]string{srcDBName + ".srcColl1", srcDBName + ".srcColl2"})
+	verifier.SetDstNamespaces([]string{dstDBName + ".dstColl1", dstDBName + ".dstColl2"})
 	verifier.SetNamespaceMap()
 
 	verifier.StartChangeEventHandler(ctx, verifier.dstChangeStreamReader)
@@ -508,14 +511,14 @@ func (suite *IntegrationTestSuite) TestRecheckDocsWithDstChangeEvents() {
 
 	coll1RecheckCount, coll2RecheckCount := 0, 0
 	for _, recheck := range rechecks {
-		suite.Require().Equal("srcDB", recheck.PrimaryKey.DatabaseName)
-		switch recheck.PrimaryKey.CollectionName {
+		suite.Require().Equal(srcDBName, recheck.PrimaryKey.SrcDatabaseName)
+		switch recheck.PrimaryKey.SrcCollectionName {
 		case "srcColl1":
 			coll1RecheckCount++
 		case "srcColl2":
 			coll2RecheckCount++
 		default:
-			suite.T().Fatalf("unknown collection name: %v", recheck.PrimaryKey.CollectionName)
+			suite.T().Fatalf("unknown collection name: %v", recheck.PrimaryKey.SrcCollectionName)
 		}
 	}
 	suite.Require().Equal(2, coll1RecheckCount)

internal/verifier/check.go

@@ -87,10 +87,10 @@ func (verifier *Verifier) CheckWorker(ctx context.Context) error {
 		select {
 		case err := <-verifier.srcChangeStreamReader.ChangeStreamErrChan:
 			cancel()
-			return errors.Wrapf(err, "got an error from %s", verifier.srcChangeStreamReader)
+			return errors.Wrapf(err, "%s failed", verifier.srcChangeStreamReader)
 		case err := <-verifier.dstChangeStreamReader.ChangeStreamErrChan:
 			cancel()
-			return errors.Wrapf(err, "got an error from %s", verifier.dstChangeStreamReader)
+			return errors.Wrapf(err, "%s failed", verifier.dstChangeStreamReader)
 		case <-ctx.Done():
 			cancel()
 			return nil
@@ -170,16 +170,18 @@ func (verifier *Verifier) CheckDriver(ctx context.Context, filter map[string]any
 		verifier.phase = Idle
 	}()
 
-	if verifier.srcChangeStreamReader.changeStreamRunning {
-		verifier.logger.Debug().Msg("Check: Source change stream already running.")
-	} else {
-		verifier.logger.Debug().Msg("Source change stream not running; starting change stream")
+	for _, csReader := range []*ChangeStreamReader{verifier.srcChangeStreamReader, verifier.dstChangeStreamReader} {
+		if csReader.changeStreamRunning {
+			verifier.logger.Debug().Msgf("Check: %s already running.", csReader)
+		} else {
+			verifier.logger.Debug().Msgf("%s not running; starting change stream", csReader)
 
-		err = verifier.srcChangeStreamReader.StartChangeStream(ctx)
-		if err != nil {
-			return errors.Wrap(err, "failed to start change stream on source")
+			err = csReader.StartChangeStream(ctx)
+			if err != nil {
+				return errors.Wrapf(err, "failed to start %s", csReader)
+			}
+			verifier.StartChangeEventHandler(ctx, csReader)
 		}
-		verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader)
 	}
 
 	if verifier.dstChangeStreamReader.changeStreamRunning {

internal/verifier/migration_verifier.go

@@ -128,8 +128,7 @@ type Verifier struct {
 
 	srcNamespaces []string
 	dstNamespaces []string
-	srcDstNsMap   map[string]string
-	dstSrcNsMap   map[string]string
+	nsMap         *NSMap
 	metaDBName    string
 
 	mux sync.RWMutex
@@ -271,7 +270,7 @@ func (verifier *Verifier) WritesOff(ctx context.Context) error {
 		return errors.Wrapf(err, "failed to fetch destination's cluster time")
 	}
 
-	// This has to happen under the lock because the change stream
+	// This has to happen outside the lock because the change stream
 	// might be inserting docs into the recheck queue, which happens
 	// under the lock.
 	select {
@@ -391,15 +390,8 @@ func (verifier *Verifier) SetDstNamespaces(arg []string) {
 }
 
 func (verifier *Verifier) SetNamespaceMap() {
-	verifier.srcDstNsMap = make(map[string]string)
-	verifier.dstSrcNsMap = make(map[string]string)
-	if len(verifier.dstNamespaces) == 0 {
-		return
-	}
-	for i, ns := range verifier.srcNamespaces {
-		verifier.srcDstNsMap[ns] = verifier.dstNamespaces[i]
-		verifier.dstSrcNsMap[verifier.dstNamespaces[i]] = ns
-	}
+	verifier.nsMap = NewNSMap()
+	verifier.nsMap.PopulateWithNamespaces(verifier.srcNamespaces, verifier.dstNamespaces)
 }
 
 func (verifier *Verifier) SetMetaDBName(arg string) {

internal/verifier/nsmap.go

@@ -0,0 +1,49 @@
+package verifier
+
+type NSMap struct {
+	srcDstNsMap map[string]string
+	dstSrcNsMap map[string]string
+}
+
+func NewNSMap() *NSMap {
+	return &NSMap{
+		srcDstNsMap: make(map[string]string),
+		dstSrcNsMap: make(map[string]string),
+	}
+}
+
+func (nsmap *NSMap) PopulateWithNamespaces(srcNamespaces []string, dstNamespaces []string) {
+	if len(srcNamespaces) != len(dstNamespaces) {
+		panic("source and destination namespaces are not the same length")
+	}
+
+	for i, srcNs := range srcNamespaces {
+		dstNs := dstNamespaces[i]
+		if _, exist := nsmap.srcDstNsMap[srcNs]; exist {
+			panic("another mapping already exists for source namespace " + srcNs)
+		}
+		if _, exist := nsmap.dstSrcNsMap[dstNs]; exist {
+			panic("another mapping already exists for destination namespace " + dstNs)
+		}
+		nsmap.srcDstNsMap[srcNs] = dstNs
+		nsmap.dstSrcNsMap[dstNs] = srcNs
+	}
+}
+
+func (nsmap *NSMap) Len() int {
+	if len(nsmap.srcDstNsMap) != len(nsmap.dstSrcNsMap) {
+		panic("source and destination namespaces are not the same length")
+	}
+
+	return len(nsmap.srcDstNsMap)
+}
+
+func (nsmap *NSMap) GetDstNamespace(srcNamespace string) (string, bool) {
+	ns, ok := nsmap.srcDstNsMap[srcNamespace]
+	return ns, ok
+}
+
+func (nsmap *NSMap) GetSrcNamespace(dstNamespace string) (string, bool) {
+	ns, ok := nsmap.dstSrcNsMap[dstNamespace]
+	return ns, ok
+}

internal/verifier/nsmap_test.go

@@ -0,0 +1,38 @@
+package verifier
+
+import (
+	"github.com/stretchr/testify/suite"
+	"testing"
+)
+
+type UnitTestSuite struct {
+	suite.Suite
+}
+
+func TestUnitTestSuite(t *testing.T) {
+	ts := new(UnitTestSuite)
+	suite.Run(t, ts)
+}
+
+func (s *UnitTestSuite) Test_EmptyNsMap() {
+	nsMap := NewNSMap()
+	srcNamespaces := []string{"srcDB.A", "srcDB.B"}
+	dstNamespaces := []string{"dstDB.B", "dstDB.A"}
+	nsMap.PopulateWithNamespaces(srcNamespaces, dstNamespaces)
+	s.Require().Equal(2, nsMap.Len())
+
+	_, ok := nsMap.GetDstNamespace("non-existent.coll")
+	s.Require().False(ok)
+
+	for i, srcNs := range srcNamespaces {
+		gotNs, ok := nsMap.GetDstNamespace(srcNs)
+		s.Require().True(ok)
+		s.Require().Equal(dstNamespaces[i], gotNs)
+	}
+
+	for i, dstNs := range dstNamespaces {
+		gotNs, ok := nsMap.GetSrcNamespace(dstNs)
+		s.Require().True(ok)
+		s.Require().Equal(srcNamespaces[i], gotNs)
+	}
+}

internal/verifier/recheck.go

@@ -24,18 +24,16 @@ const (
 
 // RecheckPrimaryKey stores the implicit type of recheck to perform
 // Currently, we only handle document mismatches/change stream updates,
-// so DatabaseName, CollectionName, and DocumentID must always be specified.
+// so SrcDatabaseName, SrcCollectionName, and DocumentID must always be specified.
 //
 // NB: Order is important here so that, within a given generation,
 // sorting by _id will guarantee that all rechecks for a given
 // namespace appear consecutively.
-//
-// DatabaseName and CollectionName should be on the source.
 type RecheckPrimaryKey struct {
-	Generation     int         `bson:"generation"`
-	DatabaseName   string      `bson:"db"`
-	CollectionName string      `bson:"coll"`
-	DocumentID     interface{} `bson:"docID"`
+	Generation        int         `bson:"generation"`
+	SrcDatabaseName   string      `bson:"db"`
+	SrcCollectionName string      `bson:"coll"`
+	DocumentID        interface{} `bson:"docID"`
 }
 
 // RecheckDoc stores the necessary information to know which documents must be rechecked.
@@ -109,10 +107,10 @@ func (verifier *Verifier) insertRecheckDocs(
 			models := make([]mongo.WriteModel, len(curThreadIndexes))
 			for m, i := range curThreadIndexes {
 				pk := RecheckPrimaryKey{
-					Generation:     generation,
-					DatabaseName:   dbNames[i],
-					CollectionName: collNames[i],
-					DocumentID:     documentIDs[i],
+					Generation:        generation,
+					SrcDatabaseName:   dbNames[i],
+					SrcCollectionName: collNames[i],
+					DocumentID:        documentIDs[i],
 				}
 
 				// The filter must exclude DataSize; otherwise, if a failed comparison
@@ -302,8 +300,8 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 		// - the buffered document IDs’ size exceeds the per-task maximum
 		// - the buffered documents exceed the partition size
 		//
-		if doc.PrimaryKey.DatabaseName != prevDBName ||
-			doc.PrimaryKey.CollectionName != prevCollName ||
+		if doc.PrimaryKey.SrcDatabaseName != prevDBName ||
+			doc.PrimaryKey.SrcCollectionName != prevCollName ||
 			int64(len(idAccum)) > maxDocsPerTask ||
 			idLenAccum >= maxIdsPerRecheckTask ||
 			dataSizeAccum >= verifier.partitionSizeInBytes {
@@ -313,8 +311,8 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 				return err
 			}
 
-			prevDBName = doc.PrimaryKey.DatabaseName
-			prevCollName = doc.PrimaryKey.CollectionName
+			prevDBName = doc.PrimaryKey.SrcDatabaseName
+			prevCollName = doc.PrimaryKey.SrcCollectionName
 			idLenAccum = 0
 			dataSizeAccum = 0
 			idAccum = idAccum[:0]