internal/manager/persistence/jobs.go

@@ -855,13 +855,24 @@ func (db *DB) UpdateJobsTaskStatusesConditional(ctx context.Context, job *Job,
 
 // TaskTouchedByWorker marks the task as 'touched' by a worker. This is used for timeout detection.
 func (db *DB) TaskTouchedByWorker(ctx context.Context, t *Task) error {
-	tx := db.gormDB.WithContext(ctx).
+	queries, err := db.queries()
-		Model(t).
+	if err != nil {
-		Select("LastTouchedAt").
+		return err
-		Updates(Task{LastTouchedAt: db.gormDB.NowFunc()})
+	}
-	if err := tx.Error; err != nil {
+
+	now := db.now()
+	err = queries.TaskTouchedByWorker(ctx, sqlc.TaskTouchedByWorkerParams{
+		UpdatedAt:     now,
+		LastTouchedAt: now,
+		ID:            int64(t.ID),
+	})
+	if err != nil {
 		return taskError(err, "saving task 'last touched at'")
 	}
+
+	// Also update the given task, so that it's consistent with the database.
+	t.LastTouchedAt = now.Time
+
 	return nil
 }
 

internal/manager/persistence/sqlc/query_jobs.sql

@@ -148,6 +148,12 @@ UPDATE tasks SET
   worker_id = @worker_id
 WHERE id=@id;
 
+-- name: TaskTouchedByWorker :exec
+UPDATE tasks SET
+  updated_at = @updated_at,
+  last_touched_at = @last_touched_at
+WHERE id=@id;
+
 -- name: JobCountTasksInStatus :one
 -- Fetch number of tasks in the given status, of the given job.
 SELECT count(*) as num_tasks FROM tasks

internal/manager/persistence/sqlc/query_jobs.sql.go

@@ -679,6 +679,24 @@ func (q *Queries) TaskAssignToWorker(ctx context.Context, arg TaskAssignToWorker
 	return err
 }
 
+const taskTouchedByWorker = `-- name: TaskTouchedByWorker :exec
+UPDATE tasks SET
+  updated_at = ?1,
+  last_touched_at = ?2
+WHERE id=?3
+`
+
+type TaskTouchedByWorkerParams struct {
+	UpdatedAt     sql.NullTime
+	LastTouchedAt sql.NullTime
+	ID            int64
+}
+
+func (q *Queries) TaskTouchedByWorker(ctx context.Context, arg TaskTouchedByWorkerParams) error {
+	_, err := q.db.ExecContext(ctx, taskTouchedByWorker, arg.UpdatedAt, arg.LastTouchedAt, arg.ID)
+	return err
+}
+
 const updateJobsTaskStatuses = `-- name: UpdateJobsTaskStatuses :exec
 UPDATE tasks SET
   updated_at = ?1,