REP-5277 Parallelize persistence of rechecks. (#41)

We receive events and enqueue rechecks in the same thread. When there are 1,000s of rechecks to enqueue at once, though, we need faster throughput than a single thread will allow.

This changeset fixes that by parallelizing the enqueueing of rechecks.

Author: FGasper

internal/verifier/change_stream.go

@@ -90,6 +90,10 @@ func (verifier *Verifier) HandleChangeStreamEvents(ctx context.Context, batch []
 		}
 	}
 
+	verifier.logger.Debug().
+		Int("count", len(docIDs)).
+		Msg("Persisting rechecks for change events.")
+
 	return verifier.insertRecheckDocs(ctx, dbNames, collNames, docIDs, dataSizes)
 }
 

internal/verifier/recheck.go

@@ -2,16 +2,22 @@ package verifier
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/10gen/migration-verifier/internal/retry"
 	"github.com/10gen/migration-verifier/internal/types"
+	"github.com/pkg/errors"
+	"github.com/samber/lo"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"golang.org/x/sync/errgroup"
 )
 
 const (
-	recheckQueue   = "recheckQueue"
-	maxBSONObjSize = 16 * 1024 * 1024
+	recheckQueue                  = "recheckQueue"
+	maxBSONObjSize                = 16 * 1024 * 1024
+	recheckInserterThreadsSoftMax = 100
 )
 
 // RecheckPrimaryKey stores the implicit type of recheck to perform
@@ -42,6 +48,10 @@ func (verifier *Verifier) InsertFailedCompareRecheckDocs(
 		collNames[i] = collName
 	}
 
+	verifier.logger.Debug().
+		Int("count", len(documentIDs)).
+		Msg("Persisting rechecks for mismatched or missing documents.")
+
 	return verifier.insertRecheckDocs(
 		context.Background(),
 		dbNames,
@@ -51,6 +61,22 @@ func (verifier *Verifier) InsertFailedCompareRecheckDocs(
 	)
 }
 
+// This will split the given slice into *roughly* the given number of chunks.
+// It may end up being more or fewer, but it should be pretty close.
+func splitToChunks[T any, Slice ~[]T](elements Slice, numChunks int) []Slice {
+	if numChunks < 1 {
+		panic(fmt.Sprintf("numChunks (%v) should be >=1", numChunks))
+	}
+
+	elsPerChunk := len(elements) / numChunks
+
+	if elsPerChunk == 0 {
+		elsPerChunk = 1
+	}
+
+	return lo.Chunk(elements, elsPerChunk)
+}
+
 func (verifier *Verifier) insertRecheckDocs(
 	ctx context.Context,
 	dbNames []string,
@@ -63,38 +89,78 @@ func (verifier *Verifier) insertRecheckDocs(
 
 	generation, _ := verifier.getGenerationWhileLocked()
 
-	models := []mongo.WriteModel{}
-	for i, documentID := range documentIDs {
-		pk := RecheckPrimaryKey{
-			Generation:     generation,
-			DatabaseName:   dbNames[i],
-			CollectionName: collNames[i],
-			DocumentID:     documentID,
-		}
+	docIDIndexes := lo.Range(len(documentIDs))
+	indexesPerThread := splitToChunks(docIDIndexes, recheckInserterThreadsSoftMax)
 
-		// The filter must exclude DataSize; otherwise, if a failed comparison
-		// and a change event happen on the same document for the same
-		// generation, the 2nd insert will fail because a) its filter won’t
-		// match anything, and b) it’ll try to insert a new document with the
-		// same _id as the one that the 1st insert already created.
-		filterDoc := bson.D{{"_id", pk}}
+	eg, groupCtx := errgroup.WithContext(ctx)
 
-		recheckDoc := RecheckDoc{
-			PrimaryKey: pk,
-			DataSize:   dataSizes[i],
-		}
+	for _, curThreadIndexes := range indexesPerThread {
+		curThreadIndexes := curThreadIndexes
+
+		eg.Go(func() error {
+			models := make([]mongo.WriteModel, len(curThreadIndexes))
+			for m, i := range curThreadIndexes {
+				pk := RecheckPrimaryKey{
+					Generation:     generation,
+					DatabaseName:   dbNames[i],
+					CollectionName: collNames[i],
+					DocumentID:     documentIDs[i],
+				}
+
+				// The filter must exclude DataSize; otherwise, if a failed comparison
+				// and a change event happen on the same document for the same
+				// generation, the 2nd insert will fail because a) its filter won’t
+				// match anything, and b) it’ll try to insert a new document with the
+				// same _id as the one that the 1st insert already created.
+				filterDoc := bson.D{{"_id", pk}}
 
-		models = append(models,
-			mongo.NewReplaceOneModel().
-				SetFilter(filterDoc).SetReplacement(recheckDoc).SetUpsert(true))
+				recheckDoc := RecheckDoc{
+					PrimaryKey: pk,
+					DataSize:   dataSizes[i],
+				}
+
+				models[m] = mongo.NewReplaceOneModel().
+					SetFilter(filterDoc).
+					SetReplacement(recheckDoc).
+					SetUpsert(true)
+			}
+
+			retryer := retry.New(retry.DefaultDurationLimit)
+			err := retryer.RunForTransientErrorsOnly(
+				groupCtx,
+				verifier.logger,
+				func(_ *retry.Info) error {
+					_, err := verifier.verificationDatabase().Collection(recheckQueue).BulkWrite(
+						groupCtx,
+						models,
+						options.BulkWrite().SetOrdered(false),
+					)
+
+					return err
+				},
+			)
+
+			return errors.Wrapf(err, "failed to persist %d recheck(s) for generation %d", len(models), generation)
+		})
 	}
-	_, err := verifier.verificationDatabase().Collection(recheckQueue).BulkWrite(ctx, models)
 
-	if err == nil {
-		verifier.logger.Debug().Msgf("Persisted %d recheck doc(s) for generation %d", len(models), generation)
+	err := eg.Wait()
+
+	if err != nil {
+		return errors.Wrapf(
+			err,
+			"failed to persist %d recheck(s) for generation %d",
+			len(documentIDs),
+			generation,
+		)
 	}
 
-	return err
+	verifier.logger.Debug().
+		Int("generation", generation).
+		Int("count", len(documentIDs)).
+		Msg("Persisted rechecks.")
+
+	return nil
 }
 
 // ClearRecheckDocs deletes the previous generation’s recheck