DEVPROD-8478 Add context to task queries (part 2)

cloud/ec2.go

@@ -349,7 +349,7 @@ func (m *ec2Manager) spawnOnDemandHost(ctx context.Context, h *host.Host, ec2Set
 		}
 
 		if h.SpawnOptions.SpawnedByTask {
-			detailErr := task.AddHostCreateDetails(h.StartedBy, h.Id, h.SpawnOptions.TaskExecutionNumber, err)
+			detailErr := task.AddHostCreateDetails(ctx, h.StartedBy, h.Id, h.SpawnOptions.TaskExecutionNumber, err)
 			grip.Error(message.WrapError(detailErr, message.Fields{
 				"message":       "error adding host create error details",
 				"host_id":       h.Id,

db/db_utils.go

@@ -233,6 +233,50 @@
 	return db.C(collection).Update(query, update)
 }
 
+// Update updates one matching document in the collection.
+func UpdateContext(ctx context.Context, collection string, query interface{}, update interface{}) error {
+	res, err := evergreen.GetEnvironment().DB().Collection(collection).UpdateOne(ctx,
+		query,
+		update,
+	)
+	if err != nil {
+		return errors.Wrapf(err, "updating task")
+	}
+	if res.MatchedCount == 0 {
+		return db.ErrNotFound
+	}
+
+	return nil
+}
+
+func UpdateAllContext(ctx context.Context, collection string, query interface{}, update interface{}) (*db.ChangeInfo, error) {
+	switch query.(type) {
+	case *Q, Q:
+		grip.EmergencyPanic(message.Fields{
+			"message":    "invalid query passed to update all",
+			"cause":      "programmer error",
+			"query":      query,
+			"collection": collection,
+		})
+	case nil:
+		grip.EmergencyPanic(message.Fields{
+			"message":    "nil query passed to update all",
+			"query":      query,
+			"collection": collection,
+		})
+	}
+
+	res, err := evergreen.GetEnvironment().DB().Collection(collection).UpdateMany(ctx,
+		query,
+		update,
+	)
+	if err != nil {
+		return nil, errors.Wrapf(err, "updating task")
+	}
+
+	return &db.ChangeInfo{Updated: int(res.ModifiedCount)}, nil
+}
+
 // UpdateId updates one _id-matching document in the collection.
 func UpdateId(collection string, id, update interface{}) error {
 	session, db, err := GetGlobalSessionFactory().GetSession()

graphql/mutation_resolver.go

@@ -944,7 +944,7 @@ func (r *mutationResolver) OverrideTaskDependencies(ctx context.Context, taskID
 	if t == nil {
 		return nil, ResourceNotFound.Send(ctx, fmt.Sprintf("cannot find task with id %s", taskID))
 	}
-	if err = t.SetOverrideDependencies(currentUser.Username()); err != nil {
+	if err = t.SetOverrideDependencies(ctx, currentUser.Username()); err != nil {
 		return nil, InternalServerError.Send(ctx, fmt.Sprintf("overriding dependencies for task '%s': %s", taskID, err.Error()))
 	}
 	return getAPITaskFromTask(ctx, r.sc.GetURL(), *t)
@@ -1322,7 +1322,7 @@ func (r *mutationResolver) ScheduleUndispatchedBaseTasks(ctx context.Context, ve
 					baseGeneratorTask, _ := generatorTask.FindTaskOnBaseCommit(ctx)
 					// If baseGeneratorTask is nil then it didn't exist on the base task and we can't do anything
 					if baseGeneratorTask != nil && baseGeneratorTask.Status == evergreen.TaskUndispatched {
-						err = baseGeneratorTask.SetGeneratedTasksToActivate(t.BuildVariant, t.DisplayName)
+						err = baseGeneratorTask.SetGeneratedTasksToActivate(ctx, t.BuildVariant, t.DisplayName)
 						if err != nil {
 							return nil, InternalServerError.Send(ctx, fmt.Sprintf("Could not activate generated task: %s", err.Error()))
 						}

graphql/patch_resolver.go

@@ -63,7 +63,7 @@ func (r *patchResolver) Builds(ctx context.Context, obj *restModel.APIPatch) ([]
 // Duration is the resolver for the duration field.
 func (r *patchResolver) Duration(ctx context.Context, obj *restModel.APIPatch) (*PatchDuration, error) {
 	query := db.Query(task.ByVersion(*obj.Id)).WithFields(task.TimeTakenKey, task.StartTimeKey, task.FinishTimeKey, task.DisplayOnlyKey, task.ExecutionKey)
-	tasks, err := task.FindAllFirstExecution(query)
+	tasks, err := task.FindAllFirstExecution(ctx, query)
 	if err != nil {
 		return nil, InternalServerError.Send(ctx, err.Error())
 	}

graphql/query_resolver.go

@@ -670,7 +670,7 @@ func (r *queryResolver) TaskTestSample(ctx context.Context, versionID string, ta
 		if dbTask.Version != versionID && dbTask.ParentPatchID != versionID {
 			return nil, InputValidationError.Send(ctx, fmt.Sprintf("task '%s' does not belong to version '%s'", dbTask.Id, versionID))
 		}
-		taskOpts, err := dbTask.CreateTestResultsTaskOptions()
+		taskOpts, err := dbTask.CreateTestResultsTaskOptions(ctx)
 		if err != nil {
 			return nil, InternalServerError.Send(ctx, fmt.Sprintf("creating test results task options for task '%s': %s", dbTask.Id, err.Error()))
 		}

graphql/task_resolver.go

@@ -227,7 +227,7 @@ func (r *taskResolver) CanSetPriority(ctx context.Context, obj *restModel.APITas
 		return true, nil
 	}
 	if len(obj.ExecutionTasks) != 0 && !evergreen.IsFinishedTaskStatus(utility.FromStringPtr(obj.Status)) {
-		tasks, err := task.FindByExecutionTasksAndMaxExecution(utility.FromStringPtrSlice(obj.ExecutionTasks), obj.Execution)
+		tasks, err := task.FindByExecutionTasksAndMaxExecution(ctx, utility.FromStringPtrSlice(obj.ExecutionTasks), obj.Execution)
 		if err != nil {
 			return false, InternalServerError.Send(ctx, fmt.Sprintf("finding execution tasks for task '%s': %s", *obj.Id, err.Error()))
 		}
@@ -256,7 +256,7 @@ func (r *taskResolver) DependsOn(ctx context.Context, obj *restModel.APITask) ([
 		depIds = append(depIds, dep.TaskId)
 	}
 
-	dependencyTasks, err := task.FindWithFields(task.ByIds(depIds), task.DisplayNameKey, task.StatusKey,
+	dependencyTasks, err := task.FindWithFields(ctx, task.ByIds(depIds), task.DisplayNameKey, task.StatusKey,
 		task.ActivatedKey, task.BuildVariantKey, task.DetailsKey, task.DependsOnKey)
 	if err != nil {
 		return nil, InternalServerError.Send(ctx, fmt.Sprintf("Cannot find dependency tasks for task %s: %s", *obj.Id, err.Error()))
@@ -346,7 +346,7 @@ func (r *taskResolver) ExecutionTasksFull(ctx context.Context, obj *restModel.AP
 	if len(obj.ExecutionTasks) == 0 {
 		return nil, nil
 	}
-	tasks, err := task.FindByExecutionTasksAndMaxExecution(utility.FromStringPtrSlice(obj.ExecutionTasks), obj.Execution)
+	tasks, err := task.FindByExecutionTasksAndMaxExecution(ctx, utility.FromStringPtrSlice(obj.ExecutionTasks), obj.Execution)
 	if err != nil {
 		return nil, InternalServerError.Send(ctx, fmt.Sprintf("finding execution tasks for task '%s': %s", utility.FromStringPtr(obj.Id), err.Error()))
 	}
@@ -387,7 +387,7 @@ func (r *taskResolver) Files(ctx context.Context, obj *restModel.APITask) (*Task
 	fileCount := 0
 
 	if obj.DisplayOnly {
-		execTasks, err := task.Find(task.ByIds(utility.FromStringPtrSlice(obj.ExecutionTasks)))
+		execTasks, err := task.Find(ctx, task.ByIds(utility.FromStringPtrSlice(obj.ExecutionTasks)))
 		if err != nil {
 			return &emptyTaskFiles, ResourceNotFound.Send(ctx, err.Error())
 		}

graphql/util.go

@@ -1008,7 +1008,7 @@ func getBaseTaskTestResultsOptions(ctx context.Context, dbTask *task.Task) ([]te
 	}
 
 	if baseTask != nil && baseTask.ResultsService == dbTask.ResultsService {
-		taskOpts, err = baseTask.CreateTestResultsTaskOptions()
+		taskOpts, err = baseTask.CreateTestResultsTaskOptions(ctx)
 		if err != nil {
 			return nil, InternalServerError.Send(ctx, fmt.Sprintf("Error creating test results task options for base task '%s': %s", baseTask.Id, err))
 		}

graphql/version_resolver.go

@@ -174,7 +174,7 @@ func (r *versionResolver) GeneratedTaskCounts(ctx context.Context, obj *restMode
 	}
 
 	var res []*GeneratedTaskCountResults
-	versionGeneratorTasks, err := task.Find(bson.M{
+	versionGeneratorTasks, err := task.Find(ctx, bson.M{
 		task.VersionKey:      versionID,
 		task.GenerateTaskKey: true,
 	})
@@ -483,7 +483,7 @@ func (r *versionResolver) VersionTiming(ctx context.Context, obj *restModel.APIV
 	if v == nil {
 		return nil, ResourceNotFound.Send(ctx, fmt.Sprintf("finding version '%s'", utility.FromStringPtr(obj.Id)))
 	}
-	timeTaken, makespan, err := v.GetTimeSpent()
+	timeTaken, makespan, err := v.GetTimeSpent(ctx)
 	if err != nil {
 		return nil, InternalServerError.Send(ctx, fmt.Sprintf("getting timing for version '%s': %s", utility.FromStringPtr(obj.Id), err.Error()))
 	}

model/build/build.go

@@ -278,9 +278,9 @@ func (b *Build) MarkFinished(status string, finishTime time.Time) error {
 	)
 }
 
-func (b *Build) GetTimeSpent() (time.Duration, time.Duration, error) {
+func (b *Build) GetTimeSpent(ctx context.Context) (time.Duration, time.Duration, error) {
 	query := db.Query(task.ByBuildId(b.Id)).WithFields(task.TimeTakenKey, task.StartTimeKey, task.FinishTimeKey, task.DisplayOnlyKey, task.ExecutionKey)
-	tasks, err := task.FindAllFirstExecution(query)
+	tasks, err := task.FindAllFirstExecution(ctx, query)
 	if err != nil {
 		return 0, 0, errors.Wrap(err, "getting tasks")
 	}

model/generate.go

@@ -323,7 +323,7 @@ func (g *GeneratedProject) saveNewBuildsAndTasks(ctx context.Context, settings *
 		return errors.Wrapf(err, "validating the number of tasks to be added by '%s'", g.Task.Id)
 	}
 	span.SetAttributes(attribute.Int(numGenerateTasksAttribute, tasksToBeGenerated))
-	if err = g.Task.SetNumGeneratedTasks(tasksToBeGenerated); err != nil {
+	if err = g.Task.SetNumGeneratedTasks(ctx, tasksToBeGenerated); err != nil {
 		return errors.Wrapf(err, "setting number of tasks generated by '%s'", g.Task.Id)
 	}
 
@@ -358,7 +358,7 @@ func (g *GeneratedProject) saveNewBuildsAndTasks(ctx context.Context, settings *
 
 	numActivatedGenerateTasks := len(activatedTasksInExistingBuilds) + len(activatedTasksInNewBuilds)
 	span.SetAttributes(attribute.Int(numActivatedGenerateTasksAttribute, numActivatedGenerateTasks))
-	if err = g.Task.SetNumActivatedGeneratedTasks(numActivatedGenerateTasks); err != nil {
+	if err = g.Task.SetNumActivatedGeneratedTasks(ctx, numActivatedGenerateTasks); err != nil {
 		return errors.Wrapf(err, "setting number of tasks generated and activated by '%s'", g.Task.Id)
 	}
 
@@ -446,7 +446,7 @@ func (g *GeneratedProject) GetNewTasksAndActivationInfo(ctx context.Context, v *
 func (g *GeneratedProject) CheckForCycles(ctx context.Context, v *Version, p *Project, projectRef *ProjectRef) error {
 	ctx, span := tracer.Start(ctx, "check-for-cycles")
 	defer span.End()
-	existingTasksGraph, err := task.VersionDependencyGraph(g.Task.Version, false)
+	existingTasksGraph, err := task.VersionDependencyGraph(ctx, g.Task.Version, false)
 	if err != nil {
 		return errors.Wrapf(err, "creating dependency graph for version '%s'", g.Task.Version)
 	}
@@ -715,7 +715,7 @@ func (g *GeneratedProject) addDependencies(ctx context.Context, newTaskIds []str
 	defer span.End()
 	statuses := []string{evergreen.TaskSucceeded, task.AllStatuses}
 	for _, status := range statuses {
-		if err := g.Task.UpdateDependsOn(status, newTaskIds); err != nil {
+		if err := g.Task.UpdateDependsOn(ctx, status, newTaskIds); err != nil {
 			return errors.Wrapf(err, "updating tasks depending on '%s'", g.Task.Id)
 		}
 	}

model/generate_test.go

@@ -1055,7 +1055,7 @@ func (s *GenerateSuite) TestSaveNewBuildsAndTasksWithBatchtime() {
 	s.NoError(err)
 	s.Require().NotZero(dbExistingBV)
 
-	tasksInExistingBV, err := task.Find(task.ByBuildId(sampleBuild.Id)) // without display
+	tasksInExistingBV, err := task.Find(ctx, task.ByBuildId(sampleBuild.Id)) // without display
 	s.NoError(err)
 	s.Len(tasksInExistingBV, 3)
 	for _, tsk := range tasksInExistingBV {

model/lifecycle.go

@@ -70,7 +70,7 @@ func SetVersionActivation(ctx context.Context, versionId string, active bool, ca
 			return errors.Wrap(err, "getting tasks to activate")
 		}
 		if len(tasksToModify) > 0 {
-			if err = task.ActivateTasks(tasksToModify, time.Now(), false, caller); err != nil {
+			if err = task.ActivateTasks(ctx, tasksToModify, time.Now(), false, caller); err != nil {
 				return errors.Wrap(err, "updating tasks for activation")
 			}
 		}
@@ -85,7 +85,7 @@ func SetVersionActivation(ctx context.Context, versionId string, active bool, ca
 			return errors.Wrap(err, "getting tasks to deactivate")
 		}
 		if len(tasksToModify) > 0 {
-			if err = task.DeactivateTasks(tasksToModify, false, caller); err != nil {
+			if err = task.DeactivateTasks(ctx, tasksToModify, false, caller); err != nil {
 				return errors.Wrap(err, "deactivating tasks")
 			}
 		}
@@ -143,7 +143,7 @@ func setTaskActivationForBuilds(ctx context.Context, buildIds []string, active,
 			return errors.Wrap(err, "getting tasks to activate")
 		}
 		if withDependencies {
-			dependOn, err := task.GetRecursiveDependenciesUp(tasksToActivate, nil)
+			dependOn, err := task.GetRecursiveDependenciesUp(ctx, tasksToActivate, nil)
 			if err != nil {
 				return errors.Wrap(err, "getting recursive dependencies")
 			}
@@ -153,7 +153,7 @@ func setTaskActivationForBuilds(ctx context.Context, buildIds []string, active,
 				}
 			}
 		}
-		if err = task.ActivateTasks(tasksToActivate, time.Now(), withDependencies, caller); err != nil {
+		if err = task.ActivateTasks(ctx, tasksToActivate, time.Now(), withDependencies, caller); err != nil {
 			return errors.Wrap(err, "updating tasks for activation")
 		}
 
@@ -171,7 +171,7 @@ func setTaskActivationForBuilds(ctx context.Context, buildIds []string, active,
 		if err != nil {
 			return errors.Wrap(err, "getting tasks to deactivate")
 		}
-		if err = task.DeactivateTasks(tasks, withDependencies, caller); err != nil {
+		if err = task.DeactivateTasks(ctx, tasks, withDependencies, caller); err != nil {
 			return errors.Wrap(err, "deactivating tasks")
 		}
 	}
@@ -213,7 +213,7 @@ func TryMarkVersionStarted(versionId string, startTime time.Time) error {
 // dependencies that have a lower priority than the one being set for this task
 // will also have their priority increased.
 func SetTaskPriority(ctx context.Context, t task.Task, priority int64, caller string) error {
-	depTasks, err := task.GetRecursiveDependenciesUp([]task.Task{t}, nil)
+	depTasks, err := task.GetRecursiveDependenciesUp(ctx, []task.Task{t}, nil)
 	if err != nil {
 		return errors.Wrap(err, "getting task dependencies")
 	}
@@ -242,7 +242,7 @@ func SetTaskPriority(ctx context.Context, t task.Task, priority int64, caller st
 	for _, taskToUpdate := range tasks {
 		taskIDs = append(taskIDs, taskToUpdate.Id)
 	}
-	_, err = task.UpdateAll(
+	_, err = task.UpdateAll(ctx,
 		bson.M{task.IdKey: bson.M{"$in": taskIDs}},
 		bson.M{"$set": bson.M{task.PriorityKey: priority}},
 	)
@@ -276,7 +276,7 @@ func SetVersionsPriority(ctx context.Context, versionIds []string, priority int6
 }
 
 func setTasksPriority(ctx context.Context, query bson.M, priority int64, caller string) error {
-	_, err := task.UpdateAll(query,
+	_, err := task.UpdateAll(ctx, query,
 		bson.M{"$set": bson.M{task.PriorityKey: priority}},
 	)
 	if err != nil {
@@ -814,7 +814,7 @@ func createTasksForBuild(ctx context.Context, creationInfo TaskCreationInfo) (ta
 		}
 
 		// update existing exec tasks
-		grip.Error(message.WrapError(task.AddDisplayTaskIdToExecTasks(id, execTasksThatNeedParentId), message.Fields{
+		grip.Error(message.WrapError(task.AddDisplayTaskIdToExecTasks(ctx, id, execTasksThatNeedParentId), message.Fields{
 			"message":              "problem adding display task ID to exec tasks",
 			"exec_tasks_to_update": execTasksThatNeedParentId,
 			"display_task_id":      id,
@@ -1029,7 +1029,7 @@ func RecomputeNumDependents(ctx context.Context, t task.Task) error {
 		taskPtrs = append(taskPtrs, &depTasks[i])
 	}
 	query := task.ByVersion(t.Version)
-	_, err = task.UpdateAll(query, bson.M{"$set": bson.M{task.NumDependentsKey: 0}})
+	_, err = task.UpdateAll(ctx, query, bson.M{"$set": bson.M{task.NumDependentsKey: 0}})
 	if err != nil {
 		return errors.Wrap(err, "resetting num dependents")
 	}
@@ -1044,7 +1044,7 @@ func RecomputeNumDependents(ctx context.Context, t task.Task) error {
 	SetNumDependents(taskPtrs)
 	catcher := grip.NewBasicCatcher()
 	for _, t := range taskPtrs {
-		catcher.Add(t.SetNumDependents())
+		catcher.Add(t.SetNumDependents(ctx))
 	}
 
 	return errors.Wrap(catcher.Resolve(), "setting num dependents")
@@ -1658,12 +1658,12 @@ func addNewBuilds(ctx context.Context, creationInfo TaskCreationInfo, existingBu
 		return nil, errors.Wrap(err, "updating version with new build IDs")
 	}
 
-	activatedTaskDependencies, err := task.GetRecursiveDependenciesUp(newActivatedTasks, nil)
+	activatedTaskDependencies, err := task.GetRecursiveDependenciesUp(ctx, newActivatedTasks, nil)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting dependencies for activated tasks")
 	}
 
-	if err = task.ActivateTasks(activatedTaskDependencies, time.Now(), true, evergreen.User); err != nil {
+	if err = task.ActivateTasks(ctx, activatedTaskDependencies, time.Now(), true, evergreen.User); err != nil {
 		return nil, errors.Wrap(err, "activating dependencies for new tasks")
 	}
 
@@ -1787,11 +1787,11 @@ func addNewTasksToExistingBuilds(ctx context.Context, creationInfo TaskCreationI
 		}
 	}
 
-	activatedTaskDependencies, err := task.GetRecursiveDependenciesUp(activatedTasks, nil)
+	activatedTaskDependencies, err := task.GetRecursiveDependenciesUp(ctx, activatedTasks, nil)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting dependencies for activated tasks")
 	}
-	if err = task.ActivateTasks(activatedTaskDependencies, time.Now(), true, evergreen.User); err != nil {
+	if err = task.ActivateTasks(ctx, activatedTaskDependencies, time.Now(), true, evergreen.User); err != nil {
 		return nil, errors.Wrap(err, "activating existing dependencies for new tasks")
 	}