0002-Honor-connect_timeout-when-connecting-with-PQcancel.patch

From 78b1e7fc2bcef3d9b1b85b88387e24918411b2e9 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <jelte.fennema@microsoft.com>
Date: Tue, 28 Dec 2021 14:29:51 +0100
Subject: [PATCH 2/2] Honor connect_timeout when connecting with PQcancel

PQcancel uses its own connection mechanism, because it needs to be
signal safe. This meant that it was not honoring connect_timeout
connection option. This change fixes that for non windows environments,
by using the timeout option of select() to ensure a connect_timeout.
---
 src/interfaces/libpq/fe-connect.c | 128 +++++++++++++++++++++++++++---
 src/interfaces/libpq/libpq-int.h  |   1 +
 2 files changed, 118 insertions(+), 11 deletions(-)

diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index ed29c64484..61185db329 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -4394,11 +4394,20 @@ PQgetCancel(PGconn *conn)
 	cancel->be_key = conn->be_key;
 	/* Set all the socket options to -1, so we can use that to see if the */
 	/* socket options are set or not. */
+	cancel->connect_timeout = -1;
 	cancel->pgtcp_user_timeout = -1;
 	cancel->keepalives = -1;
 	cancel->keepalives_idle = -1;
 	cancel->keepalives_interval = -1;
 	cancel->keepalives_count = -1;
