more context

Author: FGasper

internal/verifier/check.go

@@ -353,7 +353,7 @@ func (verifier *Verifier) CreateInitialTasks(ctx context.Context) error {
 			)
 		}
 	}
-	isPrimary, err := verifier.CheckIsPrimary()
+	isPrimary, err := verifier.CheckIsPrimary(ctx)
 	if err != nil {
 		return err
 	}
@@ -380,7 +380,7 @@ func (verifier *Verifier) CreateInitialTasks(ctx context.Context) error {
 
 	verifier.gen0PendingCollectionTasks.Store(int32(len(verifier.srcNamespaces)))
 
-	err = verifier.UpdatePrimaryTaskComplete()
+	err = verifier.UpdatePrimaryTaskComplete(ctx)
 	if err != nil {
 		return err
 	}

internal/verifier/reset_test.go

@@ -15,7 +15,7 @@ func (suite *IntegrationTestSuite) TestResetPrimaryTask() {
 
 	verifier := suite.BuildVerifier()
 
-	created, err := verifier.CheckIsPrimary()
+	created, err := verifier.CheckIsPrimary(ctx)
 	suite.Require().NoError(err)
 	suite.Require().True(created)
 
@@ -45,11 +45,11 @@ func (suite *IntegrationTestSuite) TestResetNonPrimaryTasks() {
 	verifier := suite.BuildVerifier()
 
 	// Create a primary task, and set it to complete.
-	created, err := verifier.CheckIsPrimary()
+	created, err := verifier.CheckIsPrimary(ctx)
 	suite.Require().NoError(err)
 	suite.Require().True(created)
 
-	suite.Require().NoError(verifier.UpdatePrimaryTaskComplete())
+	suite.Require().NoError(verifier.UpdatePrimaryTaskComplete(ctx))
 
 	ns1 := "foo.bar"
 	ns2 := "qux.quux"

internal/verifier/verification_task.go

@@ -90,6 +90,7 @@ type VerificationRange struct {
 }
 
 func (verifier *Verifier) insertCollectionVerificationTask(
+	ctx context.Context,
 	srcNamespace string,
 	generation int) (*VerificationTask, error) {
 
@@ -120,22 +121,22 @@ func (verifier *Verifier) insertCollectionVerificationTask(
 			To:        dstNamespace,
 		},
 	}
-	_, err := verifier.verificationTaskCollection().InsertOne(context.Background(), verificationTask)
+	_, err := verifier.verificationTaskCollection().InsertOne(ctx, verificationTask)
 	return &verificationTask, err
 }
 
 func (verifier *Verifier) InsertCollectionVerificationTask(
 	ctx context.Context,
 	srcNamespace string,
 ) (*VerificationTask, error) {
-	return verifier.insertCollectionVerificationTask(srcNamespace, verifier.generation)
+	return verifier.insertCollectionVerificationTask(ctx, srcNamespace, verifier.generation)
 }
 
 func (verifier *Verifier) InsertFailedCollectionVerificationTask(
 	ctx context.Context,
 	srcNamespace string,
 ) (*VerificationTask, error) {
-	return verifier.insertCollectionVerificationTask(srcNamespace, verifier.generation+1)
+	return verifier.insertCollectionVerificationTask(ctx, srcNamespace, verifier.generation+1)
 }
 
 func (verifier *Verifier) InsertPartitionVerificationTask(
@@ -157,7 +158,7 @@ func (verifier *Verifier) InsertPartitionVerificationTask(
 			To:        dstNamespace,
 		},
 	}
-	_, err := verifier.verificationTaskCollection().InsertOne(context.Background(), verificationTask)
+	_, err := verifier.verificationTaskCollection().InsertOne(ctx, verificationTask)
 	return &verificationTask, err
 }
 
@@ -247,7 +248,7 @@ func (verifier *Verifier) UpdateVerificationTask(ctx context.Context, task *Veri
 	return err
 }
 
-func (verifier *Verifier) CheckIsPrimary() (bool, error) {
+func (verifier *Verifier) CheckIsPrimary(ctx context.Context) (bool, error) {
 	ownerSetId := primitive.NewObjectID()
 	filter := bson.M{"type": verificationTaskPrimary}
 	opts := options.Update()
@@ -259,7 +260,7 @@ func (verifier *Verifier) CheckIsPrimary() (bool, error) {
 			"status": verificationTaskAdded,
 		},
 	}
-	result, err := verifier.verificationTaskCollection().UpdateOne(context.Background(), filter, update, opts)
+	result, err := verifier.verificationTaskCollection().UpdateOne(ctx, filter, update, opts)
 	if err != nil {
 		return false, err
 	}
@@ -268,8 +269,7 @@ func (verifier *Verifier) CheckIsPrimary() (bool, error) {
 	return isPrimary, nil
 }
 
-func (verifier *Verifier) UpdatePrimaryTaskComplete() error {
-	var ctx = context.Background()
+func (verifier *Verifier) UpdatePrimaryTaskComplete(ctx context.Context) error {
 	updateFields := bson.M{
 		"$set": bson.M{
 			"status": verificationTaskCompleted,