check runner

Author: FGasper

internal/verifier/change_stream_test.go

@@ -246,26 +246,17 @@ func (suite *IntegrationTestSuite) TestEventBeforeWritesOff() {
 
 	verifier := suite.BuildVerifier()
 
-	checkDoneChan := make(chan struct{})
-	checkContinueChan := make(chan struct{})
-
 	db := suite.srcMongoClient.Database(suite.DBNameForTest())
 	coll := db.Collection("mycoll")
 	suite.Require().NoError(
 		db.CreateCollection(ctx, coll.Name()),
 	)
 
 	// start verifier
-	verifierDoneChan := make(chan struct{})
-	go func() {
-		err := verifier.CheckDriver(ctx, nil, checkDoneChan, checkContinueChan)
-		suite.Require().NoError(err)
+	verifierRunner := RunVerifierCheck(suite.Context(), suite.T(), verifier)
 
-		close(verifierDoneChan)
-	}()
-
-	// wait for generation 1
-	<-checkDoneChan
+	// wait for generation 0 to end
+	verifierRunner.AwaitGenerationEnd()
 
 	docsCount := 10_000
 	docs := lo.RepeatBy(docsCount, func(_ int) bson.D { return bson.D{} })
@@ -277,15 +268,7 @@ func (suite *IntegrationTestSuite) TestEventBeforeWritesOff() {
 
 	suite.Require().NoError(verifier.WritesOff(ctx))
 
-	verifierDone := false
-	for !verifierDone {
-		select {
-		case <-verifierDoneChan:
-			verifierDone = true
-		case <-checkDoneChan:
-		case checkContinueChan <- struct{}{}:
-		}
-	}
+	suite.Require().NoError(verifierRunner.Await())
 
 	generation := verifier.generation
 	failedTasks, incompleteTasks, err := FetchFailedAndIncompleteTasks(

internal/verifier/check_runner.go

@@ -0,0 +1,55 @@
+package verifier
+
+import (
+	"context"
+	"testing"
+)
+
+type CheckRunner struct {
+	checkDoneChan        chan error
+	generationDoneChan   chan struct{}
+	doNextGenerationChan chan struct{}
+}
+
+func RunVerifierCheck(ctx context.Context, t *testing.T, verifier *Verifier) *CheckRunner {
+	verifierDoneChan := make(chan error)
+
+	generationDoneChan := make(chan struct{})
+	doNextGenerationChan := make(chan struct{})
+
+	go func() {
+		err := verifier.CheckDriver(ctx, nil, generationDoneChan, doNextGenerationChan)
+		verifierDoneChan <- err
+	}()
+
+	return &CheckRunner{
+		checkDoneChan:        verifierDoneChan,
+		generationDoneChan:   generationDoneChan,
+		doNextGenerationChan: doNextGenerationChan,
+	}
+}
+
+// AwaitGenerationEnd blocks until the check’s current generation ends.
+func (cr *CheckRunner) AwaitGenerationEnd() {
+	<-cr.generationDoneChan
+}
+
+// StartNextGeneration blocks until it can tell the check to start
+// the next generation.
+func (cr *CheckRunner) StartNextGeneration() {
+	cr.doNextGenerationChan <- struct{}{}
+}
+
+// Await will await generations and start new ones until the check
+// finishes. It returns the error that verifier.CheckDriver() returns.
+func (cr *CheckRunner) Await() error {
+	for {
+		select {
+		case err := <-cr.checkDoneChan:
+			return err
+
+		case <-cr.generationDoneChan:
+		case cr.doNextGenerationChan <- struct{}{}:
+		}
+	}
+}

internal/verifier/migration_verifier.go

@@ -254,6 +254,10 @@ func (verifier *Verifier) WritesOff(ctx context.Context) error {
 		verifier.writesOffTimestamp = &finalTs
 
 		verifier.mux.Unlock()
+
+		// This has to happen under the lock because the change stream
+		// might be inserting docs into the recheck queue, which happens
+		// under the lock.
 		verifier.changeStreamWritesOffTsChan <- finalTs
 	} else {
 		verifier.mux.Unlock()

internal/verifier/migration_verifier_test.go

@@ -20,7 +20,6 @@ import (
 	"github.com/10gen/migration-verifier/internal/testutil"
 	"github.com/cespare/permute/v2"
 	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/log"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -29,7 +28,6 @@ import (
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
-	"golang.org/x/sync/errgroup"
 )
 
 func TestIntegration(t *testing.T) {
@@ -1292,7 +1290,7 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 	verifier.SetDstNamespaces([]string{"testDb2.testColl3"})
 	verifier.SetNamespaceMap()
 
-	ctx := context.Background()
+	ctx := suite.Context()
 
 	srcColl := suite.srcMongoClient.Database("testDb1").Collection("testColl1")
 	dstColl := suite.dstMongoClient.Database("testDb2").Collection("testColl3")
@@ -1303,18 +1301,7 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 	_, err = dstColl.InsertOne(ctx, bson.M{"_id": 1, "x": 42})
 	suite.Require().NoError(err)
 
-	checkDoneChan := make(chan struct{})
-	checkContinueChan := make(chan struct{})
-
-	errGroup, errGrpCtx := errgroup.WithContext(context.Background())
-	errGroup.Go(func() error {
-		checkDriverErr := verifier.CheckDriver(errGrpCtx, nil, checkDoneChan, checkContinueChan)
-		// Log this as fatal error so that the test doesn't hang.
-		if checkDriverErr != nil {
-			log.Fatal().Err(checkDriverErr).Msg("check driver error")
-		}
-		return checkDriverErr
-	})
+	runner := RunVerifierCheck(ctx, suite.T(), verifier)
 
 	waitForTasks := func() *VerificationStatus {
 		status, err := verifier.GetVerificationStatus()
@@ -1326,16 +1313,16 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 			suite.T().Logf("TotalTasks is 0 (generation=%d); waiting %s then will run another generation …", verifier.generation, delay)
 
 			time.Sleep(delay)
-			checkContinueChan <- struct{}{}
-			<-checkDoneChan
+			runner.StartNextGeneration()
+			runner.AwaitGenerationEnd()
 			status, err = verifier.GetVerificationStatus()
 			suite.Require().NoError(err)
 		}
 		return status
 	}
 
 	// wait for one generation to finish
-	<-checkDoneChan
+	runner.AwaitGenerationEnd()
 	status := waitForTasks()
 	suite.Require().Equal(VerificationStatus{TotalTasks: 2, FailedTasks: 1, CompletedTasks: 1}, *status)
 
@@ -1344,10 +1331,10 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 	suite.Require().NoError(err)
 
 	// tell check to start the next generation
-	checkContinueChan <- struct{}{}
+	runner.StartNextGeneration()
 
 	// wait for generation to finish
-	<-checkDoneChan
+	runner.AwaitGenerationEnd()
 	status = waitForTasks()
 	// there should be no failures now, since they are are equivalent at this point in time
 	suite.Require().Equal(VerificationStatus{TotalTasks: 1, CompletedTasks: 1}, *status)
@@ -1357,10 +1344,10 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 	suite.Require().NoError(err)
 
 	// tell check to start the next generation
-	checkContinueChan <- struct{}{}
+	runner.StartNextGeneration()
 
 	// wait for one generation to finish
-	<-checkDoneChan
+	runner.AwaitGenerationEnd()
 	status = waitForTasks()
 
 	// there should be a failure from the src insert
@@ -1371,19 +1358,20 @@ func (suite *IntegrationTestSuite) TestGenerationalRechecking() {
 	suite.Require().NoError(err)
 
 	// continue
-	checkContinueChan <- struct{}{}
+	runner.StartNextGeneration()
 
 	// wait for it to finish again, this should be a clean run
-	<-checkDoneChan
+	runner.AwaitGenerationEnd()
 	status = waitForTasks()
 
 	// there should be no failures now, since they are are equivalent at this point in time
 	suite.Assert().Equal(VerificationStatus{TotalTasks: 1, CompletedTasks: 1}, *status)
 
+	// We could just abandon this verifier, but we might as well shut it down
+	// gracefully. That prevents a spurious error in the log from “drop”
+	// change events.
 	suite.Require().NoError(verifier.WritesOff(ctx))
-
-	checkContinueChan <- struct{}{}
-	require.NoError(suite.T(), errGroup.Wait())
+	suite.Require().NoError(runner.Await())
 }
 
 func (suite *IntegrationTestSuite) TestVerifierWithFilter() {