Merge branch 'pks-datastore-test-race' into 'master'

Fix context cancellation issues with datastore

See merge request gitlab-org/gitaly!2397

3 files changed, 20 insertions, 5 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()
diff --git a/internal/praefect/datastore/queue_test.go b/internal/praefect/datastore/queue_test.go
index c8be612e8..3ea4de70d 100644
--- a/internal/praefect/datastore/queue_test.go
+++ b/internal/praefect/datastore/queue_test.go
@@ -5,6 +5,7 @@ package datastore
 import (
 	"context"
 	"testing"
+	"sync"
 	"time"
 
 	"github.com/stretchr/testify/assert"
@@ -927,8 +928,12 @@ func TestPostgresReplicationEventQueue_StartHealthUpdate(t *testing.T) {
 	t.Run("triggers all passed in events", func(t *testing.T) {
 		db.TruncateAll(t)
 
+		var wg sync.WaitGroup
 		ctx, cancel := testhelper.Context()
-		defer cancel()
+		defer func() {
+			cancel()
+			wg.Wait()
+		}()
 
 		queue := PostgresReplicationEventQueue{qc: db}
 		events := []ReplicationEvent{eventType1, eventType2, eventType3, eventType4}
@@ -950,11 +955,13 @@ func TestPostgresReplicationEventQueue_StartHealthUpdate(t *testing.T) {
 		initialJobLocks := fetchJobLocks(t, ctx, db)
 
 		trigger := make(chan time.Time, 1)
+		wg.Add(1)
 		go func() {
-			trigger <- time.Time{}
+			defer wg.Done()
 			assert.NoError(t, queue.StartHealthUpdate(ctx, trigger, dequeuedEventsToTrigger))
 		}()
 
+		trigger <- time.Time{}
 		time.Sleep(time.Millisecond) // we should sleep as the processing is too fast and won't give different time
 		trigger <- time.Time{}       // once this consumed we are sure that the previous update has been executed