From a99e68ec99c01fdad7fafd946c1a501ad1c9894f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Mon, 20 May 2024 21:55:12 +0200
Subject: [PATCH] Manager: Convert TaskTouchedByWorker to sqlc

No functional changes.
---
 internal/manager/persistence/jobs.go          | 21 ++++++++++++++-----
 .../manager/persistence/sqlc/query_jobs.sql   |  6 ++++++
 .../persistence/sqlc/query_jobs.sql.go        | 18 ++++++++++++++++
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/internal/manager/persistence/jobs.go b/internal/manager/persistence/jobs.go
index d0443297..0be9a613 100644
--- a/internal/manager/persistence/jobs.go
+++ b/internal/manager/persistence/jobs.go
@@ -873,13 +873,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).
-		Model(t).
-		Select("LastTouchedAt").
-		Updates(Task{LastTouchedAt: db.gormDB.NowFunc()})
-	if err := tx.Error; err != nil {
+	queries, err := db.queries()
+	if err != nil {
+		return err
+	}
+
+	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
 }
 
diff --git a/internal/manager/persistence/sqlc/query_jobs.sql b/internal/manager/persistence/sqlc/query_jobs.sql
index ff81a65d..5547ffee 100644
--- a/internal/manager/persistence/sqlc/query_jobs.sql
+++ b/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
diff --git a/internal/manager/persistence/sqlc/query_jobs.sql.go b/internal/manager/persistence/sqlc/query_jobs.sql.go
index ba2287f4..85616d77 100644
--- a/internal/manager/persistence/sqlc/query_jobs.sql.go
+++ b/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,