psql not responding to SIGINT upon db reconnection

From b9ccfc3c84a25b8616fd40495954bb6f77788e28 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v2] 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.

Restore the default SIGINT handler while polling in PQconnectdbParams(),
and put the custom handler back afterward.
---
 src/bin/psql/command.c | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 6733f008fd..e40413a229 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -40,6 +40,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3530,8 +3531,11 @@ do_connect(enum trivalue reuse_previous_specification,
 	{
 		const char **keywords = pg_malloc((nconnopts + 1) * sizeof(*keywords));
 		const char **values = pg_malloc((nconnopts + 1) * sizeof(*values));
+		int			rc;
 		int			paramnum = 0;
 		PQconninfoOption *ci;
+		struct sigaction oldact;
+		struct sigaction newact = { 0 };
 
 		/*
 		 * Copy non-default settings into the PQconnectdbParams parameter
@@ -3576,9 +3580,24 @@ do_connect(enum trivalue reuse_previous_specification,
 		keywords[paramnum] = NULL;
 		values[paramnum] = NULL;
 
+		/*
+		 * Restore the default SIGINT behavior while within libpq. Otherwise, we
+		 * can never exit from polling for the database connection. Failure to
+		 * restore is non-fatal.
+		 */
+		newact.sa_handler = SIG_DFL;
+		rc = sigaction(SIGINT, &newact, &oldact);
+
 		/* Note we do not want libpq to re-expand the dbname parameter */
 		n_conn = PQconnectdbParams(keywords, values, false);
 
+		/*
+		 * Only attempt to restore the old handler if we successfully overwrote
+		 * it initially.
+		 */
+		if (rc == 0)
+			sigaction(SIGINT, &oldact, NULL);
+
 		pg_free(keywords);
 		pg_free(values);
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 7f9554944911c77aa1a1900537a91e1e7bd75d93 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v3] 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.

Restore the default SIGINT handler while polling in PQconnectdbParams(),
and put the custom handler back afterward.
---
 src/bin/psql/command.c | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 6733f008fd..11dedbb4c6 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -40,6 +40,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3530,8 +3531,11 @@ do_connect(enum trivalue reuse_previous_specification,
 	{
 		const char **keywords = pg_malloc((nconnopts + 1) * sizeof(*keywords));
 		const char **values = pg_malloc((nconnopts + 1) * sizeof(*values));
+		int			rc;
 		int			paramnum = 0;
 		PQconninfoOption *ci;
+		struct sigaction oldact;
+		struct sigaction newact = { 0 };
 
 		/*
 		 * Copy non-default settings into the PQconnectdbParams parameter
@@ -3576,9 +3580,24 @@ do_connect(enum trivalue reuse_previous_specification,
 		keywords[paramnum] = NULL;
 		values[paramnum] = NULL;
 
+		/*
+		 * Restore the default SIGINT behavior while within libpq. Otherwise,
+		 * SIGINT can never exit from polling for the database connection.
+		 * Failure to restore the default is non-fatal.
+		 */
+		newact.sa_handler = SIG_DFL;
+		rc = sigaction(SIGINT, &newact, &oldact);
+
 		/* Note we do not want libpq to re-expand the dbname parameter */
 		n_conn = PQconnectdbParams(keywords, values, false);
 
+		/*
+		 * Only attempt to restore the old handler if we successfully overrode
+		 * it initially.
+		 */
+		if (rc == 0)
+			sigaction(SIGINT, &oldact, NULL);
+
 		pg_free(keywords);
 		pg_free(values);
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 73a95bc537f205cdac567f3f1860b08cf8323e17 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v4] 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 CancelRequested check
into the polling loop.
---
 src/bin/psql/command.c | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 6733f008fd..bfecc25801 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -40,6 +40,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3577,11 +3578,33 @@ 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);
 
+		while (true) {
+			int		socket;
+
+			if (CancelRequested)
+				break;
+
+			socket = PQsocket(n_conn);
+			if (socket == -1)
+				break;
+
+			switch (PQconnectPoll(n_conn)) {
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				break;
+			case PGRES_POLLING_READING:
+			case PGRES_POLLING_WRITING:
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+			}
+		}
+
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 1b4303df1200c561cf478f9c7037ff940f6cd741 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v4] 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 CancelRequested check
into the polling loop.
---
 src/bin/psql/command.c | 129 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 128 insertions(+), 1 deletion(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 82cc091568..588eed9351 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -11,6 +11,11 @@
 #include <time.h>
 #include <pwd.h>
 #include <utime.h>
+#if HAVE_POLL
+#include <poll.h>
+#else
+#include <sys/select.h>
+#endif
 #ifndef WIN32
 #include <sys/stat.h>			/* for stat() */
 #include <sys/time.h>			/* for setitimer() */
@@ -40,6 +45,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3251,6 +3257,90 @@ copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf)
 	return false;
 }
 
+/*
+ * Check a file descriptor for read and/or write data, possibly waiting.
+ * If neither forRead nor forWrite are set, immediately return a timeout
+ * condition (without waiting).  Return >0 if condition is met, 0
+ * if a timeout occurred, -1 if an error or interrupt occurred.
+ *
+ * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
+ * if end_time is 0 (or indeed, any time before now).
+ */
+static int
+pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+{
+	/* We use poll(2) if available, otherwise select(2) */
+#ifdef HAVE_POLL
+	struct pollfd input_fd;
+	int			timeout_ms;
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	input_fd.fd = sock;
+	input_fd.events = POLLERR;
+	input_fd.revents = 0;
+
+	if (forRead)
+		input_fd.events |= POLLIN;
+	if (forWrite)
+		input_fd.events |= POLLOUT;
+
+	/* Compute appropriate timeout interval */
+	if (end_time == ((time_t) -1))
+		timeout_ms = -1;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout_ms = (end_time - now) * 1000;
+		else
+			timeout_ms = 0;
+	}
+
+	return poll(&input_fd, 1, timeout_ms);
+#else							/* !HAVE_POLL */
+
+	fd_set		input_mask;
+	fd_set		output_mask;
+	fd_set		except_mask;
+	struct timeval timeout;
+	struct timeval *ptr_timeout;
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	FD_ZERO(&input_mask);
+	FD_ZERO(&output_mask);
+	FD_ZERO(&except_mask);
+	if (forRead)
+		FD_SET(sock, &input_mask);
+
+	if (forWrite)
+		FD_SET(sock, &output_mask);
+	FD_SET(sock, &except_mask);
+
+	/* Compute appropriate timeout interval */
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout.tv_sec = end_time - now;
+		else
+			timeout.tv_sec = 0;
+		timeout.tv_usec = 0;
+		ptr_timeout = &timeout;
+	}
+
+	return select(sock + 1, &input_mask, &output_mask,
+				  &except_mask, ptr_timeout);
+#endif							/* HAVE_POLL */
+}
+
 /*
  * Ask the user for a password; 'username' is the username the
  * password is for, if one has been explicitly specified.
@@ -3324,6 +3414,7 @@ do_connect(enum trivalue reuse_previous_specification,
 	bool		keep_password = true;
 	bool		has_connection_string;
 	bool		reuse_previous;
+	bool		for_read;
 
 	has_connection_string = dbname ?
 		recognized_connection_string(dbname) : false;
@@ -3614,11 +3705,47 @@ 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);
 
+		for_read = false;
+		while (true) {
+			int		rc;
+			int		sock;
+
+			if (cancel_pressed)
+				break;
+
+			sock = PQsocket(n_conn);
+			if (sock == -1)
+				break;
+
+			rc = pqSocketPoll(sock, for_read, !for_read, (time_t)-1);
+			Assert(rc != 0); /* Timeout is impossible. */
+			if (rc == -1)
+			{
+				success = false;
+				break;
+			}
+
+			switch (PQconnectPoll(n_conn)) {
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				goto finish;
+			case PGRES_POLLING_READING:
+				for_read = true;
+				continue;
+			case PGRES_POLLING_WRITING:
+				for_read = false;
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+			}
+		}
+
+	finish:
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 4fdccf9211ac78bbd7430488d515646f9e9cce9b Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v5] 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 CancelRequested check
into the polling loop.
---
 meson.build                |   1 +
 src/bin/psql/command.c     | 222 ++++++++++++++++++++++++++++++++++++-
 src/include/pg_config.h.in |   3 +
 src/tools/msvc/Solution.pm |   1 +
 4 files changed, 226 insertions(+), 1 deletion(-)

