final timestsamp in change stream

Author: FGasper

internal/verifier/change_stream.go

@@ -160,15 +160,21 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 			eventsRead++
 		}
 
-		if eventsRead > 0 {
-			verifier.logger.Debug().Int("eventsCount", eventsRead).Msgf("Received a batch of events.")
-			err := verifier.HandleChangeStreamEvents(ctx, changeEventBatch)
-			if err != nil {
-				return false, errors.Wrap(err, "failed to handle change events")
-			}
+		if cs.Err() != nil {
+			return false, errors.Wrap(cs.Err(), "change stream iteration failed")
+		}
+
+		if eventsRead == 0 {
+			return false, nil
 		}
 
-		return eventsRead > 0, errors.Wrap(cs.Err(), "change stream iteration failed")
+		verifier.logger.Debug().Int("eventsCount", eventsRead).Msgf("Received a batch of events.")
+		err := verifier.HandleChangeStreamEvents(ctx, changeEventBatch)
+		if err != nil {
+			return false, errors.Wrap(err, "failed to handle change events")
+		}
+
+		return true, nil
 	}
 
 	for {
@@ -184,15 +190,31 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 		// If the changeStreamEnderChan has a message, the user has indicated that
 		// source writes are ended. This means we should exit rather than continue
 		// reading the change stream since there should be no more events.
-		case <-verifier.changeStreamEnderChan:
+		case finalTs := <-verifier.changeStreamFinalTsChan:
 			verifier.logger.Debug().
-				Msg("Change stream thread received shutdown request.")
+				Interface("finalTimestamp", finalTs).
+				Msg("Change stream thread received final timestamp.")
 
 			changeStreamEnded = true
 
 			// Read all change events until the source reports no events.
 			// (i.e., the `getMore` call returns empty)
 			for {
+				curTs, err := extractTimestampFromResumeToken(cs.ResumeToken())
+				if err != nil {
+					err = errors.Wrap(err, "failed to extract timestamp from change stream's resume token")
+					break
+				}
+
+				if curTs.Compare(finalTs) >= 0 {
+					verifier.logger.Debug().
+						Interface("currentTimestamp", curTs).
+						Interface("finalTimestamp", finalTs).
+						Msg("Change stream has reached the final timestamp. Shutting down.")
+
+					break
+				}
+
 				var gotEvent bool
 				gotEvent, err = readAndHandleOneChangeEventBatch()
 

internal/verifier/change_stream_test.go

@@ -141,7 +141,7 @@ func (suite *MultiSourceVersionTestSuite) TestStartAtTimeNoChanges() {
 	err = verifier.StartChangeStream(ctx)
 	suite.Require().NoError(err)
 	suite.Require().Equal(verifier.srcStartAtTs, origStartTs)
-	verifier.changeStreamEnderChan <- struct{}{}
+	verifier.changeStreamFinalTsChan <- *origStartTs
 	<-verifier.changeStreamDoneChan
 	suite.Require().Equal(verifier.srcStartAtTs, origStartTs)
 }
@@ -176,7 +176,7 @@ func (suite *MultiSourceVersionTestSuite) TestStartAtTimeWithChanges() {
 	newStartTs := sess.OperationTime()
 	suite.Require().NotNil(newStartTs)
 	suite.Require().Negative(origStartTs.Compare(*newStartTs))
-	verifier.changeStreamEnderChan <- struct{}{}
+	verifier.changeStreamFinalTsChan <- *newStartTs
 	<-verifier.changeStreamDoneChan
 	suite.Require().Equal(verifier.srcStartAtTs, newStartTs)
 }

internal/verifier/check.go

@@ -42,13 +42,24 @@ func (verifier *Verifier) Check(ctx context.Context, filter map[string]any) {
 	verifier.MaybeStartPeriodicHeapProfileCollection(ctx)
 }
 
-func (verifier *Verifier) waitForChangeStream() error {
+func (verifier *Verifier) waitForChangeStream(ctx context.Context) error {
 	verifier.mux.RLock()
 	csRunning := verifier.changeStreamRunning
 	verifier.mux.RUnlock()
 	if csRunning {
 		verifier.logger.Debug().Msg("Changestream still running, signalling that writes are done and waiting for change stream to exit")
-		verifier.changeStreamEnderChan <- struct{}{}
+
+		finalTs, err := GetClusterTime(
+			ctx,
+			verifier.logger,
+			verifier.srcClient,
+		)
+
+		if err != nil {
+			return errors.Wrapf(err, "failed to fetch source's cluster time")
+		}
+
+		verifier.changeStreamFinalTsChan <- finalTs
 		select {
 		case err := <-verifier.changeStreamErrChan:
 			verifier.logger.Warn().Err(err).
@@ -242,7 +253,7 @@ func (verifier *Verifier) CheckDriver(ctx context.Context, filter map[string]any
 			// It's necessary to wait for the change stream to finish before incrementing the
 			// generation number, or the last changes will not be checked.
 			verifier.mux.Unlock()
-			err := verifier.waitForChangeStream()
+			err := verifier.waitForChangeStream(ctx)
 			if err != nil {
 				return err
 			}

internal/verifier/migration_verifier.go

@@ -123,12 +123,12 @@ type Verifier struct {
 	metaDBName    string
 	srcStartAtTs  *primitive.Timestamp
 
-	mux                   sync.RWMutex
-	changeStreamRunning   bool
-	changeStreamEnderChan chan struct{}
-	changeStreamErrChan   chan error
-	changeStreamDoneChan  chan struct{}
-	lastChangeEventTime   *primitive.Timestamp
+	mux                     sync.RWMutex
+	changeStreamRunning     bool
+	changeStreamFinalTsChan chan primitive.Timestamp
+	changeStreamErrChan     chan error
+	changeStreamDoneChan    chan struct{}
+	lastChangeEventTime     *primitive.Timestamp
 
 	readConcernSetting ReadConcernSetting
 
@@ -188,15 +188,15 @@ func NewVerifier(settings VerifierSettings) *Verifier {
 	}
 
 	return &Verifier{
-		phase:                 Idle,
-		numWorkers:            NumWorkers,
-		readPreference:        readpref.Primary(),
-		partitionSizeInBytes:  400 * 1024 * 1024,
-		failureDisplaySize:    DefaultFailureDisplaySize,
-		changeStreamEnderChan: make(chan struct{}),
-		changeStreamErrChan:   make(chan error),
-		changeStreamDoneChan:  make(chan struct{}),
-		readConcernSetting:    readConcern,
+		phase:                   Idle,
+		numWorkers:              NumWorkers,
+		readPreference:          readpref.Primary(),
+		partitionSizeInBytes:    400 * 1024 * 1024,
+		failureDisplaySize:      DefaultFailureDisplaySize,
+		changeStreamFinalTsChan: make(chan primitive.Timestamp),
+		changeStreamErrChan:     make(chan error),
+		changeStreamDoneChan:    make(chan struct{}),
+		readConcernSetting:      readConcern,
 
 		// This will get recreated once gen0 starts, but we want it
 		// here in case the change streams gets an event before then.