Allow SIGINT to cancel psql database reconnections.
authorRobert Haas <rhaas@postgresql.org>
Tue, 2 Apr 2024 14:26:10 +0000 (10:26 -0400)
committerRobert Haas <rhaas@postgresql.org>
Tue, 2 Apr 2024 14:26:10 +0000 (10:26 -0400)
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.

Tristan Partin, reviewed by Gurjeet Singh, Heikki Linnakangas, Jelte
Fennema-Nio, and me.

Discussion: http://postgr.es/m/D08WWCPVHKHN.3QELIKZJ2D9RZ@neon.tech

src/bin/psql/command.c

index 9b0fa041f730c55f4890c08689c7ec99a4d47afd..1e00b0d4869aa565cd16ee25c2e3a9d8acfd90ff 100644 (file)
@@ -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)