wait for change event handler

Author: tdq45gj

internal/verifier/change_stream.go

@@ -3,6 +3,7 @@ package verifier
 import (
 	"context"
 	"fmt"
+	"golang.org/x/sync/errgroup"
 	"time"
 
 	"github.com/10gen/migration-verifier/internal/keystring"
@@ -82,7 +83,6 @@ func (verifier *Verifier) initializeChangeStreamReaders() {
 		watcherClient:               verifier.srcClient,
 		buildInfo:                   *verifier.srcBuildInfo,
 		changeStreamRunning:         false,
-		ChangeEventBatchChan:        make(chan []ParsedEvent),
 		ChangeStreamWritesOffTsChan: make(chan primitive.Timestamp),
 		ChangeStreamErrChan:         make(chan error),
 		ChangeStreamDoneChan:        make(chan struct{}),
@@ -95,28 +95,31 @@ func (verifier *Verifier) initializeChangeStreamReaders() {
 		watcherClient:               verifier.dstClient,
 		buildInfo:                   *verifier.dstBuildInfo,
 		changeStreamRunning:         false,
-		ChangeEventBatchChan:        make(chan []ParsedEvent),
 		ChangeStreamWritesOffTsChan: make(chan primitive.Timestamp),
 		ChangeStreamErrChan:         make(chan error),
 		ChangeStreamDoneChan:        make(chan struct{}),
 	}
 }
 
 // StartChangeEventHandler starts a goroutine that handles change event batches from the reader.
-func (verifier *Verifier) StartChangeEventHandler(ctx context.Context, reader *ChangeStreamReader) {
-	go func() {
+func (verifier *Verifier) StartChangeEventHandler(ctx context.Context, reader *ChangeStreamReader, errGroup *errgroup.Group) {
+	errGroup.Go(func() error {
 		for {
 			select {
 			case <-ctx.Done():
-				return
-			case batch := <-reader.ChangeEventBatchChan:
+				return ctx.Err()
+			case batch, more := <-reader.ChangeEventBatchChan:
+				if !more {
+					// Change stream reader has closed the event batch channel because it has finished.
+					return nil
+				}
 				err := verifier.HandleChangeStreamEvents(ctx, batch, reader.readerType)
 				if err != nil {
-					reader.ChangeStreamErrChan <- err
+					return err
 				}
 			}
 		}
-	}()
+	})
 }
 
 // HandleChangeStreamEvents performs the necessary work for change stream events after receiving a batch.
@@ -298,6 +301,8 @@ func (csr *ChangeStreamReader) iterateChangeStream(
 ) error {
 	var lastPersistedTime time.Time
 
+	defer close(csr.ChangeEventBatchChan)
+
 	persistResumeTokenIfNeeded := func() error {
 		if time.Since(lastPersistedTime) <= minChangeStreamPersistInterval {
 			return nil
@@ -328,7 +333,7 @@ func (csr *ChangeStreamReader) iterateChangeStream(
 		case writesOffTs := <-csr.ChangeStreamWritesOffTsChan:
 			csr.logger.Debug().
 				Interface("writesOffTimestamp", writesOffTs).
-				Msg("Change stream thread received writesOff timestamp. Finalizing change stream.")
+				Msgf("%s thread received writesOff timestamp. Finalizing change stream.", csr)
 
 			gotwritesOffTimestamp = true
 
@@ -477,6 +482,11 @@ func (csr *ChangeStreamReader) StartChangeStream(ctx context.Context) error {
 	// there's no chance of "nonsense" like both channels returning a payload.
 	initialCreateResultChan := make(chan mo.Result[primitive.Timestamp])
 
+	// Change stream reader closes ChangeEventBatchChan each time after
+	// finish reading change events.
+	// It needs to be re-initialized when a change stream starts.
+	csr.ChangeEventBatchChan = make(chan []ParsedEvent)
+
 	go func() {
 		retryer := retry.New(retry.DefaultDurationLimit)
 		retryer = retryer.WithErrorCodes(util.CursorKilled)

internal/verifier/change_stream_test.go

@@ -2,6 +2,7 @@ package verifier
 
 import (
 	"context"
+	"golang.org/x/sync/errgroup"
 	"strings"
 	"time"
 
@@ -54,7 +55,7 @@ func (suite *IntegrationTestSuite) TestChangeStreamResumability() {
 		verifier1 := suite.BuildVerifier()
 		ctx, cancel := context.WithCancel(suite.Context())
 		defer cancel()
-		verifier1.StartChangeEventHandler(ctx, verifier1.srcChangeStreamReader)
+		verifier1.StartChangeEventHandler(ctx, verifier1.srcChangeStreamReader, &errgroup.Group{})
 		err := verifier1.srcChangeStreamReader.StartChangeStream(ctx)
 		suite.Require().NoError(err)
 	}()
@@ -80,7 +81,7 @@ func (suite *IntegrationTestSuite) TestChangeStreamResumability() {
 
 	newTime := suite.getClusterTime(ctx, suite.srcMongoClient)
 
-	verifier2.StartChangeEventHandler(ctx, verifier2.srcChangeStreamReader)
+	verifier2.StartChangeEventHandler(ctx, verifier2.srcChangeStreamReader, &errgroup.Group{})
 	err = verifier2.srcChangeStreamReader.StartChangeStream(ctx)
 	suite.Require().NoError(err)
 
@@ -155,7 +156,7 @@ func (suite *IntegrationTestSuite) TestStartAtTimeNoChanges() {
 	suite.Require().NoError(err)
 	origStartTs := sess.OperationTime()
 	suite.Require().NotNil(origStartTs)
-	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader)
+	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader, &errgroup.Group{})
 	err = verifier.srcChangeStreamReader.StartChangeStream(ctx)
 	suite.Require().NoError(err)
 	suite.Require().Equal(verifier.srcChangeStreamReader.startAtTs, origStartTs)
@@ -176,7 +177,7 @@ func (suite *IntegrationTestSuite) TestStartAtTimeWithChanges() {
 
 	origSessionTime := sess.OperationTime()
 	suite.Require().NotNil(origSessionTime)
-	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader)
+	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader, &errgroup.Group{})
 	err = verifier.srcChangeStreamReader.StartChangeStream(ctx)
 	suite.Require().NoError(err)
 
@@ -228,7 +229,7 @@ func (suite *IntegrationTestSuite) TestNoStartAtTime() {
 	suite.Require().NoError(err)
 	origStartTs := sess.OperationTime()
 	suite.Require().NotNil(origStartTs)
-	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader)
+	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader, &errgroup.Group{})
 	err = verifier.srcChangeStreamReader.StartChangeStream(ctx)
 	suite.Require().NoError(err)
 	suite.Require().NotNil(verifier.srcChangeStreamReader.startAtTs)
@@ -248,7 +249,7 @@ func (suite *IntegrationTestSuite) TestWithChangeEventsBatching() {
 
 	verifier := suite.BuildVerifier()
 
-	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader)
+	verifier.StartChangeEventHandler(ctx, verifier.srcChangeStreamReader, &errgroup.Group{})
 	suite.Require().NoError(verifier.srcChangeStreamReader.StartChangeStream(ctx))
 
 	_, err := coll1.InsertOne(ctx, bson.D{{"_id", 1}})
@@ -389,6 +390,7 @@ func (suite *IntegrationTestSuite) testInsertsBeforeWritesOff(docsCount int) {
 		lo.ToAnySlice(docs),
 	)
 	suite.Require().NoError(err)
+	//fmt.Println(fmt.Sprintf("src cluster time %v", suite.getClusterTime(ctx, suite.srcMongoClient)))
 
 	suite.Require().NoError(verifier.WritesOff(ctx))
 
@@ -472,7 +474,7 @@ func (suite *IntegrationTestSuite) TestRecheckDocsWithDstChangeEvents() {
 	verifier.SetDstNamespaces([]string{dstDBName + ".dstColl1", dstDBName + ".dstColl2"})
 	verifier.SetNamespaceMap()
 
-	verifier.StartChangeEventHandler(ctx, verifier.dstChangeStreamReader)
+	verifier.StartChangeEventHandler(ctx, verifier.dstChangeStreamReader, &errgroup.Group{})
 	suite.Require().NoError(verifier.dstChangeStreamReader.StartChangeStream(ctx))
 
 	_, err := coll1.InsertOne(ctx, bson.D{{"_id", 1}})

internal/verifier/check.go

@@ -207,6 +207,7 @@ func (verifier *Verifier) CheckDriver(ctx context.Context, filter map[string]any
 		verifier.phase = Idle
 	}()
 
+	ceHandlerGroup := &errgroup.Group{}
 	for _, csReader := range []*ChangeStreamReader{verifier.srcChangeStreamReader, verifier.dstChangeStreamReader} {
 		if csReader.changeStreamRunning {
 			verifier.logger.Debug().Msgf("Check: %s already running.", csReader)
@@ -217,7 +218,7 @@ func (verifier *Verifier) CheckDriver(ctx context.Context, filter map[string]any
 			if err != nil {
 				return errors.Wrapf(err, "failed to start %s", csReader)
 			}
-			verifier.StartChangeEventHandler(ctx, csReader)
+			verifier.StartChangeEventHandler(ctx, csReader, ceHandlerGroup)
 		}
 	}
 
@@ -293,6 +294,7 @@ func (verifier *Verifier) CheckDriver(ctx context.Context, filter map[string]any
 			if err = verifier.waitForChangeStream(ctx, verifier.dstChangeStreamReader); err != nil {
 				return err
 			}
+			ceHandlerGroup.Wait()
 			verifier.mux.Lock()
 			verifier.lastGeneration = true
 		}