diff --git a/meson.build b/meson.build
index ee58ee7a06..2d63485c53 100644
--- a/meson.build
+++ b/meson.build
@@ -2440,6 +2440,7 @@ func_checks = [
   ['posix_fadvise'],
   ['posix_fallocate'],
   ['ppoll'],
+  ['pselect'],
   ['pstat'],
   ['pthread_barrier_wait', {'dependencies': [thread_dep]}],
   ['pthread_is_threaded_np', {'dependencies': [thread_dep]}],
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 82cc091568..c325fedb51 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -11,6 +11,11 @@
 #include <time.h>
 #include <pwd.h>
 #include <utime.h>
+#if HAVE_POLL
+#include <poll.h>
+#else
+#include <sys/select.h>
+#endif
 #ifndef WIN32
 #include <sys/stat.h>			/* for stat() */
 #include <sys/time.h>			/* for setitimer() */
@@ -40,6 +45,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3251,6 +3257,169 @@ copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf)
 	return false;
 }
 
+/*
+ * Check a file descriptor for read and/or write data, possibly waiting.
+ * If neither forRead nor forWrite are set, immediately return a timeout
+ * condition (without waiting).  Return >0 if condition is met, 0
+ * if a timeout occurred, -1 if an error or interrupt occurred.
+ *
+ * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
+ * if end_time is 0 (or indeed, any time before now).
+ */
+static int
+pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+{
+	/*
+	 * We use functions in the following order if available:
+	 *   - ppoll(2) OR pselect(2)
+	 *   - poll(2) OR select(2)
+	 */
+#ifdef HAVE_POLL
+	struct pollfd input_fd;
+#ifdef HAVE_PPOLL
+	int				rc;
+	sigset_t		emptyset;
+	sigset_t		blockset;
+	sigset_t		origset;
+	struct timespec timeout;
+	struct timespec *ptr_timeout;
+#else
+	int			timeout_ms;
+#endif
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	input_fd.fd = sock;
+	input_fd.events = POLLERR;
+	input_fd.revents = 0;
+
+	if (forRead)
+		input_fd.events |= POLLIN;
+	if (forWrite)
+		input_fd.events |= POLLOUT;
+
+	/* Compute appropriate timeout interval */
+#ifdef HAVE_PPOLL
+	sigemptyset(&blockset);
+	sigaddset(&blockset, SIGINT);
+	sigprocmask(SIG_BLOCK, &blockset, &origset);
+
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		timeout.tv_sec = end_time;
+		timeout.tv_nsec = 0;
+		ptr_timeout = &timeout;
+	}
+#else
+	if (end_time == ((time_t) -1))
+		timeout_ms = -1;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout_ms = (end_time - now) * 1000;
+		else
+			timeout_ms = 0;
+	}
+#endif
+
+#ifdef HAVE_PPOLL
+	sigemptyset(&emptyset);
+	if (cancel_pressed)
+	{
+		sigprocmask(SIG_SETMASK, &origset, NULL);
+		return 1;
+	}
+
+	rc = ppoll(&input_fd, 1, ptr_timeout, &emptyset);
+	sigprocmask(SIG_SETMASK, &origset, NULL);
+	return rc;
+#else
+	return poll(&input_fd, 1, timeout_ms);
+#endif
+#else							/* !HAVE_POLL */
+
+	fd_set		input_mask;
+	fd_set		output_mask;
+	fd_set		except_mask;
+#ifdef HAVE_PSELECT
+	int 			rc;
+	sigset_t		emptyset;
+	sigset_t		blockset;
+	sigset_t		origset;
+	struct timespec timeout;
+	struct timespec *ptr_timeout;
+#else
+	struct timeval timeout;
+	struct timeval *ptr_timeout;
+#endif
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	FD_ZERO(&input_mask);
+	FD_ZERO(&output_mask);
+	FD_ZERO(&except_mask);
+	if (forRead)
+		FD_SET(sock, &input_mask);
+
+	if (forWrite)
+		FD_SET(sock, &output_mask);
+	FD_SET(sock, &except_mask);
+
+	/* Compute appropriate timeout interval */
+#ifdef HAVE_PSELECT
+	sigemptyset(&blockset);
+	sigaddset(&blockset, SIGINT);
+	sigprocmask(SIG_BLOCK, &blockset, &origset);
+
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		timeout.tv_sec = end_time;
+		timeout.tv_nsec = 0;
+		ptr_timeout = &timeout;
+	}
+#else
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout.tv_sec = end_time - now;
+		else
+			timeout.tv_sec = 0;
+		timeout.tv_usec = 0;
+		ptr_timeout = &timeout;
+	}
+#endif
+
+#ifdef HAVE_PSELECT
+	sigemptyset(&emptyset);
+	if (cancel_pressed)
+	{
+		sigprocmask(SIG_SETMASK, &origset, NULL);
+		return 1;
+	}
+
+	rc = pselect(sock + 1, &input_mask, &output_mask,
+				 &except_mask, ptr_timeout, &emptyset);
+	sigprocmask(SIG_SETMASK, &origset, NULL);
+	return rc;
+#else
+	return select(sock + 1, &input_mask, &output_mask,
+				  &except_mask, ptr_timeout);
+#endif
+#endif							/* HAVE_POLL */
+}
+
 /*
  * Ask the user for a password; 'username' is the username the
  * password is for, if one has been explicitly specified.
@@ -3324,6 +3493,7 @@ do_connect(enum trivalue reuse_previous_specification,
 	bool		keep_password = true;
 	bool		has_connection_string;
 	bool		reuse_previous;
+	bool		for_read;
 
 	has_connection_string = dbname ?
 		recognized_connection_string(dbname) : false;
@@ -3614,11 +3784,61 @@ 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);
 
+		for_read = false;
+		while (true) {
+			int		rc;
+			int		sock;
+			time_t	timeout;
+
+			if (cancel_pressed)
+				break;
+
+			sock = PQsocket(n_conn);
+			if (sock == -1)
+				break;
+
+			/*
+			 * We use ppoll(2)/pselect(2) to account for the race condition in
+			 * which SIGINT is sent after checking cancel_pressed. But on
+			 * platforms that don't have either function, we can just spin the
+			 * CPU a bit polling, so set the timeout to 1 second if we don't
+			 * have the aforementioned functions, otherwise set timeout to
+			 * a negative value indicating we will sit and wait forever.
+			 */
+#if defined(HAVE_PPOLL) || defined(HAVE_PSELECT)
+			timeout = -1;
+#else
+			timeout = 1;
+#endif
+
+			rc = pqSocketPoll(sock, for_read, !for_read, timeout);
+			if (rc == -1)
+			{
+				success = false;
+				break;
+			}
+
+			switch (PQconnectPoll(n_conn)) {
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				goto finish;
+			case PGRES_POLLING_READING:
+				for_read = true;
+				continue;
+			case PGRES_POLLING_WRITING:
+				for_read = false;
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+			}
+		}
+
+	finish:
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
diff --git a/src/include/pg_config.h.in b/src/include/pg_config.h.in
index d8a2985567..f9fd7d0de7 100644
--- a/src/include/pg_config.h.in
+++ b/src/include/pg_config.h.in
@@ -333,6 +333,9 @@
 /* Define to 1 if you have the `ppoll' function. */
 #undef HAVE_PPOLL
 
+/* Define to 1 if you have the `pselect' function. */
+#undef HAVE_PSELECT
+
 /* Define if you have POSIX threads libraries and header files. */
 #undef HAVE_PTHREAD
 
diff --git a/src/tools/msvc/Solution.pm b/src/tools/msvc/Solution.pm
index 98a5b5d872..d035f44f73 100644
--- a/src/tools/msvc/Solution.pm
+++ b/src/tools/msvc/Solution.pm
@@ -308,6 +308,7 @@ sub GenerateFiles
 		HAVE_POSIX_FADVISE => undef,
 		HAVE_POSIX_FALLOCATE => undef,
 		HAVE_PPOLL => undef,
+		HAVE_PSELECT => undef,
 		HAVE_PTHREAD => undef,
 		HAVE_PTHREAD_BARRIER_WAIT => undef,
 		HAVE_PTHREAD_IS_THREADED_NP => undef,
-- 
Tristan Partin
Neon (https://neon.tech)

From b22f286d3733d6d5ec2a924682679f6884b3600c Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v6] 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 CancelRequested check
into the polling loop.
---
 meson.build                |   1 +
 src/bin/psql/command.c     | 224 ++++++++++++++++++++++++++++++++++++-
 src/include/pg_config.h.in |   3 +
 src/tools/msvc/Solution.pm |   1 +
 4 files changed, 228 insertions(+), 1 deletion(-)

