change writesOff to send the finalTs

Author: FGasper

internal/verifier/change_stream.go

@@ -126,11 +126,26 @@ func (verifier *Verifier) GetChangeStreamFilter() []bson.D {
 	return []bson.D{stage}
 }
 
-func (verifier *Verifier) readAndHandleOneChangeEventBatch(ctx context.Context, cs *mongo.ChangeStream) error {
+func (verifier *Verifier) readAndHandleOneChangeEventBatch(
+	ctx context.Context,
+	cs *mongo.ChangeStream,
+	finalTs *primitive.Timestamp,
+) error {
 	eventsRead := 0
 	var changeEventBatch []ParsedEvent
 
 	for hasEventInBatch := true; hasEventInBatch; hasEventInBatch = cs.RemainingBatchLength() > 0 {
+		// Once the change stream reaches the final timestamp we should stop reading.
+		if finalTs != nil {
+			csTimestamp, err := extractTimestampFromResumeToken(cs.ResumeToken())
+			if err != nil {
+				return errors.Wrap(err, "failed to extract timestamp from change stream's resume token")
+			}
+			if !csTimestamp.Before(*finalTs) {
+				break
+			}
+		}
+
 		gotEvent := cs.TryNext(ctx)
 
 		if cs.Err() != nil {
@@ -184,7 +199,7 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 
 	for {
 		var err error
-		var changeStreamEnded bool
+		var gotFinalTimestamp bool
 
 		select {
 
@@ -203,7 +218,7 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 				Interface("finalTimestamp", finalTs).
 				Msg("Change stream thread received final timestamp. Finalizing change stream.")
 
-			changeStreamEnded = true
+			gotFinalTimestamp = true
 
 			// Read all change events until the source reports no events.
 			// (i.e., the `getMore` call returns empty)
@@ -224,15 +239,15 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 					break
 				}
 
-				err = verifier.readAndHandleOneChangeEventBatch(ctx, cs)
+				err = verifier.readAndHandleOneChangeEventBatch(ctx, cs, &finalTs)
 
 				if err != nil {
 					break
 				}
 			}
 
 		default:
-			err = verifier.readAndHandleOneChangeEventBatch(ctx, cs)
+			err = verifier.readAndHandleOneChangeEventBatch(ctx, cs, nil)
 
 			if err == nil {
 				err = persistResumeTokenIfNeeded()
@@ -246,12 +261,12 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 
 			verifier.changeStreamErrChan <- err
 
-			if !changeStreamEnded {
+			if !gotFinalTimestamp {
 				break
 			}
 		}
 
-		if changeStreamEnded {
+		if gotFinalTimestamp {
 			verifier.mux.Lock()
 			verifier.changeStreamRunning = false
 			if verifier.lastChangeEventTime != nil {
@@ -260,7 +275,7 @@ func (verifier *Verifier) iterateChangeStream(ctx context.Context, cs *mongo.Cha
 			verifier.mux.Unlock()
 			// since we have started Recheck, we must signal that we have
 			// finished the change stream changes so that Recheck can continue.
-			verifier.changeStreamDoneChan <- struct{}{}
+			close(verifier.changeStreamDoneChan)
 			break
 		}
 	}

internal/verifier/check.go

@@ -43,34 +43,19 @@ func (verifier *Verifier) Check(ctx context.Context, filter map[string]any) {
 }
 
 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")
-
-		finalTs, err := GetNewClusterTime(
-			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).
-				Msg("Received error from change stream.")
-			return err
-		case <-verifier.changeStreamDoneChan:
-			verifier.logger.Debug().
-				Msg("Received completion signal from change stream.")
-			break
-		}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-verifier.changeStreamErrChan:
+		verifier.logger.Warn().Err(err).
+			Msg("Received error from change stream.")
+		return err
+	case <-verifier.changeStreamDoneChan:
+		verifier.logger.Debug().
+			Msg("Received completion signal from change stream.")
+		break
 	}
+
 	return nil
 }
 

internal/verifier/migration_verifier.go

@@ -127,6 +127,7 @@ type Verifier struct {
 	changeStreamErrChan     chan error
 	changeStreamDoneChan    chan struct{}
 	lastChangeEventTime     *primitive.Timestamp
+	writesOffTimestamp      *primitive.Timestamp
 
 	readConcernSetting ReadConcernSetting
 
@@ -226,13 +227,33 @@ func (verifier *Verifier) SetFailureDisplaySize(size int64) {
 	verifier.failureDisplaySize = size
 }
 
-func (verifier *Verifier) WritesOff(ctx context.Context) {
+func (verifier *Verifier) WritesOff(ctx context.Context) error {
 	verifier.logger.Debug().
 		Msg("WritesOff called.")
 
 	verifier.mux.Lock()
+	defer verifier.mux.Unlock()
 	verifier.writesOff = true
-	verifier.mux.Unlock()
+
+	if verifier.writesOffTimestamp == nil {
+		verifier.logger.Debug().Msg("Change stream still running. Signalling that writes are done.")
+
+		finalTs, err := GetNewClusterTime(
+			ctx,
+			verifier.logger,
+			verifier.srcClient,
+		)
+
+		if err != nil {
+			return errors.Wrapf(err, "failed to fetch source's cluster time")
+		}
+
+		verifier.writesOffTimestamp = &finalTs
+
+		verifier.changeStreamFinalTsChan <- finalTs
+	}
+
+	return nil
 }
 
 func (verifier *Verifier) WritesOn(ctx context.Context) {

internal/verifier/migration_verifier_test.go

@@ -1390,8 +1390,19 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 	// because of the calls to WritesOff
 	status, err = verifier.GetVerificationStatus()
 	suite.Require().NoError(err)
+
 	// there should be a failure from the src insert
-	suite.Require().Equal(VerificationStatus{TotalTasks: 1, FailedTasks: 1}, *status)
+	suite.T().Logf("status: %+v", *status)
+	suite.Assert().Equal(VerificationStatus{TotalTasks: 1, FailedTasks: 1}, *status)
+
+	failedTasks, incompleteTasks, err := FetchFailedAndIncompleteTasks(
+		ctx,
+		verifier.verificationTaskCollection(),
+		verificationTaskVerifyDocuments,
+		verifier.generation,
+	)
+	suite.T().Logf("failed: %+v", failedTasks)
+	suite.T().Logf("incomplete: %+v", incompleteTasks)
 
 	checkContinueChan <- struct{}{}
 	require.NoError(suite.T(), errGroup.Wait())

internal/verifier/web_server.go

@@ -22,7 +22,7 @@ const RequestInProgressErrorDescription = "Another request is currently in progr
 // MigrationVerifierAPI represents the interaction webserver with mongosync
 type MigrationVerifierAPI interface {
 	Check(ctx context.Context, filter map[string]any)
-	WritesOff(ctx context.Context)
+	WritesOff(ctx context.Context) error
 	WritesOn(ctx context.Context)
 	GetProgress(ctx context.Context) (Progress, error)
 }
@@ -214,7 +214,11 @@ func (server *WebServer) writesOffEndpoint(c *gin.Context) {
 		return
 	}
 
-	server.Mapi.WritesOff(context.Background())
+	err := server.Mapi.WritesOff(context.Background())
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{"error": err.Error()})
+		return
+	}
 	successResponse(c)
 }
 

internal/verifier/web_server_test.go

@@ -30,8 +30,8 @@ func NewMockVerifier() *MockVerifier {
 func (verifier *MockVerifier) Check(ctx context.Context, filter map[string]any) {
 	verifier.filter = filter
 }
-func (verifier *MockVerifier) WritesOff(ctx context.Context) {}
-func (verifier *MockVerifier) WritesOn(ctx context.Context)  {}
+func (verifier *MockVerifier) WritesOff(ctx context.Context) error { return nil }
+func (verifier *MockVerifier) WritesOn(ctx context.Context)        {}
 func (verifier *MockVerifier) GetProgress(ctx context.Context) (Progress, error) {
 	return Progress{}, nil
 }