+	if (conn->connect_timeout != NULL)
+	{
+		if (!parse_int_param(conn->connect_timeout, &cancel->connect_timeout, conn,
+							 "connect_timeout"))
+		{
+			return NULL;
+		}
+	}
 	if (conn->pgtcp_user_timeout != NULL)
 	{
 		if (!parse_int_param(conn->pgtcp_user_timeout, &cancel->pgtcp_user_timeout, conn,
@@ -4479,6 +4488,14 @@ internal_cancel(PGcancel *cancel, char *errbuf, int errbufsize)
 	pgsocket	tmpsock = PGINVALID_SOCKET;
 	char		sebuf[PG_STRERROR_R_BUFLEN];
 	int			maxlen;
+#ifndef WIN32
+	int			tmpsockflags = 0;
+	fd_set		fdset;
+	bool		nonblocking_connect = cancel->connect_timeout > 0;
+	struct timeval connect_timeout;
+#else
+	bool		nonblocking_connect = false;
+#endif
 	struct
 	{
 		uint32		packetlen;
@@ -4581,21 +4598,110 @@ internal_cancel(PGcancel *cancel, char *errbuf, int errbufsize)
 	}
 #endif
 
+#ifndef WIN32
+	if (cancel->connect_timeout > 0)
+	{
+		/*
+		 * The connect() system call does not allow a timeout to be specified.
+		 * So to honor the connect_timeout that the user specified, we connect
+		 * in non blocking mode. Then we provide the connecti select() to
+		 */
+		nonblocking_connect = true;
+		tmpsockflags = fcntl(tmpsock, F_GETFL);
+		if (tmpsockflags < 0)
+		{
+			strlcpy(errbuf, "PQcancel() -- fcntl(F_GETFL) failed: ", errbufsize);
+			goto cancel_errReturn;
+		}
+		if (fcntl(tmpsock, F_SETFL, (tmpsockflags | O_NONBLOCK)) == -1)
+		{
+			strlcpy(errbuf, "PQcancel() -- fcntl(F_SETFL, O_NONBLOCK) failed: ", errbufsize);
+			goto cancel_errReturn;
+		}
+	}
+#endif
 
 retry3:
 	if (connect(tmpsock, (struct sockaddr *) &cancel->raddr.addr,
 				cancel->raddr.salen) < 0)
 	{
-		if (SOCK_ERRNO == EINTR)
-			/* Interrupted system call - we'll just try again */
+		if (nonblocking_connect && (
+									SOCK_ERRNO == EINPROGRESS ||
+#ifdef WIN32
+									SOCK_ERRNO == EWOULDBLOCK ||
+#endif
+									SOCK_ERRNO == EINTR))
+		{
+			/*
+			 * This is fine - we're in non-blocking mode, and the connection
+			 * is in progress.
+			 */
+		}
+		else if (SOCK_ERRNO == EINTR)
+		{
+			/*
+			 * Interrupted system call while in blocking mode - we'll just try
+			 * again
+			 */
 			goto retry3;
-		strlcpy(errbuf, "PQcancel() -- connect() failed: ", errbufsize);
-		goto cancel_errReturn;
+		}
+		else
+		{
+			strlcpy(errbuf, "PQcancel() -- connect() failed: ", errbufsize);
+			goto cancel_errReturn;
+		}
 	}
 
-	/*
-	 * We needn't set nonblocking I/O or NODELAY options here.
-	 */
+
+#ifndef WIN32
+	if (cancel->connect_timeout > 0)
+	{
+		connect_timeout.tv_sec = cancel->connect_timeout;
+		connect_timeout.tv_usec = 0;
+
+		FD_ZERO(&fdset);
+		FD_SET(tmpsock, &fdset);
+retry4:
+		if (select(tmpsock + 1, NULL, &fdset, NULL, &connect_timeout) == 1)
+		{
+			int			so_error;
+			socklen_t	len = sizeof so_error;
+
+			if (SOCK_ERRNO == EINTR)
+				/* Interrupted system call - we'll just try again */
+				goto retry4;
+			if (getsockopt(tmpsock, SOL_SOCKET, SO_ERROR, &so_error, &len))
+			{
+				strlcpy(errbuf, "PQcancel() -- getsockopt(SO_ERROR) failed: ", errbufsize);
+				goto cancel_errReturn;
+			}
+			if (so_error != 0)
+			{
+				SOCK_ERRNO = so_error;
+				strlcpy(errbuf, "PQcancel() -- select() failed: ", errbufsize);
+				goto cancel_errReturn;
+			}
+		}
+		else
+		{
+			strlcpy(errbuf, "PQcancel() -- Connection timed out\n", errbufsize);
+			closesocket(tmpsock);
+			SOCK_ERRNO_SET(save_errno);
+			return false;
+		}
+
+		/*
+		 * Now that we're connected we want a blocking socket again. So we
+		 * reset the flags to the value they had before changing to non
+		 * blocking mode.
+		 */
+		if (fcntl(tmpsock, F_SETFL, tmpsockflags) == -1)
+		{
+			strlcpy(errbuf, "PQcancel() -- fcntl(F_SETFL, !O_NONBLOCK) failed: ", errbufsize);
+			goto cancel_errReturn;
+		}
+	}
+#endif
 
 	/* Create and send the cancel request packet. */
 
@@ -4604,12 +4710,12 @@ retry3:
 	crp.cp.backendPID = pg_hton32(cancel->be_pid);
 	crp.cp.cancelAuthCode = pg_hton32(cancel->be_key);
 
-retry4:
+retry5:
 	if (send(tmpsock, (char *) &crp, sizeof(crp), 0) != (int) sizeof(crp))
 	{
 		if (SOCK_ERRNO == EINTR)
 			/* Interrupted system call - we'll just try again */
-			goto retry4;
+			goto retry5;
 		strlcpy(errbuf, "PQcancel() -- send() failed: ", errbufsize);
 		goto cancel_errReturn;
 	}
@@ -4621,12 +4727,12 @@ retry4:
 	 * one we thought we were canceling.  Note we don't actually expect this
 	 * read to obtain any data, we are just waiting for EOF to be signaled.
 	 */
-retry5:
+retry6:
 	if (recv(tmpsock, (char *) &crp, 1, 0) < 0)
 	{
 		if (SOCK_ERRNO == EINTR)
 			/* Interrupted system call - we'll just try again */
-			goto retry5;
+			goto retry6;
 		/* we ignore other error conditions */
 	}
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 6dabb14451..c08ecbd06f 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -583,6 +583,7 @@ struct pg_cancel
 	SockAddr	raddr;			/* Remote address */
 	int			be_pid;			/* PID of backend --- needed for cancels */
 	int			be_key;			/* key of backend --- needed for cancels */
+	int			connect_timeout;	/* connection timeout */
 	int			pgtcp_user_timeout; /* tcp user timeout */
 	int			keepalives;		/* use TCP keepalives? */
 	int			keepalives_idle;	/* time between TCP keepalives */
-- 
2.17.1