diff --git a/meson.build b/meson.build
index ee58ee7a06..2d63485c53 100644
--- a/meson.build
+++ b/meson.build
@@ -2440,6 +2440,7 @@ func_checks = [
   ['posix_fadvise'],
   ['posix_fallocate'],
   ['ppoll'],
+  ['pselect'],
   ['pstat'],
   ['pthread_barrier_wait', {'dependencies': [thread_dep]}],
   ['pthread_is_threaded_np', {'dependencies': [thread_dep]}],
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 82cc091568..e87401b790 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -11,6 +11,11 @@
 #include <time.h>
 #include <pwd.h>
 #include <utime.h>
+#if HAVE_POLL
+#include <poll.h>
+#else
+#include <sys/select.h>
+#endif
 #ifndef WIN32
 #include <sys/stat.h>			/* for stat() */
 #include <sys/time.h>			/* for setitimer() */
@@ -40,6 +45,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3251,6 +3257,169 @@ copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf)
 	return false;
 }
 
+/*
+ * Check a file descriptor for read and/or write data, possibly waiting.
+ * If neither forRead nor forWrite are set, immediately return a timeout
+ * condition (without waiting).  Return >0 if condition is met, 0
+ * if a timeout occurred, -1 if an error or interrupt occurred.
+ *
+ * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
+ * if end_time is 0 (or indeed, any time before now).
+ */
+static int
+pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+{
+	/*
+	 * We use functions in the following order if available:
+	 *   - ppoll(2) OR pselect(2)
+	 *   - poll(2) OR select(2)
+	 */
+#ifdef HAVE_POLL
+	struct pollfd input_fd;
+#ifdef HAVE_PPOLL
+	int			rc;
+	sigset_t	emptyset;
+	sigset_t	blockset;
+	sigset_t	origset;
+	struct timespec timeout;
+	struct timespec *ptr_timeout;
+#else
+	int			timeout_ms;
+#endif
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	input_fd.fd = sock;
+	input_fd.events = POLLERR;
+	input_fd.revents = 0;
+
+	if (forRead)
+		input_fd.events |= POLLIN;
+	if (forWrite)
+		input_fd.events |= POLLOUT;
+
+	/* Compute appropriate timeout interval */
+#ifdef HAVE_PPOLL
+	sigemptyset(&blockset);
+	sigaddset(&blockset, SIGINT);
+	sigprocmask(SIG_BLOCK, &blockset, &origset);
+
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		timeout.tv_sec = end_time;
+		timeout.tv_nsec = 0;
+		ptr_timeout = &timeout;
+	}
+#else
+	if (end_time == ((time_t) -1))
+		timeout_ms = -1;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout_ms = (end_time - now) * 1000;
+		else
+			timeout_ms = 0;
+	}
+#endif
+
+#ifdef HAVE_PPOLL
+	sigemptyset(&emptyset);
+	if (cancel_pressed)
+	{
+		sigprocmask(SIG_SETMASK, &origset, NULL);
+		return 1;
+	}
+
+	rc = ppoll(&input_fd, 1, ptr_timeout, &emptyset);
+	sigprocmask(SIG_SETMASK, &origset, NULL);
+	return rc;
+#else
+	return poll(&input_fd, 1, timeout_ms);
+#endif
+#else							/* !HAVE_POLL */
+
+	fd_set		input_mask;
+	fd_set		output_mask;
+	fd_set		except_mask;
+#ifdef HAVE_PSELECT
+	int			rc;
+	sigset_t	emptyset;
+	sigset_t	blockset;
+	sigset_t	origset;
+	struct timespec timeout;
+	struct timespec *ptr_timeout;
+#else
+	struct timeval timeout;
+	struct timeval *ptr_timeout;
+#endif
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	FD_ZERO(&input_mask);
+	FD_ZERO(&output_mask);
+	FD_ZERO(&except_mask);
+	if (forRead)
+		FD_SET(sock, &input_mask);
+
+	if (forWrite)
+		FD_SET(sock, &output_mask);
+	FD_SET(sock, &except_mask);
+
+	/* Compute appropriate timeout interval */
+#ifdef HAVE_PSELECT
+	sigemptyset(&blockset);
+	sigaddset(&blockset, SIGINT);
+	sigprocmask(SIG_BLOCK, &blockset, &origset);
+
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		timeout.tv_sec = end_time;
+		timeout.tv_nsec = 0;
+		ptr_timeout = &timeout;
+	}
+#else
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout.tv_sec = end_time - now;
+		else
+			timeout.tv_sec = 0;
+		timeout.tv_usec = 0;
+		ptr_timeout = &timeout;
+	}
+#endif
+
+#ifdef HAVE_PSELECT
+	sigemptyset(&emptyset);
+	if (cancel_pressed)
+	{
+		sigprocmask(SIG_SETMASK, &origset, NULL);
+		return 1;
+	}
+
+	rc = pselect(sock + 1, &input_mask, &output_mask,
+				 &except_mask, ptr_timeout, &emptyset);
+	sigprocmask(SIG_SETMASK, &origset, NULL);
+	return rc;
+#else
+	return select(sock + 1, &input_mask, &output_mask,
+				  &except_mask, ptr_timeout);
+#endif
+#endif							/* HAVE_POLL */
+}
+
 /*
  * Ask the user for a password; 'username' is the username the
  * password is for, if one has been explicitly specified.
@@ -3324,6 +3493,7 @@ do_connect(enum trivalue reuse_previous_specification,
 	bool		keep_password = true;
 	bool		has_connection_string;
 	bool		reuse_previous;
+	bool		for_read;
 
 	has_connection_string = dbname ?
 		recognized_connection_string(dbname) : false;
@@ -3614,11 +3784,63 @@ 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);
 
+		for_read = false;
+		while (true)
+		{
+			int			rc;
+			int			sock;
+			time_t		timeout;
+
+			if (cancel_pressed)
+				break;
+
+			sock = PQsocket(n_conn);
+			if (sock == -1)
+				break;
+
+			/*
+			 * We use ppoll(2)/pselect(2) to account for the race condition in
+			 * which SIGINT is sent after checking cancel_pressed. But on
+			 * platforms that don't have either function, we can just spin the
+			 * CPU a bit polling, so set the timeout to 1 second if we don't
+			 * have the aforementioned functions, otherwise set timeout to a
+			 * negative value indicating we will sit and wait forever.
+			 */
+#if defined(HAVE_PPOLL) || defined(HAVE_PSELECT)
+			timeout = -1;
+#else
+			timeout = 1;
+#endif
+
+			rc = pqSocketPoll(sock, for_read, !for_read, timeout);
+			if (rc == -1)
+			{
+				success = false;
+				break;
+			}
+
+			switch (PQconnectPoll(n_conn))
+			{
+				case PGRES_POLLING_OK:
+				case PGRES_POLLING_FAILED:
+					goto finish;
+				case PGRES_POLLING_READING:
+					for_read = true;
+					continue;
+				case PGRES_POLLING_WRITING:
+					for_read = false;
+					continue;
+				case PGRES_POLLING_ACTIVE:
+					pg_unreachable();
+			}
+		}
+
+finish:
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
diff --git a/src/include/pg_config.h.in b/src/include/pg_config.h.in
index d8a2985567..f9fd7d0de7 100644
--- a/src/include/pg_config.h.in
+++ b/src/include/pg_config.h.in
@@ -333,6 +333,9 @@
 /* Define to 1 if you have the `ppoll' function. */
 #undef HAVE_PPOLL
 
+/* Define to 1 if you have the `pselect' function. */
+#undef HAVE_PSELECT
+
 /* Define if you have POSIX threads libraries and header files. */
 #undef HAVE_PTHREAD
 
diff --git a/src/tools/msvc/Solution.pm b/src/tools/msvc/Solution.pm
index 98a5b5d872..d035f44f73 100644
--- a/src/tools/msvc/Solution.pm
+++ b/src/tools/msvc/Solution.pm
@@ -308,6 +308,7 @@ sub GenerateFiles
 		HAVE_POSIX_FADVISE => undef,
 		HAVE_POSIX_FALLOCATE => undef,
 		HAVE_PPOLL => undef,
+		HAVE_PSELECT => undef,
 		HAVE_PTHREAD => undef,
 		HAVE_PTHREAD_BARRIER_WAIT => undef,
 		HAVE_PTHREAD_IS_THREADED_NP => undef,
