From ca13d7733b8a2f0f7113b579a23be617e6accaf9 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 13 Mar 2024 10:49:05 +0100
Subject: [PATCH v37 1/2] Hopefully make cancel test more reliable

The newly introduced cancel test in libpq_pipeline was flaky. It's not
completely clear why, but one option is that the check for "active" was
actually seeing the active state for the previous query. This change
should address any such race condition by first waiting until the
connection is reported as idle.
---
 .../modules/libpq_pipeline/libpq_pipeline.c   | 62 ++++++++++++-------
 1 file changed, 39 insertions(+), 23 deletions(-)

diff --git a/src/test/modules/libpq_pipeline/libpq_pipeline.c b/src/test/modules/libpq_pipeline/libpq_pipeline.c
index 1fe15ee8899..1d1549e7a1d 100644
--- a/src/test/modules/libpq_pipeline/libpq_pipeline.c
+++ b/src/test/modules/libpq_pipeline/libpq_pipeline.c
@@ -114,48 +114,34 @@ confirm_query_canceled_impl(int line, PGconn *conn)
 		PQconsumeInput(conn);
 }
 
-#define send_cancellable_query(conn, monitorConn) \
-	send_cancellable_query_impl(__LINE__, conn, monitorConn)
 static void
-send_cancellable_query_impl(int line, PGconn *conn, PGconn *monitorConn)
+wait_for_connection_state(int line, PGconn *conn, char *state, PGconn *monitorConn)
 {
-	const char *env_wait;
-	const Oid	paramTypes[1] = {INT4OID};
+	const Oid	paramTypes[] = {INT4OID, TEXTOID};
 	int			procpid = PQbackendPID(conn);
 
-	env_wait = getenv("PG_TEST_TIMEOUT_DEFAULT");
-	if (env_wait == NULL)
-		env_wait = "180";
-
-	if (PQsendQueryParams(conn, "SELECT pg_sleep($1)", 1, paramTypes,
-						  &env_wait, NULL, NULL, 0) != 1)
-		pg_fatal_impl(line, "failed to send query: %s", PQerrorMessage(conn));
-
-	/*
-	 * Wait until the query is actually running. Otherwise sending a
-	 * cancellation request might not cancel the query due to race conditions.
-	 */
 	while (true)
 	{
 		char	   *value;
 		PGresult   *res;
-		const char *paramValues[1];
+		const char *paramValues[2];
 		char		pidval[16];
 
 		snprintf(pidval, 16, "%d", procpid);
 		paramValues[0] = pidval;
+		paramValues[1] = state;
 
 		res = PQexecParams(monitorConn,
 						   "SELECT count(*) FROM pg_stat_activity WHERE "
-						   "pid = $1 AND state = 'active'",
-						   1, NULL, paramValues, NULL, NULL, 1);
+						   "pid = $1 AND state = $2",
+						   2, paramTypes, paramValues, NULL, NULL, 1);
 
 		if (PQresultStatus(res) != PGRES_TUPLES_OK)
-			pg_fatal("could not query pg_stat_activity: %s", PQerrorMessage(monitorConn));
+			pg_fatal_impl(line, "could not query pg_stat_activity: %s", PQerrorMessage(monitorConn));
 		if (PQntuples(res) != 1)
-			pg_fatal("unexpected number of rows received: %d", PQntuples(res));
+			pg_fatal_impl(line, "unexpected number of rows received: %d", PQntuples(res));
 		if (PQnfields(res) != 1)
-			pg_fatal("unexpected number of columns received: %d", PQnfields(res));
+			pg_fatal_impl(line, "unexpected number of columns received: %d", PQnfields(res));
 		value = PQgetvalue(res, 0, 0);
 		if (*value != '0')
 		{
@@ -169,6 +155,36 @@ send_cancellable_query_impl(int line, PGconn *conn, PGconn *monitorConn)
 	}
 }
 
+#define send_cancellable_query(conn, monitorConn) \
+	send_cancellable_query_impl(__LINE__, conn, monitorConn)
+static void
+send_cancellable_query_impl(int line, PGconn *conn, PGconn *monitorConn)
+{
+	const char *env_wait;
+	const Oid	paramTypes[1] = {INT4OID};
+
+	/*
+	 * Wait for the connection to be idle, so that our check for an active
+	 * connection below is reliable, instead of possibly seeing an outdated
+	 * state.
+	 */
+	wait_for_connection_state(line, conn, "idle", monitorConn);
+
+	env_wait = getenv("PG_TEST_TIMEOUT_DEFAULT");
+	if (env_wait == NULL)
+		env_wait = "180";
+
+	if (PQsendQueryParams(conn, "SELECT pg_sleep($1)", 1, paramTypes,
+						  &env_wait, NULL, NULL, 0) != 1)
+		pg_fatal_impl(line, "failed to send query: %s", PQerrorMessage(conn));
+
+	/*
+	 * Wait for the query to start, because if the query is not running yet
+	 * the cancel request that we send won't have any effect.
+	 */
+	wait_for_connection_state(line, conn, "active", monitorConn);
+}
+
 /*
  * Create a new connection with the same conninfo as the given one.
  */

base-commit: a3da95deee38ee067b0bead639c830eacbe894d5
-- 
2.34.1