Refactor query cancellation code into src/fe_utils/

Originally, this code was duplicated in src/bin/psql/ and
src/bin/scripts/, but it can be useful for other frontend applications,
like pgbench.  This refactoring offers the possibility to setup a custom
callback which would get called in the signal handler for SIGINT or when
the interruption console events happen on Windows.

Author: Fabien Coelho, with contributions from Michael Paquier
Reviewed-by: Álvaro Herrera, Ibrar Ahmed
Discussion: https://postgr.es/m/alpine.DEB.2.21.1910311939430.27369@lancre

src/fe_utils/cancel.c

@@ -0,0 +1,225 @@
+/*------------------------------------------------------------------------
+ *
+ * Query cancellation support for frontend code
+ *
+ * Assorted utility functions to control query cancellation with signal
+ * handler for SIGINT.
+ *
+ *
+ * Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/fe-utils/cancel.c
+ *
+ *------------------------------------------------------------------------
+ */
+
+#include "postgres_fe.h"
+
+#include <signal.h>
+#include <unistd.h>
+
+#include "fe_utils/cancel.h"
+#include "fe_utils/connect.h"
+#include "fe_utils/string_utils.h"
+
+
+/*
+ * Write a simple string to stderr --- must be safe in a signal handler.
+ * We ignore the write() result since there's not much we could do about it.
+ * Certain compilers make that harder than it ought to be.
+ */
+#define write_stderr(str) \
+	do { \
+		const char *str_ = (str); \
+		int		rc_; \
+		rc_ = write(fileno(stderr), str_, strlen(str_)); \
+		(void) rc_; \
+	} while (0)
+
+static PGcancel *volatile cancelConn = NULL;
+bool		CancelRequested = false;
+
+#ifdef WIN32
+static CRITICAL_SECTION cancelConnLock;
+#endif
+
+/*
+ * Additional callback for cancellations.
+ */
+static void (*cancel_callback) (void) = NULL;
+
+
+/*
+ * SetCancelConn
+ *
+ * Set cancelConn to point to the current database connection.
+ */
+void
+SetCancelConn(PGconn *conn)
+{
+	PGcancel   *oldCancelConn;
+
+#ifdef WIN32
+	EnterCriticalSection(&cancelConnLock);
+#endif
+
+	/* Free the old one if we have one */
+	oldCancelConn = cancelConn;
+
+	/* be sure handle_sigint doesn't use pointer while freeing */
+	cancelConn = NULL;
+
+	if (oldCancelConn != NULL)
+		PQfreeCancel(oldCancelConn);
+
+	cancelConn = PQgetCancel(conn);
+
+#ifdef WIN32
+	LeaveCriticalSection(&cancelConnLock);
+#endif
+}
+
+/*
+ * ResetCancelConn
+ *
+ * Free the current cancel connection, if any, and set to NULL.
+ */
+void
+ResetCancelConn(void)
+{
+	PGcancel   *oldCancelConn;
+
+#ifdef WIN32
+	EnterCriticalSection(&cancelConnLock);
+#endif
+
+	oldCancelConn = cancelConn;
+
+	/* be sure handle_sigint doesn't use pointer while freeing */
+	cancelConn = NULL;
+
+	if (oldCancelConn != NULL)
+		PQfreeCancel(oldCancelConn);
+
+#ifdef WIN32
+	LeaveCriticalSection(&cancelConnLock);
+#endif
+}
+
+
+/*
+ * Code to support query cancellation
+ *
+ * Note that sending the cancel directly from the signal handler is safe
+ * because PQcancel() is written to make it so.  We use write() to report
+ * to stderr because it's better to use simple facilities in a signal
+ * handler.
+ *
+ * On Windows, the signal canceling happens on a separate thread, because
+ * that's how SetConsoleCtrlHandler works.  The PQcancel function is safe
+ * for this (unlike PQrequestCancel).  However, a CRITICAL_SECTION is required
+ * to protect the PGcancel structure against being changed while the signal
+ * thread is using it.
+ */
+
+#ifndef WIN32
+
+/*
+ * handle_sigint
+ *
+ * Handle interrupt signals by canceling the current command, if cancelConn
+ * is set.
+ */
+static void
+handle_sigint(SIGNAL_ARGS)
+{
+	int			save_errno = errno;
+	char		errbuf[256];
+
+	if (cancel_callback != NULL)
+		cancel_callback();
+
+	/* Send QueryCancel if we are processing a database query */
+	if (cancelConn != NULL)
+	{
+		if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
+		{
+			CancelRequested = true;
+			write_stderr(_("Cancel request sent\n"));
+		}
+		else
+		{
+			write_stderr(_("Could not send cancel request: "));
+			write_stderr(errbuf);
+		}
+	}
+	else
+		CancelRequested = true;
+
+	errno = save_errno;			/* just in case the write changed it */
+}
+
+/*
+ * setup_cancel_handler
+ *
+ * Register query cancellation callback for SIGINT.
+ */
+void
+setup_cancel_handler(void (*callback) (void))
+{
+	cancel_callback = callback;
+	pqsignal(SIGINT, handle_sigint);
+}
+
+#else							/* WIN32 */
+
+static BOOL WINAPI
+consoleHandler(DWORD dwCtrlType)
+{
+	char		errbuf[256];
+
+	if (dwCtrlType == CTRL_C_EVENT ||
+		dwCtrlType == CTRL_BREAK_EVENT)
+	{
+		if (cancel_callback != NULL)
+			cancel_callback();
+
+		/* Send QueryCancel if we are processing a database query */
+		EnterCriticalSection(&cancelConnLock);
+		if (cancelConn != NULL)
+		{
+			if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
+			{
+				write_stderr(_("Cancel request sent\n"));
+				CancelRequested = true;
+			}
+			else
+			{
+				write_stderr(_("Could not send cancel request: %s"));
+				write_stderr(errbuf);
+			}
+		}
+		else
+			CancelRequested = true;
+
+		LeaveCriticalSection(&cancelConnLock);
+
+		return TRUE;
+	}
+	else
+		/* Return FALSE for any signals not being handled */
+		return FALSE;
+}
+
+void
+setup_cancel_handler(void (*callback) (void))
+{
+	cancel_callback = callback;
+
+	InitializeCriticalSection(&cancelConnLock);
+
+	SetConsoleCtrlHandler(consoleHandler, TRUE);
+}
+
+#endif							/* WIN32 */