-- 
Tristan Partin
Neon (https://neon.tech)

From 5e19b26111708b654c7b49c3b9fd7dc18b94c0e7 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 24 Jul 2023 11:12:59 -0500
Subject: [PATCH v7 1/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 CancelRequested check
into the polling loop.
---
 meson.build                |   1 +
 src/bin/psql/command.c     | 156 ++++++++++++++++++++++++++++++++++++-
 src/include/pg_config.h.in |   3 +
 src/tools/msvc/Solution.pm |   1 +
 4 files changed, 160 insertions(+), 1 deletion(-)

diff --git a/meson.build b/meson.build
index ee58ee7a06..2d63485c53 100644
--- a/meson.build
+++ b/meson.build
@@ -2440,6 +2440,7 @@ func_checks = [
   ['posix_fadvise'],
   ['posix_fallocate'],
   ['ppoll'],
+  ['pselect'],
   ['pstat'],
   ['pthread_barrier_wait', {'dependencies': [thread_dep]}],
   ['pthread_is_threaded_np', {'dependencies': [thread_dep]}],
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 82cc091568..3a76623b05 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -11,6 +11,7 @@
 #include <time.h>
 #include <pwd.h>
 #include <utime.h>
+#include <sys/select.h>
 #ifndef WIN32
 #include <sys/stat.h>			/* for stat() */
 #include <sys/time.h>			/* for setitimer() */
@@ -40,6 +41,7 @@
 #include "large_obj.h"
 #include "libpq-fe.h"
 #include "libpq/pqcomm.h"
+#include "libpq/pqsignal.h"
 #include "mainloop.h"
 #include "portability/instr_time.h"
 #include "pqexpbuffer.h"
@@ -3298,6 +3300,157 @@ param_is_newly_set(const char *old_val, const char *new_val)
 	return false;
 }
 
+/*
+ * Check a file descriptor for read and/or write data, possibly waiting.
+ * If neither forRead nor forWrite are set, immediately return a timeout
+ * condition (without waiting).  Return >0 if condition is met, 0
+ * if a timeout occurred, -1 if an error or interrupt occurred.
+ *
+ * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
+ * if end_time is 0 (or indeed, any time before now).
+ *
+ * Uses the select(2) family of functions because it is available on every
+ * platform. It is unlikely that psql would be holding enough file descriptors
+ * that would necessitate using poll(2) or ppoll(2) for example.
+ */
+static int
+pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+{
+	/*
+	 * We use functions in the following order if available: pselect(2) or,
+	 * select(2).
+	 */
+	fd_set		input_mask;
+	fd_set		output_mask;
+	fd_set		except_mask;
+#ifdef HAVE_PSELECT
+	int			rc;
+	sigset_t	emptyset;
+	sigset_t	blockset;
+	sigset_t	origset;
+	struct timespec timeout;
+	struct timespec *ptr_timeout;
+#else
+	struct timeval timeout;
+	struct timeval *ptr_timeout;
+#endif
+
+	if (!forRead && !forWrite)
+		return 0;
+
+	FD_ZERO(&input_mask);
+	FD_ZERO(&output_mask);
+	FD_ZERO(&except_mask);
+	if (forRead)
+		FD_SET(sock, &input_mask);
+
+	if (forWrite)
+		FD_SET(sock, &output_mask);
+	FD_SET(sock, &except_mask);
+
+	/* Compute appropriate timeout interval */
+#ifdef HAVE_PSELECT
+	sigemptyset(&blockset);
+	sigaddset(&blockset, SIGINT);
+	sigprocmask(SIG_BLOCK, &blockset, &origset);
+
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		timeout.tv_sec = end_time;
+		timeout.tv_nsec = 0;
+		ptr_timeout = &timeout;
+	}
+#else
+	if (end_time == ((time_t) -1))
+		ptr_timeout = NULL;
+	else
+	{
+		time_t		now = time(NULL);
+
+		if (end_time > now)
+			timeout.tv_sec = end_time - now;
+		else
+			timeout.tv_sec = 0;
+		timeout.tv_usec = 0;
+		ptr_timeout = &timeout;
+	}
+#endif
+
+#ifdef HAVE_PSELECT
+	sigemptyset(&emptyset);
+	if (cancel_pressed)
+	{
+		sigprocmask(SIG_SETMASK, &origset, NULL);
+		return 1;
+	}
+
+	rc = pselect(sock + 1, &input_mask, &output_mask,
+				 &except_mask, ptr_timeout, &emptyset);
+	sigprocmask(SIG_SETMASK, &origset, NULL);
+	return rc;
+#else
+	return select(sock + 1, &input_mask, &output_mask,
+				  &except_mask, ptr_timeout);
+#endif
+}
+
+static void
+process_connection_state_machine(PGconn *conn)
+{
+	bool		for_read = false;
+
+	while (true)
+	{
+		int			rc;
+		int			sock;
+		time_t		timeout;
+
+		if (cancel_pressed)
+			break;
+
+		sock = PQsocket(conn);
+		if (sock == -1)
+			break;
+
+		/*
+		 * We use pselect(2) to account for the race condition in which SIGINT
+		 * is sent after checking cancel_pressed. But on platforms that don't
+		 * have either function, we can just spin the CPU a bit polling, so
+		 * set the timeout to 1 second if we don't have the aforementioned
+		 * function. Otherwise, set timeout to a negative value indicating we
+		 * will sit and wait forever.
+		 */
+#if defined(HAVE_PSELECT)
+		timeout = -1;
+#else
+		timeout = 1;
+#endif
+
+		rc = pqSocketPoll(sock, for_read, !for_read, timeout);
+		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
  *
@@ -3614,11 +3767,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/include/pg_config.h.in b/src/include/pg_config.h.in
index d8a2985567..f9fd7d0de7 100644
--- a/src/include/pg_config.h.in
+++ b/src/include/pg_config.h.in
@@ -333,6 +333,9 @@
 /* Define to 1 if you have the `ppoll' function. */
 #undef HAVE_PPOLL
 
+/* Define to 1 if you have the `pselect' function. */
+#undef HAVE_PSELECT
+
 /* Define if you have POSIX threads libraries and header files. */
 #undef HAVE_PTHREAD
 
diff --git a/src/tools/msvc/Solution.pm b/src/tools/msvc/Solution.pm
index 98a5b5d872..d035f44f73 100644
--- a/src/tools/msvc/Solution.pm
+++ b/src/tools/msvc/Solution.pm
@@ -308,6 +308,7 @@ sub GenerateFiles
 		HAVE_POSIX_FADVISE => undef,
 		HAVE_POSIX_FALLOCATE => undef,
 		HAVE_PPOLL => undef,
+		HAVE_PSELECT => undef,
 		HAVE_PTHREAD => undef,
 		HAVE_PTHREAD_BARRIER_WAIT => undef,
 		HAVE_PTHREAD_IS_THREADED_NP => undef,
-- 
Tristan Partin
Neon (https://neon.tech)

From 4fa6db1900a355a171d3e16019d02f3a415764d0 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:40:57 -0600
Subject: [PATCH v8 1/2] Expose PQsocketPoll for frontends

PQsocketPoll should help with state machine processing when connecting
to a database asynchronously via PQconnectStart().
---
 doc/src/sgml/libpq.sgml         | 5 ++++-
 src/interfaces/libpq/fe-misc.c  | 7 +++----
 src/interfaces/libpq/libpq-fe.h | 4 ++++
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index d0d5aefadc..aa26c2cc8d 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -358,7 +358,10 @@ PostgresPollingStatusType PQconnectPoll(PGconn *conn);
        Loop thus: If <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_READING</symbol>, wait until the socket is ready to
        read (as indicated by <function>select()</function>, <function>poll()</function>, or
-       similar system function).
+       similar system function).  Note that <function>PQsocketPoll</function>
+       can help reduce boilerplate by abstracting the setup of
+       <function>select()</function>, or <function>poll()</function> if it is
+       available on your system.
        Then call <function>PQconnectPoll(conn)</function> again.
        Conversely, if <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_WRITING</symbol>, wait until the socket is ready
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index 47a28b0a3a..ee917d375d 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -55,7 +55,6 @@ static int	pqPutMsgBytes(const void *buf, size_t len, PGconn *conn);
 static int	pqSendSome(PGconn *conn, int len);
 static int	pqSocketCheck(PGconn *conn, int forRead, int forWrite,
 						  time_t end_time);
-static int	pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time);
 
 /*
  * PQlibVersion: return the libpq version number
@@ -1059,7 +1058,7 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 
 	/* We will retry as long as we get EINTR */
 	do
-		result = pqSocketPoll(conn->sock, forRead, forWrite, end_time);
+		result = PQsocketPoll(conn->sock, forRead, forWrite, end_time);
 	while (result < 0 && SOCK_ERRNO == EINTR);
 
 	if (result < 0)
@@ -1083,8 +1082,8 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
  * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
  * if end_time is 0 (or indeed, any time before now).
  */
-static int
-pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+int
+PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 {
 	/* We use poll(2) if available, otherwise select(2) */
 #ifdef HAVE_POLL
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index defc415fa3..11a7fd32b6 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -21,6 +21,7 @@ extern "C"
 #endif
 
 #include <stdio.h>
+#include <time.h>
 
 /*
  * postgres_ext.h defines the backend's externally visible types,
@@ -644,6 +645,9 @@ extern int	lo_export(PGconn *conn, Oid lobjId, const char *filename);
 /* Get the version of the libpq library in use */
 extern int	PQlibVersion(void);
 
+/* Poll a file descriptor for reading and/or writing with a timeout */
+extern int	PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+
 /* Determine length of multibyte encoded char at *s */
 extern int	PQmblen(const char *s, int encoding);
 
-- 
Tristan Partin
Neon (https://neon.tech)

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)

From 7f8bf7250ecf79f65c129ccc42643c36bc53f882 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:40:57 -0600
Subject: [PATCH v9 1/2] Expose PQsocketPoll for frontends

PQsocketPoll should help with state machine processing when connecting
to a database asynchronously via PQconnectStart().
---
 doc/src/sgml/libpq.sgml          | 38 +++++++++++++++++++++++++++++++-
 src/interfaces/libpq/exports.txt |  1 +
 src/interfaces/libpq/fe-misc.c   |  7 +++---
 src/interfaces/libpq/libpq-fe.h  |  4 ++++
 4 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index d3e87056f2c..10f191f6b88 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -262,6 +262,39 @@ PGconn *PQsetdb(char *pghost,
      </listitem>
     </varlistentry>
 
+    <varlistentry id="libpq-PQsocketPoll">
+     <term><function>PQsocketPoll</function><indexterm><primary>PQsocketPoll</primary></indexterm></term>
+     <listitem>
+      <para>
+       <indexterm><primary>nonblocking connection</primary></indexterm>
+       Poll a connection&apos;s underlying socket descriptor retrieved with <xref linkend="libpq-PQsocket"/>.
+<synopsis>
+int PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+</synopsis>
+      </para>
+
+      <para>
+       This function takes care of the setup for monitoring a file descriptor. The underlying
+       function is either <function>poll(2)</function> or <function>select(2)</function>,
+       depending on platform support. The primary use of this function is working through the state
+       machine instantiated by <xref linkend="libpq-PQconnectStartParams"/>.
+      </para>
+
+      <para>
+       The function returns a value greater than <literal>0</literal> if the specified condition
+       is met, <literal>0</literal> if a timeout occurred, or <literal>-1</literal> if an error
+       or interrupt occurred. In the event <literal>forRead</literal> and
+       <literal>forWrite</literal> are not set, the function immediately retuns a timeout
+       condition.
+      </para>
+
+      <para>
+       <literal>end_time</literal> is the time in the future in seconds starting from the UNIX
+       epoch in which you would like the function to return if the condition is not met.
+      </para>
+     </listitem>
+    </varlistentry>
+
     <varlistentry id="libpq-PQconnectStartParams">
      <term><function>PQconnectStartParams</function><indexterm><primary>PQconnectStartParams</primary></indexterm></term>
      <term><function>PQconnectStart</function><indexterm><primary>PQconnectStart</primary></indexterm></term>
@@ -358,7 +391,10 @@ PostgresPollingStatusType PQconnectPoll(PGconn *conn);
        Loop thus: If <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_READING</symbol>, wait until the socket is ready to
        read (as indicated by <function>select()</function>, <function>poll()</function>, or
-       similar system function).
+       similar system function).  Note that <function>PQsocketPoll</function>
+       can help reduce boilerplate by abstracting the setup of
+       <function>select(2)</function> or <function>poll(2)</function> if it is
+       available on your system.
        Then call <function>PQconnectPoll(conn)</function> again.
        Conversely, if <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_WRITING</symbol>, wait until the socket is ready
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index 9fbd3d34074..1e48d37677d 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -202,3 +202,4 @@ PQcancelSocket            199
 PQcancelErrorMessage      200
 PQcancelReset             201
 PQcancelFinish            202
+PQsocketPoll              203
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index f2fc78a481c..f562cd8d344 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -55,7 +55,6 @@ static int	pqPutMsgBytes(const void *buf, size_t len, PGconn *conn);
 static int	pqSendSome(PGconn *conn, int len);
 static int	pqSocketCheck(PGconn *conn, int forRead, int forWrite,
 						  time_t end_time);
-static int	pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time);
 
 /*
  * PQlibVersion: return the libpq version number
@@ -1059,7 +1058,7 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 
 	/* We will retry as long as we get EINTR */
 	do
-		result = pqSocketPoll(conn->sock, forRead, forWrite, end_time);
+		result = PQsocketPoll(conn->sock, forRead, forWrite, end_time);
 	while (result < 0 && SOCK_ERRNO == EINTR);
 
 	if (result < 0)
@@ -1083,8 +1082,8 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
  * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
  * if end_time is 0 (or indeed, any time before now).
  */
-static int
-pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+int
+PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 {
 	/* We use poll(2) if available, otherwise select(2) */
 #ifdef HAVE_POLL
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 09b485bd2bc..98c27415281 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -21,6 +21,7 @@ extern "C"
 #endif
 
 #include <stdio.h>
+#include <time.h>
 
 /*
  * postgres_ext.h defines the backend's externally visible types,
@@ -670,6 +671,9 @@ extern int	lo_export(PGconn *conn, Oid lobjId, const char *filename);
 /* Get the version of the libpq library in use */
 extern int	PQlibVersion(void);
 
+/* Poll a file descriptor for reading and/or writing with a timeout */
+extern int	PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+
 /* Determine length of multibyte encoded char at *s */
 extern int	PQmblen(const char *s, int encoding);
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 779e642df8013b3fc8c63f9faf370cbbf2950392 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:53:10 -0600
Subject: [PATCH v9 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.
---
 src/bin/psql/command.c | 50 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 9b0fa041f73..00a20e47328 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -159,6 +159,7 @@ static void discard_query_text(PsqlScanState scan_state, ConditionalStack cstack
 static bool copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf);
 static bool do_connect(enum trivalue reuse_previous_specification,
 					   char *dbname, char *user, char *host, char *port);
+static void wait_until_connected(PGconn *conn);
 static bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
 					int lineno, bool discard_on_quit, bool *edited);
 static bool do_shell(const char *command);
@@ -3595,11 +3596,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);
 
+		wait_until_connected(n_conn);
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
@@ -3748,6 +3750,52 @@ do_connect(enum trivalue reuse_previous_specification,
 	return true;
 }
 
+static void
+wait_until_connected(PGconn *conn)
+{
+	bool		for_read = false;
+
+	while (true)
+	{
+		int			rc;
+		int			sock;
+		time_t		timeout;
+
+		/*
+		 * On every iteration of the state machine, let's check if the user has
+		 * requested a cancellation of the connection process.
+		 */
+		if (cancel_pressed)
+			break;
+
+		sock = PQsocket(conn);
+		if (sock == -1)
+			break;
+
+		/* Use a 1 second timeout */
+		timeout = time(NULL) + 1;
+		rc = PQsocketPoll(sock, for_read, !for_read, timeout);
+		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();
+}
 
 void
 connection_warnings(bool in_startup)
-- 
Tristan Partin
Neon (https://neon.tech)

From dc1eb3866975680b55451b79b40f109e3104eee8 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:53:10 -0600
Subject: [PATCH v10 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.
---
 src/bin/psql/command.c | 72 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 9b0fa041f73..1e00b0d4869 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -159,6 +159,7 @@ static void discard_query_text(PsqlScanState scan_state, ConditionalStack cstack
 static bool copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf);
 static bool do_connect(enum trivalue reuse_previous_specification,
 					   char *dbname, char *user, char *host, char *port);
+static void wait_until_connected(PGconn *conn);
 static bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
 					int lineno, bool discard_on_quit, bool *edited);
 static bool do_shell(const char *command);
@@ -3595,11 +3596,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);
 
+		wait_until_connected(n_conn);
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
@@ -3748,6 +3750,74 @@ do_connect(enum trivalue reuse_previous_specification,
 	return true;
 }
 
+/*
+ * Processes the connection sequence described by PQconnectStartParams(). Don't
+ * worry about reporting errors in this function. Our caller will check the
+ * connection's status, and report appropriately.
+ */
+static void
+wait_until_connected(PGconn *conn)
+{
+	bool		forRead = false;
+
+	while (true)
+	{
+		int			rc;
+		int			sock;
+		time_t		end_time;
+
+		/*
+		 * On every iteration of the connection sequence, let's check if the
+		 * user has requested a cancellation.
+		 */
+		if (cancel_pressed)
+			break;
+
+		/*
+		 * Do not assume that the socket remains the same across
+		 * PQconnectPoll() calls.
+		 */
+		sock = PQsocket(conn);
+		if (sock == -1)
+			break;
+
+		/*
+		 * If the user sends SIGINT between the cancel_pressed check, and
+		 * polling of the socket, it will not be recognized. Instead, we will
+		 * just wait until the next step in the connection sequence or forever,
+		 * which might require users to send SIGTERM or SIGQUIT.
+		 *
+		 * Some solutions would include the "self-pipe trick," using
+		 * pselect(2) and ppoll(2), or using a timeout.
+		 *
+		 * The self-pipe trick requires a bit of code to setup. pselect(2) and
+		 * ppoll(2) are not on all the platforms we support. The simplest
+		 * solution happens to just be adding a timeout, so let's wait for 1
+		 * second and check cancel_pressed again.
+		 */
+		end_time = time(NULL) + 1;
+		rc = PQsocketPoll(sock, forRead, !forRead, end_time);
+		if (rc == -1)
+			return;
+
+		switch (PQconnectPoll(conn))
+		{
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				return;
+			case PGRES_POLLING_READING:
+				forRead = true;
+				continue;
+			case PGRES_POLLING_WRITING:
+				forRead = false;
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+		}
+	}
+
+	pg_unreachable();
+}
 
 void
 connection_warnings(bool in_startup)
-- 
Tristan Partin
Neon (https://neon.tech)

From dc45e95ab443d973845dd840600d6719dcd4cae2 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:40:57 -0600
Subject: [PATCH v10 1/2] Expose PQsocketPoll for frontends

PQsocketPoll should help with state machine processing when connecting
to a database asynchronously via PQconnectStart().
---
 doc/src/sgml/libpq.sgml          | 43 +++++++++++++++++++++++++++++++-
 src/interfaces/libpq/exports.txt |  1 +
 src/interfaces/libpq/fe-misc.c   |  7 +++---
 src/interfaces/libpq/libpq-fe.h  |  4 +++
 4 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index d3e87056f2c..774b57ba70b 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -262,6 +262,44 @@ PGconn *PQsetdb(char *pghost,
      </listitem>
     </varlistentry>
 
+    <varlistentry id="libpq-PQsocketPoll">
+     <term><function>PQsocketPoll</function><indexterm><primary>PQsocketPoll</primary></indexterm></term>
+     <listitem>
+      <para>
+       <indexterm><primary>nonblocking connection</primary></indexterm>
+       Poll a connection&apos;s underlying socket descriptor retrieved with <xref linkend="libpq-PQsocket"/>.
+<synopsis>
+int PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+</synopsis>
+      </para>
+
+      <para>
+       This function sets up polling of a file descriptor. The underlying function is either
+       <function>poll(2)</function> or <function>select(2)</function>, depending on platform
+       support. The primary use of this function is iterating through the connection sequence
+       described in the documentation of <xref linkend="libpq-PQconnectStartParams"/>. If
+       <parameter>forRead</parameter> is specified, the function waits for the socket to be ready
+       for reading. If <parameter>forWrite</parameter> is specified, the function waits for the
+       socket to be ready for write. See <literal>POLLIN</literal> and <literal>POLLOUT</literal>
+       from <function>poll(2)</function>, or <parameter>readfds</parameter> and
+       <parameter>writefds</parameter> from <function>select(2)</function> for more information.
+      </para>
+
+      <para>
+       The function returns a value greater than <literal>0</literal> if the specified condition
+       is met, <literal>0</literal> if a timeout occurred, or <literal>-1</literal> if an error
+       or interrupt occurred. In the event <literal>forRead</literal> and
+       <literal>forWrite</literal> are not set, the function immediately returns a timeout
+       condition.
+      </para>
+
+      <para>
+       <literal>end_time</literal> is the time in the future in seconds starting from the UNIX
+       epoch in which you would like the function to return if the condition is not met.
+      </para>
+     </listitem>
+    </varlistentry>
+
     <varlistentry id="libpq-PQconnectStartParams">
      <term><function>PQconnectStartParams</function><indexterm><primary>PQconnectStartParams</primary></indexterm></term>
      <term><function>PQconnectStart</function><indexterm><primary>PQconnectStart</primary></indexterm></term>
@@ -358,7 +396,10 @@ PostgresPollingStatusType PQconnectPoll(PGconn *conn);
        Loop thus: If <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_READING</symbol>, wait until the socket is ready to
        read (as indicated by <function>select()</function>, <function>poll()</function>, or
-       similar system function).
+       similar system function).  Note that <function>PQsocketPoll</function>
+       can help reduce boilerplate by abstracting the setup of
+       <function>select(2)</function> or <function>poll(2)</function> if it is
+       available on your system.
        Then call <function>PQconnectPoll(conn)</function> again.
        Conversely, if <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_WRITING</symbol>, wait until the socket is ready
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index 9fbd3d34074..1e48d37677d 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -202,3 +202,4 @@ PQcancelSocket            199
 PQcancelErrorMessage      200
 PQcancelReset             201
 PQcancelFinish            202
+PQsocketPoll              203
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index f2fc78a481c..f562cd8d344 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -55,7 +55,6 @@ static int	pqPutMsgBytes(const void *buf, size_t len, PGconn *conn);
 static int	pqSendSome(PGconn *conn, int len);
 static int	pqSocketCheck(PGconn *conn, int forRead, int forWrite,
 						  time_t end_time);
-static int	pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time);
 
 /*
  * PQlibVersion: return the libpq version number
@@ -1059,7 +1058,7 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 
 	/* We will retry as long as we get EINTR */
 	do
-		result = pqSocketPoll(conn->sock, forRead, forWrite, end_time);
+		result = PQsocketPoll(conn->sock, forRead, forWrite, end_time);
 	while (result < 0 && SOCK_ERRNO == EINTR);
 
 	if (result < 0)
@@ -1083,8 +1082,8 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
  * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
  * if end_time is 0 (or indeed, any time before now).
  */
-static int
-pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+int
+PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 {
 	/* We use poll(2) if available, otherwise select(2) */
 #ifdef HAVE_POLL
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 09b485bd2bc..8d3c5c6f662 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -21,6 +21,7 @@ extern "C"
 #endif
 
 #include <stdio.h>
+#include <time.h>
 
 /*
  * postgres_ext.h defines the backend's externally visible types,
@@ -670,6 +671,9 @@ extern int	lo_export(PGconn *conn, Oid lobjId, const char *filename);
 /* Get the version of the libpq library in use */
 extern int	PQlibVersion(void);
 
+/* Poll a socket for reading and/or writing with an optional timeout */
+extern int	PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+
 /* Determine length of multibyte encoded char at *s */
 extern int	PQmblen(const char *s, int encoding);
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 14c794d9699f7b9f27d1a4ec026f748c6b7d8853 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:40:57 -0600
Subject: [PATCH v11 1/2] Expose PQsocketPoll for frontends

PQsocketPoll should help with state machine processing when connecting
to a database asynchronously via PQconnectStart().
---
 doc/src/sgml/libpq.sgml          | 40 +++++++++++++++++++++++++++++++-
 src/interfaces/libpq/exports.txt |  1 +
 src/interfaces/libpq/fe-misc.c   |  7 +++---
 src/interfaces/libpq/libpq-fe.h  |  4 ++++
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index d3e87056f2c..e69feacfe6a 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -262,6 +262,41 @@ PGconn *PQsetdb(char *pghost,
      </listitem>
     </varlistentry>
 
+    <varlistentry id="libpq-PQsocketPoll">
+     <term><function>PQsocketPoll</function><indexterm><primary>PQsocketPoll</primary></indexterm></term>
+     <listitem>
+      <para>
+       <indexterm><primary>nonblocking connection</primary></indexterm>
+       Poll a connection&apos;s underlying socket descriptor retrieved with <xref linkend="libpq-PQsocket"/>.
+<synopsis>
+int PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+</synopsis>
+      </para>
+
+      <para>
+       This function sets up polling of a file descriptor. The underlying function is either
+       <function>poll(2)</function> or <function>select(2)</function>, depending on platform
+       support. The primary use of this function is iterating through the connection sequence
+       described in the documentation of <xref linkend="libpq-PQconnectStartParams"/>. If
+       <parameter>forRead</parameter> is specified, the function waits for the socket to be ready
+       for reading. If <parameter>forWrite</parameter> is specified, the function waits for the
+       socket to be ready for write. See <literal>POLLIN</literal> and <literal>POLLOUT</literal>
+       from <function>poll(2)</function>, or <parameter>readfds</parameter> and
+       <parameter>writefds</parameter> from <function>select(2)</function> for more information. If
+       <parameter>end_time</parameter> is not <literal>-1</literal>, it specifies the time at which
+       this function should stop waiting for the condition to be met.
+      </para>
+
+      <para>
+       The function returns a value greater than <literal>0</literal> if the specified condition
+       is met, <literal>0</literal> if a timeout occurred, or <literal>-1</literal> if an error
+       occurred. The error can be retrieved by checking the <literal>errno(3)</literal> value. In
+       the event <literal>forRead</literal> and <literal>forWrite</literal> are not set, the
+       function immediately returns a timeout condition.
+      </para>
+     </listitem>
+    </varlistentry>
+
     <varlistentry id="libpq-PQconnectStartParams">
      <term><function>PQconnectStartParams</function><indexterm><primary>PQconnectStartParams</primary></indexterm></term>
      <term><function>PQconnectStart</function><indexterm><primary>PQconnectStart</primary></indexterm></term>
@@ -358,7 +393,10 @@ PostgresPollingStatusType PQconnectPoll(PGconn *conn);
        Loop thus: If <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_READING</symbol>, wait until the socket is ready to
        read (as indicated by <function>select()</function>, <function>poll()</function>, or
-       similar system function).
+       similar system function).  Note that <function>PQsocketPoll</function>
+       can help reduce boilerplate by abstracting the setup of
+       <function>select(2)</function> or <function>poll(2)</function> if it is
+       available on your system.
        Then call <function>PQconnectPoll(conn)</function> again.
        Conversely, if <function>PQconnectPoll(conn)</function> last returned
        <symbol>PGRES_POLLING_WRITING</symbol>, wait until the socket is ready
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index 9fbd3d34074..1e48d37677d 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -202,3 +202,4 @@ PQcancelSocket            199
 PQcancelErrorMessage      200
 PQcancelReset             201
 PQcancelFinish            202
+PQsocketPoll              203
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index f2fc78a481c..f562cd8d344 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -55,7 +55,6 @@ static int	pqPutMsgBytes(const void *buf, size_t len, PGconn *conn);
 static int	pqSendSome(PGconn *conn, int len);
 static int	pqSocketCheck(PGconn *conn, int forRead, int forWrite,
 						  time_t end_time);
-static int	pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time);
 
 /*
  * PQlibVersion: return the libpq version number
@@ -1059,7 +1058,7 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 
 	/* We will retry as long as we get EINTR */
 	do
-		result = pqSocketPoll(conn->sock, forRead, forWrite, end_time);
+		result = PQsocketPoll(conn->sock, forRead, forWrite, end_time);
 	while (result < 0 && SOCK_ERRNO == EINTR);
 
 	if (result < 0)
@@ -1083,8 +1082,8 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
  * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
  * if end_time is 0 (or indeed, any time before now).
  */
-static int
-pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+int
+PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 {
 	/* We use poll(2) if available, otherwise select(2) */
 #ifdef HAVE_POLL
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 09b485bd2bc..8d3c5c6f662 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -21,6 +21,7 @@ extern "C"
 #endif
 
 #include <stdio.h>
+#include <time.h>
 
 /*
  * postgres_ext.h defines the backend's externally visible types,
@@ -670,6 +671,9 @@ extern int	lo_export(PGconn *conn, Oid lobjId, const char *filename);
 /* Get the version of the libpq library in use */
 extern int	PQlibVersion(void);
 
+/* Poll a socket for reading and/or writing with an optional timeout */
+extern int	PQsocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+
 /* Determine length of multibyte encoded char at *s */
 extern int	PQmblen(const char *s, int encoding);
 
-- 
Tristan Partin
Neon (https://neon.tech)

From 6e3b08aa89e7a0268c3ade83c708adf1e1637ffa Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:53:10 -0600
Subject: [PATCH v11 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.
---
 src/bin/psql/command.c | 72 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 9b0fa041f73..1e00b0d4869 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -159,6 +159,7 @@ static void discard_query_text(PsqlScanState scan_state, ConditionalStack cstack
 static bool copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf);
 static bool do_connect(enum trivalue reuse_previous_specification,
 					   char *dbname, char *user, char *host, char *port);
+static void wait_until_connected(PGconn *conn);
 static bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
 					int lineno, bool discard_on_quit, bool *edited);
 static bool do_shell(const char *command);
@@ -3595,11 +3596,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);
 
+		wait_until_connected(n_conn);
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
@@ -3748,6 +3750,74 @@ do_connect(enum trivalue reuse_previous_specification,
 	return true;
 }
 
+/*
+ * Processes the connection sequence described by PQconnectStartParams(). Don't
+ * worry about reporting errors in this function. Our caller will check the
+ * connection's status, and report appropriately.
+ */
+static void
+wait_until_connected(PGconn *conn)
+{
+	bool		forRead = false;
+
+	while (true)
+	{
+		int			rc;
+		int			sock;
+		time_t		end_time;
+
+		/*
+		 * On every iteration of the connection sequence, let's check if the
+		 * user has requested a cancellation.
+		 */
+		if (cancel_pressed)
+			break;
+
+		/*
+		 * Do not assume that the socket remains the same across
+		 * PQconnectPoll() calls.
+		 */
+		sock = PQsocket(conn);
+		if (sock == -1)
+			break;
+
+		/*
+		 * If the user sends SIGINT between the cancel_pressed check, and
+		 * polling of the socket, it will not be recognized. Instead, we will
+		 * just wait until the next step in the connection sequence or forever,
+		 * which might require users to send SIGTERM or SIGQUIT.
+		 *
+		 * Some solutions would include the "self-pipe trick," using
+		 * pselect(2) and ppoll(2), or using a timeout.
+		 *
+		 * The self-pipe trick requires a bit of code to setup. pselect(2) and
+		 * ppoll(2) are not on all the platforms we support. The simplest
+		 * solution happens to just be adding a timeout, so let's wait for 1
+		 * second and check cancel_pressed again.
+		 */
+		end_time = time(NULL) + 1;
+		rc = PQsocketPoll(sock, forRead, !forRead, end_time);
+		if (rc == -1)
+			return;
+
+		switch (PQconnectPoll(conn))
+		{
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				return;
+			case PGRES_POLLING_READING:
+				forRead = true;
+				continue;
+			case PGRES_POLLING_WRITING:
+				forRead = false;
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+		}
+	}
+
+	pg_unreachable();
+}
 
 void
 connection_warnings(bool in_startup)
-- 
Tristan Partin
Neon (https://neon.tech)

From 7f4279bb939e2d2707fdd0f471893d734959ab91 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:21:52 +0200
Subject: [PATCH v11 2/2] Remove PGRES_POLLING_ACTIVE

This enum variant has been unused for at least 21 years
44aba280207740d0956160c0288e61f28f024a71. It was left for backwards
compatibility for "awhile". I think that time has come.
---
 src/bin/psql/command.c          | 2 --
 src/interfaces/libpq/libpq-fe.h | 2 --
 2 files changed, 4 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index dc3cc7c10b7..9a163cf22cd 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3811,8 +3811,6 @@ wait_until_connected(PGconn *conn)
 			case PGRES_POLLING_WRITING:
 				forRead = false;
 				continue;
-			case PGRES_POLLING_ACTIVE:
-				pg_unreachable();
 		}
 	}
 }
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 8d3c5c6f662..2459b0cf5e1 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -91,8 +91,6 @@ typedef enum
 	PGRES_POLLING_READING,		/* These two indicate that one may	  */
 	PGRES_POLLING_WRITING,		/* use select before polling again.   */
 	PGRES_POLLING_OK,
-	PGRES_POLLING_ACTIVE		/* unused; keep for awhile for backwards
-								 * compatibility */
 } PostgresPollingStatusType;
 
 typedef enum
-- 
2.34.1

From 29100578b5f0b4cec68ee1b8c572c4f280335da3 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:19:04 +0200
Subject: [PATCH v11 1/2] Fix actually reachable pg_unreachable call

In cafe1056558fe07cdc52b95205588fcd80870362 a call to pg_unreachable was
introduced that was actually reachable, because there were break
statements in the loop. This addresses that by both changing the break
statements to return and removing the call to pg_unreachable, which
seemed of dubious necessity anyway.
---
 src/bin/psql/command.c | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index c005624e9c3..dc3cc7c10b7 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3771,7 +3771,7 @@ wait_until_connected(PGconn *conn)
 		 * user has requested a cancellation.
 		 */
 		if (cancel_pressed)
-			break;
+			return;
 
 		/*
 		 * Do not assume that the socket remains the same across
@@ -3779,7 +3779,7 @@ wait_until_connected(PGconn *conn)
 		 */
 		sock = PQsocket(conn);
 		if (sock == -1)
-			break;
+			return;
 
 		/*
 		 * If the user sends SIGINT between the cancel_pressed check, and
@@ -3815,8 +3815,6 @@ wait_until_connected(PGconn *conn)
 				pg_unreachable();
 		}
 	}
-
-	pg_unreachable();
 }
 
 void

base-commit: 936e3fa3787a51397280c1081587586e83c20399
-- 
2.34.1

From 29100578b5f0b4cec68ee1b8c572c4f280335da3 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:19:04 +0200
Subject: [PATCH v12 1/2] Fix actually reachable pg_unreachable call

In cafe1056558fe07cdc52b95205588fcd80870362 a call to pg_unreachable was
introduced that was actually reachable, because there were break
statements in the loop. This addresses that by both changing the break
statements to return and removing the call to pg_unreachable, which
seemed of dubious necessity anyway.
---
 src/bin/psql/command.c | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index c005624e9c3..dc3cc7c10b7 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3771,7 +3771,7 @@ wait_until_connected(PGconn *conn)
 		 * user has requested a cancellation.
 		 */
 		if (cancel_pressed)
-			break;
+			return;
 
 		/*
 		 * Do not assume that the socket remains the same across
@@ -3779,7 +3779,7 @@ wait_until_connected(PGconn *conn)
 		 */
 		sock = PQsocket(conn);
 		if (sock == -1)
-			break;
+			return;
 
 		/*
 		 * If the user sends SIGINT between the cancel_pressed check, and
@@ -3815,8 +3815,6 @@ wait_until_connected(PGconn *conn)
 				pg_unreachable();
 		}
 	}
-
-	pg_unreachable();
 }
 
 void

base-commit: 936e3fa3787a51397280c1081587586e83c20399
-- 
2.34.1

From f71634361a0c9ba3416c75fc3ff6d55536f134c4 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:21:52 +0200
Subject: [PATCH v12 2/2] Remove PGRES_POLLING_ACTIVE case

This enum variant has been unused for at least 21 years
44aba280207740d0956160c0288e61f28f024a71. It's been there only for
backwards compatibility of the ABI, so no need to check for it.

Also update the comment on PGRES_POLLING_ACTIVE, since we're stuck with
it forever due to ABI compatibility guarantees.
---
 src/bin/psql/command.c          | 2 --
 src/interfaces/libpq/libpq-fe.h | 3 +--
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index dc3cc7c10b7..9a163cf22cd 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3811,8 +3811,6 @@ wait_until_connected(PGconn *conn)
 			case PGRES_POLLING_WRITING:
 				forRead = false;
 				continue;
-			case PGRES_POLLING_ACTIVE:
-				pg_unreachable();
 		}
 	}
 }
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 8d3c5c6f662..c184e853889 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -91,8 +91,7 @@ typedef enum
 	PGRES_POLLING_READING,		/* These two indicate that one may	  */
 	PGRES_POLLING_WRITING,		/* use select before polling again.   */
 	PGRES_POLLING_OK,
-	PGRES_POLLING_ACTIVE		/* unused; keep for awhile for backwards
-								 * compatibility */
+	PGRES_POLLING_ACTIVE		/* unused; keep for backwards compatibility */
 } PostgresPollingStatusType;
 
 typedef enum
