v8-0002-Allow-SIGINT-to-cancel-psql-database-reconnection.patch

From e4d61216b822852940631ada65b903b65780fb71 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:53:10 -0600
Subject: [PATCH v8 2/2] Allow SIGINT to cancel psql database reconnections

After installing the SIGINT handler in psql, SIGINT can no longer cancel
database reconnections. For instance, if the user starts a reconnection
and then needs to do some form of interaction (ie psql is polling),
there is no way to cancel the reconnection process currently.

Use PQconnectStartParams() in order to insert a cancel_pressed check
into the polling loop.
---
 leak.supp                        |  2 ++
 src/bin/psql/command.c           | 43 +++++++++++++++++++++++++++++++-
 src/interfaces/libpq/exports.txt |  1 +
 3 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 leak.supp

diff --git a/leak.supp b/leak.supp
new file mode 100644
index 0000000000..f3fc27a1d8
--- /dev/null
+++ b/leak.supp
@@ -0,0 +1,2 @@
+leak:save_ps_display_args
+leak:parse_psql_options
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 5c906e4806..1900657324 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3279,6 +3279,46 @@ param_is_newly_set(const char *old_val, const char *new_val)
 	return false;
 }
 
+static void
+process_connection_state_machine(PGconn *conn)
+{
+	bool		for_read = false;
+
+	while (true)
+	{
+		int			rc;
+		int			sock;
+
+		if (cancel_pressed)
+			break;
+
+		sock = PQsocket(conn);
+		if (sock == -1)
+			break;
+
+		rc = PQsocketPoll(sock, for_read, !for_read, 1);
+		if (rc == -1)
+			return;
+
+		switch (PQconnectPoll(conn))
+		{
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				return;
+			case PGRES_POLLING_READING:
+				for_read = true;
+				continue;
+			case PGRES_POLLING_WRITING:
+				for_read = false;
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+		}
+	}
+
+	pg_unreachable();
+}
+
 /*
  * do_connect -- handler for \connect
  *
@@ -3595,11 +3635,12 @@ do_connect(enum trivalue reuse_previous_specification,
 		values[paramnum] = NULL;
 
 		/* Note we do not want libpq to re-expand the dbname parameter */
-		n_conn = PQconnectdbParams(keywords, values, false);
+		n_conn = PQconnectStartParams(keywords, values, false);
 
 		pg_free(keywords);
 		pg_free(values);
 
+		process_connection_state_machine(n_conn);
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index 088592deb1..7a7a6ddbab 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -193,3 +193,4 @@ PQsendClosePrepared       190
 PQsendClosePortal         191
 PQchangePassword          192
 PQsendPipelineSync        193
+PQsocketPoll              194
-- 
Tristan Partin
Neon (https://neon.tech)