libpq-be-fe-helpers.h: wrap new cancel APIs

Commit 61461a300c1c introduced new functions to libpq for cancelling
queries.  This commit introduces a helper function that backend-side
libraries and extensions can use to invoke those.  This function takes a
timeout and can itself be interrupted while it is waiting for a cancel
request to be sent and processed, instead of being blocked.

This replaces the usage of the old functions in postgres_fdw and dblink.

Finally, it also adds some test coverage for the cancel support in
postgres_fdw.

Author: Jelte Fennema-Nio <postgres@jeltef.nl>
Discussion: https://postgr.es/m/CAGECzQT_VgOWWENUqvUV9xQmbaCyXjtRRAYO8W07oqashk_N+g@mail.gmail.com

1 files changed, 89 insertions, 0 deletions

diff --git a/src/include/libpq/libpq-be-fe-helpers.h b/src/include/libpq/libpq-be-fe-helpers.h
index 5d33bcf32f7..2adf92030af 100644
--- a/src/include/libpq/libpq-be-fe-helpers.h
+++ b/src/include/libpq/libpq-be-fe-helpers.h
@@ -44,6 +44,8 @@
 #include "miscadmin.h"
 #include "storage/fd.h"
 #include "storage/latch.h"
+#include "utils/timestamp.h"
+#include "utils/wait_event.h"
 
 
 static inline void libpqsrv_connect_prepare(void);
@@ -365,4 +367,91 @@ libpqsrv_get_result(PGconn *conn, uint32 wait_event_info)
 	return PQgetResult(conn);
 }
 
+/*
+ * Submit a cancel request to the given connection, waiting only until
+ * the given time.
+ *
+ * We sleep interruptibly until we receive confirmation that the cancel
+ * request has been accepted, and if it is, return NULL; if the cancel
+ * request fails, return an error message string (which is not to be
+ * freed).
+ *
+ * For other problems (to wit: OOM when strdup'ing an error message from
+ * libpq), this function can ereport(ERROR).
+ *
+ * Note: this function leaks a string's worth of memory when reporting
+ * libpq errors.  Make sure to call it in a transient memory context.
+ */
+static inline char *
+libpqsrv_cancel(PGconn *conn, TimestampTz endtime)
+{
+	PGcancelConn *cancel_conn;
+	char	   *error = NULL;
+
+	cancel_conn = PQcancelCreate(conn);
+	if (cancel_conn == NULL)
+		return _("out of memory");
+
+	/* In what follows, do not leak any PGcancelConn on any errors. */
+
+	PG_TRY();
+	{
+		if (!PQcancelStart(cancel_conn))
+		{
+			error = pchomp(PQcancelErrorMessage(cancel_conn));
+			goto exit;
+		}
+
+		for (;;)
+		{
+			PostgresPollingStatusType pollres;
+			TimestampTz now;
+			long		cur_timeout;
+			int			waitEvents = WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH;
+
+			pollres = PQcancelPoll(cancel_conn);
+			if (pollres == PGRES_POLLING_OK)
+				break;			/* success! */
+
+			/* If timeout has expired, give up, else get sleep time. */
+			now = GetCurrentTimestamp();
+			cur_timeout = TimestampDifferenceMilliseconds(now, endtime);
+			if (cur_timeout <= 0)
+			{
+				error = _("cancel request timed out");
+				break;
+			}
+
+			switch (pollres)
+			{
+				case PGRES_POLLING_READING:
+					waitEvents |= WL_SOCKET_READABLE;
+					break;
+				case PGRES_POLLING_WRITING:
+					waitEvents |= WL_SOCKET_WRITEABLE;
+					break;
+				default:
+					error = pchomp(PQcancelErrorMessage(cancel_conn));
+					goto exit;
+			}
+
+			/* Sleep until there's something to do */
+			WaitLatchOrSocket(MyLatch, waitEvents, PQcancelSocket(cancel_conn),
+							  cur_timeout, PG_WAIT_CLIENT);
+
+			ResetLatch(MyLatch);
+
+			CHECK_FOR_INTERRUPTS();
+		}
+exit:	;
+	}
+	PG_FINALLY();
+	{
+		PQcancelFinish(cancel_conn);
+	}
+	PG_END_TRY();
+
+	return error;
+}
+
 #endif							/* LIBPQ_BE_FE_HELPERS_H */