From 0b8e21375ce3b55048503c51948c979cd443613c Mon Sep 17 00:00:00 2001
From: Jelte Fennema <jelte.fennema@microsoft.com>
Date: Thu, 30 Sep 2021 16:29:53 +0200
Subject: [PATCH] Use tcp_user_timeout in PQcancel

PGcancel would not adhere to the timeout specified in the
tcp_user_timeout connection option. Which means a call to PGcancel could
take much longer than expected. This can especially be an issue because
there's no non blocking version of PGcancel.
---
 src/interfaces/libpq/fe-connect.c | 26 +++++++++++++++++++++++---
 src/interfaces/libpq/libpq-int.h  |  1 +
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 56755f0796..0c1ec9d64d 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -4363,13 +4363,19 @@ PQgetCancel(PGconn *conn)
 	if (conn->sock == PGINVALID_SOCKET)
 		return NULL;
 
-	cancel = malloc(sizeof(PGcancel));
+	cancel = calloc(1, sizeof(PGcancel));
 	if (cancel == NULL)
 		return NULL;
 
 	memcpy(&cancel->raddr, &conn->raddr, sizeof(SockAddr));
 	cancel->be_pid = conn->be_pid;
 	cancel->be_key = conn->be_key;
+	if (conn->pgtcp_user_timeout != NULL) {
+		if (!parse_int_param(conn->pgtcp_user_timeout, &cancel->pgtcp_user_timeout, conn,
+							 "tcp_user_timeout")) {
+			return NULL;
+		}
+	}
 
 	return cancel;
 }
@@ -4404,7 +4410,7 @@ PQfreeCancel(PGcancel *cancel)
  * between the two versions of the cancel function possible.
  */
 static int
-internal_cancel(SockAddr *raddr, int be_pid, int be_key,
+internal_cancel(SockAddr *raddr, int be_pid, int be_key, int pgtcp_user_timeout,
 				char *errbuf, int errbufsize)
 {
 	int			save_errno = SOCK_ERRNO;
@@ -4426,6 +4432,19 @@ internal_cancel(SockAddr *raddr, int be_pid, int be_key,
 		strlcpy(errbuf, "PQcancel() -- socket() failed: ", errbufsize);
 		goto cancel_errReturn;
 	}
+
+#ifdef TCP_USER_TIMEOUT
+	if (!IS_AF_UNIX(raddr->addr.ss_family)) {
+		if (pgtcp_user_timeout < 0) {
+			pgtcp_user_timeout = 0;
+		}
+		if (setsockopt(tmpsock, IPPROTO_TCP, TCP_USER_TIMEOUT,
+					   (char *) &pgtcp_user_timeout, sizeof(pgtcp_user_timeout)) < 0) {
+			goto cancel_errReturn;
+		}
+	}
+#endif
+
 retry3:
 	if (connect(tmpsock, (struct sockaddr *) &raddr->addr,
 				raddr->salen) < 0)
@@ -4517,6 +4536,7 @@ PQcancel(PGcancel *cancel, char *errbuf, int errbufsize)
 	}
 
 	return internal_cancel(&cancel->raddr, cancel->be_pid, cancel->be_key,
+						   cancel->pgtcp_user_timeout,
 						   errbuf, errbufsize);
 }
 
@@ -4551,7 +4571,7 @@ PQrequestCancel(PGconn *conn)
 		return false;
 	}
 
-	r = internal_cancel(&conn->raddr, conn->be_pid, conn->be_key,
+	r = internal_cancel(&conn->raddr, conn->be_pid, conn->be_key, 0,
 						conn->errorMessage.data, conn->errorMessage.maxlen);
 
 	if (!r)
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 334aea4b6e..ef27a0bf23 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -581,6 +581,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			pgtcp_user_timeout; /* tcp user timeout */
 };
 
 
-- 
2.17.1