-- 
2.34.1

From 29100578b5f0b4cec68ee1b8c572c4f280335da3 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:19:04 +0200
Subject: [PATCH v13 1/2] Fix actually reachable pg_unreachable call

In cafe1056558fe07cdc52b95205588fcd80870362 a call to pg_unreachable was
introduced that was actually reachable, because there were break
statements in the loop. This addresses that by both changing the break
statements to return and removing the call to pg_unreachable, which
seemed of dubious necessity anyway.
---
 src/bin/psql/command.c | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index c005624e9c3..dc3cc7c10b7 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3771,7 +3771,7 @@ wait_until_connected(PGconn *conn)
 		 * user has requested a cancellation.
 		 */
 		if (cancel_pressed)
-			break;
+			return;
 
 		/*
 		 * Do not assume that the socket remains the same across
@@ -3779,7 +3779,7 @@ wait_until_connected(PGconn *conn)
 		 */
 		sock = PQsocket(conn);
 		if (sock == -1)
-			break;
+			return;
 
 		/*
 		 * If the user sends SIGINT between the cancel_pressed check, and
@@ -3815,8 +3815,6 @@ wait_until_connected(PGconn *conn)
 				pg_unreachable();
 		}
 	}
-
-	pg_unreachable();
 }
 
 void

base-commit: 936e3fa3787a51397280c1081587586e83c20399
-- 
2.34.1

From 76cef25162b44adc20172afee47836ca765d3b5c Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:21:52 +0200
Subject: [PATCH v13 2/2] Remove PGRES_POLLING_ACTIVE case

This enum variant has been unused for at least 21 years
44aba280207740d0956160c0288e61f28f024a71. It's been there only for
backwards compatibility of the ABI, so no need to check for it.

Also update the comment on PGRES_POLLING_ACTIVE, since we're stuck with
it forever due to ABI compatibility guarantees.
---
 src/bin/psql/command.c          | 7 ++-----
 src/interfaces/libpq/libpq-fe.h | 3 +--
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index dc3cc7c10b7..5b31d08bf70 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3802,17 +3802,14 @@ wait_until_connected(PGconn *conn)
 
 		switch (PQconnectPoll(conn))
 		{
-			case PGRES_POLLING_OK:
-			case PGRES_POLLING_FAILED:
-				return;
 			case PGRES_POLLING_READING:
 				forRead = true;
 				continue;
 			case PGRES_POLLING_WRITING:
 				forRead = false;
 				continue;
-			case PGRES_POLLING_ACTIVE:
-				pg_unreachable();
+			default:
+				return;
 		}
 	}
 }
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 8d3c5c6f662..c184e853889 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -91,8 +91,7 @@ typedef enum
 	PGRES_POLLING_READING,		/* These two indicate that one may	  */
 	PGRES_POLLING_WRITING,		/* use select before polling again.   */
 	PGRES_POLLING_OK,
