Parallelize insertion of recheck tasks (#166)

This reduces the time between generations by parallelizing the writes with the iteration of rechecks.

Author: FGasper

internal/verifier/migration_verifier.go

@@ -627,11 +627,6 @@ func (verifier *Verifier) ProcessVerifyTask(ctx context.Context, workerNum int,
 	return nil
 }
 
-const (
-	maxRecheckIDsBytes = 1024 * 1024 // 1 MiB
-	maxRecheckIDs      = 10_000
-)
-
 func (verifier *Verifier) logChunkInfo(ctx context.Context, namespaceAndUUID *uuidutil.NamespaceAndUUID) {
 	// Only log full chunk info in debug mode
 	debugMsg := verifier.logger.Debug()

internal/verifier/recheck.go

@@ -20,13 +20,23 @@ import (
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
+	"golang.org/x/exp/slices"
 )
 
 const (
 	recheckBatchByteLimit  = 1024 * 1024
 	recheckBatchCountLimit = 1000
 
 	recheckQueueCollectionNameBase = "recheckQueue"
+
+	maxTasksPerRequest = 500
+	maxInsertSize      = 256 << 10
+
+	// TODO: Try something beneath 256 KiB instead.
+	maxRecheckIDsBytes = 1024 * 1024 // 1 MiB
+
+	// The max # of docs that we want each recheck task’s cursor to return.
+	maxRecheckIDs = 10_000
 )
 
 // InsertFailedCompareRecheckDocs is for inserting RecheckDocs based on failures during Check.
@@ -298,15 +308,20 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 	}
 	defer cursor.Close(ctx)
 
+	var curTasks []bson.Raw
+	var curTasksBytes int
+
+	eg, egCtx := contextplus.ErrGroup(ctx)
+
+	var totalTasks, totalInserts int
 	persistBufferedRechecks := func() error {
 		if len(idAccum) == 0 {
 			return nil
 		}
 
 		namespace := prevDBName + "." + prevCollName
 
-		task, err := verifier.InsertDocumentRecheckTask(
-			ctx,
+		task, err := verifier.createDocumentRecheckTask(
 			idAccum,
 			types.ByteCount(dataSizeAccum),
 			namespace,
@@ -319,6 +334,32 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 				namespace,
 			)
 		}
+		totalTasks++
+
+		taskRaw, err := bson.Marshal(task)
+		if err != nil {
+			return errors.Wrapf(
+				err,
+				"failed to marshal a %d-document recheck task for collection %#q",
+				len(idAccum),
+				namespace,
+			)
+		}
+
+		curTasks = append(curTasks, taskRaw)
+		curTasksBytes += len(taskRaw)
+		if len(curTasks) == maxTasksPerRequest || curTasksBytes >= maxInsertSize {
+			tasksClone := slices.Clone(curTasks)
+			curTasks = curTasks[:0]
+
+			eg.Go(
+				func() error {
+					return verifier.insertDocumentRecheckTasks(egCtx, tasksClone)
+				},
+			)
+
+			totalInserts++
+		}
 
 		verifier.logger.Debug().
 			Any("task", task.PrimaryKey).
@@ -393,11 +434,29 @@ func (verifier *Verifier) GenerateRecheckTasks(ctx context.Context) error {
 	}
 
 	err = persistBufferedRechecks()
+	if err != nil {
+		return err
+	}
+
+	if len(curTasks) > 0 {
+		eg.Go(
+			func() error {
+				return verifier.insertDocumentRecheckTasks(egCtx, curTasks)
+			},
+		)
+	}
+
+	err = eg.Wait()
+	if err != nil {
+		return errors.Wrapf(err, "persisting document recheck tasks")
+	}
 
-	if err == nil && totalDocs > 0 {
+	if totalDocs > 0 {
 		verifier.logger.Info().
 			Int("generation", generation).
 			Int64("totalDocs", int64(totalDocs)).
+			Int("tasks", totalTasks).
+			Int("insertRequests", totalInserts).
 			Str("totalData", reportutils.FmtBytes(totalRecheckData)).
 			Stringer("timeElapsed", time.Since(startTime)).
 			Msg("Scheduled documents for recheck in the new generation.")

internal/verifier/verification_task.go

@@ -199,8 +199,7 @@ func (verifier *Verifier) InsertPartitionVerificationTask(
 	return &task, err
 }
 
-func (verifier *Verifier) InsertDocumentRecheckTask(
-	ctx context.Context,
+func (verifier *Verifier) createDocumentRecheckTask(
 	ids []any,
 	dataSize types.ByteCount,
 	srcNamespace string,
@@ -214,7 +213,7 @@ func (verifier *Verifier) InsertDocumentRecheckTask(
 		}
 	}
 
-	task := VerificationTask{
+	return &VerificationTask{
 		PrimaryKey: bson.NewObjectID(),
 		Generation: verifier.generation,
 		Ids:        ids,
@@ -226,20 +225,24 @@ func (verifier *Verifier) InsertDocumentRecheckTask(
 		},
 		SourceDocumentCount: types.DocumentCount(len(ids)),
 		SourceByteCount:     dataSize,
-	}
+	}, nil
+}
 
+func (verifier *Verifier) insertDocumentRecheckTasks(
+	ctx context.Context,
+	tasks []bson.Raw,
+) error {
 	err := retry.New().WithCallback(
 		func(ctx context.Context, _ *retry.FuncInfo) error {
-			_, err := verifier.verificationTaskCollection().InsertOne(ctx, &task)
+			_, err := verifier.verificationTaskCollection().InsertMany(ctx, tasks)
 
 			return err
 		},
-		"persisting recheck task for namespace %#q (%d document(s))",
-		task.QueryFilter.Namespace,
-		len(ids),
+		"persisting %d recheck tasks",
+		len(tasks),
 	).Run(ctx, verifier.logger)
 
-	return &task, err
+	return err
 }
 
 func (verifier *Verifier) FindNextVerifyTaskAndUpdate(