datastore: Fix detection of context cancellation in health update

When the context for `StartHealthUpdate()` got canceled, we do not want
to return an error but instead just stop the health updating loop and
return. We need to treat context cancellation specially when executing
SQL queries, though, as the pq driver will not return `context.Canceled`
or `context.DeadlineExceeded` errors, but instead a custom error code
"query_canceled".

Fix detection of context cancellation by checking for any `pq.Error`s
with the mentioned error code.

2 files changed, 11 insertions, 3 deletions

diff --git a/changelogs/unreleased/pks-datastore-test-race.yml b/changelogs/unreleased/pks-datastore-test-race.yml
new file mode 100644
index 000000000..8068d3b53
--- /dev/null
+++ b/changelogs/unreleased/pks-datastore-test-race.yml
@@ -0,0 +1,5 @@
+---
+title: Fix detection of context cancellation for health updater
+merge_request: 2397
+author:
+type: fixed
diff --git a/internal/praefect/datastore/queue.go b/internal/praefect/datastore/queue.go
index 8d1d6e941..0d7083de4 100644
--- a/internal/praefect/datastore/queue.go
+++ b/internal/praefect/datastore/queue.go
@@ -430,10 +430,13 @@ func (rq PostgresReplicationEventQueue) StartHealthUpdate(ctx context.Context, t
 		case <-trigger:
 			res, err := rq.qc.ExecContext(ctx, query, jobIDs, lockIDs)
 			if err != nil {
-				if !(errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) {
-					return err
+				if pqError, ok := err.(*pq.Error); ok && pqError.Code.Name() == "query_canceled" {
+					return nil
 				}
-				return nil
+				if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+					return nil
+				}
+				return err
 			}
 
 			affected, err := res.RowsAffected()