-	PGRES_POLLING_ACTIVE		/* unused; keep for awhile for backwards
-								 * compatibility */
+	PGRES_POLLING_ACTIVE		/* unused; keep for backwards compatibility */
 } PostgresPollingStatusType;
 
 typedef enum
-- 
2.34.1

From b0ccfb6fca3e4ae9e1e62ac3b6a4c5f428b56e04 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:19:04 +0200
Subject: [PATCH v14 1/2] Fix actually reachable pg_unreachable call

In cafe1056558fe07cdc52b95205588fcd80870362 a call to pg_unreachable was
introduced that was actually reachable, because there were break
statements in the loop. This addresses that by removing the call to
pg_unreachable, which seemed of dubious necessity anyway.
---
 src/bin/psql/command.c | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index c005624e9c3..479f9f2be59 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -3815,8 +3815,6 @@ wait_until_connected(PGconn *conn)
 				pg_unreachable();
 		}
 	}
-
-	pg_unreachable();
 }
 
 void

base-commit: 936e3fa3787a51397280c1081587586e83c20399
-- 
2.34.1

From 799eb6989d481ab617c473b8acbbfc1a405c3c7d Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Wed, 3 Apr 2024 15:21:52 +0200
Subject: [PATCH v14 2/2] Change comment on PGRES_POLLING_ACTIVE

Updates the comment on PGRES_POLLING_ACTIVE, since we're stuck with
it forever due to ABI compatibility guarantees.
---
 src/interfaces/libpq/libpq-fe.h | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 8d3c5c6f662..c184e853889 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -91,8 +91,7 @@ typedef enum
 	PGRES_POLLING_READING,		/* These two indicate that one may	  */
 	PGRES_POLLING_WRITING,		/* use select before polling again.   */
 	PGRES_POLLING_OK,
-	PGRES_POLLING_ACTIVE		/* unused; keep for awhile for backwards
-								 * compatibility */
+	PGRES_POLLING_ACTIVE		/* unused; keep for backwards compatibility */
 } PostgresPollingStatusType;
 
 typedef enum
-- 
2.34.1