Replace WaitForQueries with WaitForBlockedQuery

WaitForQueries is currently just checking the database has run the
expected number of queries as the last query from the given number of
connections. The tests that use this helper want to wait for a
query to be waiting on a lock. This ensures there is another
instance of the query actually blocked. This commit changes the function
to block until there is a given query blocked on a lock to better ensure
the correct case is hit. Additionally, the helper is not currently checking
the query is in the same database as the test is running. Parallel tests
running a matching query in a different database can cause this condition
to flake. This commit thus includes and additional check to ensure the query
is running in the same database as the current test.

3 files changed, 17 insertions, 13 deletions

diff --git a/internal/praefect/datastore/glsql/testing.go b/internal/praefect/datastore/glsql/testing.go
index b362a6e65..630562f0d 100644
--- a/internal/praefect/datastore/glsql/testing.go
+++ b/internal/praefect/datastore/glsql/testing.go
@@ -348,21 +348,25 @@ func getDatabaseEnvironment(t testing.TB) map[string]string {
 	return databaseEnv
 }
 
-// WaitForQueries is a helper that waits until a certain number of queries matching the prefix are present in the
-// database. This is useful for ensuring multiple transactions are executing the query when testing concurrent
-// execution.
-func WaitForQueries(ctx context.Context, t testing.TB, db Querier, queryPrefix string, count int) {
+// WaitForBlockedQuery is a helper that waits until a blocked query matching the prefix is present in the
+// database. This is useful for ensuring another transaction is blocking a query when testing concurrent
+// execution of multiple queries.
+func WaitForBlockedQuery(ctx context.Context, t testing.TB, db Querier, queryPrefix string) {
 	t.Helper()
 
 	for {
-		var queriesPresent bool
+		var queryBlocked bool
 		require.NoError(t, db.QueryRowContext(ctx, `
-			SELECT COUNT(*) = $2
-			FROM pg_stat_activity
-			WHERE TRIM(e'\n' FROM query) LIKE $1
-		`, queryPrefix+"%", count).Scan(&queriesPresent))
-
-		if queriesPresent {
+			SELECT EXISTS (
+				SELECT FROM pg_stat_activity
+				WHERE TRIM(e'\n' FROM query) LIKE $1
+				AND state = 'active'
+				AND wait_event_type = 'Lock'
+				AND datname = current_database()
+			)
+		`, queryPrefix+"%").Scan(&queryBlocked))
+
+		if queryBlocked {
 			return
 		}
 
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index 251b99491..0eb44c163 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -214,7 +214,7 @@ func TestRepositoryStore_incrementGenerationConcurrently(t *testing.T) {
 			require.NoError(t, err)
 
 			go func() {
-				glsql.WaitForQueries(ctx, t, db, "WITH updated_replicas AS (", 2)
+				glsql.WaitForBlockedQuery(ctx, t, db, "WITH updated_replicas AS (")
 				firstTx.Commit(t)
 			}()
 
diff --git a/internal/praefect/nodes/health_manager_test.go b/internal/praefect/nodes/health_manager_test.go
index b0a7a7a15..ad5c94cb8 100644
--- a/internal/praefect/nodes/health_manager_test.go
+++ b/internal/praefect/nodes/health_manager_test.go
@@ -601,7 +601,7 @@ func TestHealthManager_orderedWrites(t *testing.T) {
 	}()
 
 	// Wait for tx2 to be blocked on the gitaly-1 lock acquired by tx1
-	glsql.WaitForQueries(ctx, t, db, "INSERT INTO node_status", 1)
+	glsql.WaitForBlockedQuery(ctx, t, db, "INSERT INTO node_status")
 
 	// Ensure tx1 can acquire lock on gitaly-2.
 	require.NoError(t, hm1.updateHealthChecks(ctx, []string{virtualStorage}, []string{"gitaly-2"}, []bool{true}))