Invent a new, more thread-safe version of PQrequestCancel, called PQcancel.

Use this new function in psql.  Implement query cancellation in psql for
Windows.  Code by Magnus Hagander, documentation and minor editorialization
by Tom Lane.

7 files changed, 264 insertions, 69 deletions

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 0ca6fa4f557..0089d5d2c81 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2004, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.92 2004/10/10 23:37:40 neilc Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.93 2004/10/30 23:10:50 tgl Exp $
  */
 #include "postgres_fe.h"
 #include "common.h"
@@ -223,23 +223,33 @@ NoticeProcessor(void *arg, const char *message)
  *
  * Before we start a query, we enable a SIGINT signal catcher that sends a
  * cancel request to the backend. Note that sending the cancel directly from
- * the signal handler is safe because PQrequestCancel() is written to make it
- * so. We use write() to print to stdout because it's better to use simple
+ * the signal handler is safe because PQcancel() is written to make it
+ * so. We use write() to print to stderr because it's better to use simple
  * facilities in a signal handler.
+ *
+ * On win32, the signal cancelling 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 other
+ * thread is using it.
  */
-static PGconn *volatile cancelConn = NULL;
+static PGcancel *cancelConn = NULL;
+#ifdef WIN32
+static CRITICAL_SECTION cancelConnLock;
+#endif
 
 volatile bool cancel_pressed = false;
 
+#define write_stderr(str)	write(fileno(stderr), str, strlen(str))
 
-#ifndef WIN32
 
-#define write_stderr(String) write(fileno(stderr), String, strlen(String))
+#ifndef WIN32
 
 void
 handle_sigint(SIGNAL_ARGS)
 {
 	int			save_errno = errno;
+	char        errbuf[256];
 
 	/* Don't muck around if prompting for a password. */
 	if (prompt_state)
@@ -250,17 +260,60 @@ handle_sigint(SIGNAL_ARGS)
 
 	cancel_pressed = true;
 
-	if (PQrequestCancel(cancelConn))
+	if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 		write_stderr("Cancel request sent\n");
 	else
 	{
 		write_stderr("Could not send cancel request: ");
-		write_stderr(PQerrorMessage(cancelConn));
+		write_stderr(errbuf);
 	}
 	errno = save_errno;			/* just in case the write changed it */
 }
-#endif   /* not WIN32 */
 
+#else /* WIN32 */
+
+static BOOL WINAPI
+consoleHandler(DWORD dwCtrlType)
+{
+	char        errbuf[256];
+
+	if (dwCtrlType == CTRL_C_EVENT ||
+		dwCtrlType == CTRL_BREAK_EVENT)
+	{
+		if (prompt_state)
+			return TRUE;
+
+		/* Perform query cancel */
+		EnterCriticalSection(&cancelConnLock);
+		if (cancelConn != NULL)
+		{
+			cancel_pressed = true;
+
+			if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
+				write_stderr("Cancel request sent\n");
+			else
+			{
+				write_stderr("Could not send cancel request: ");
+				write_stderr(errbuf);
+			}
+		}
+		LeaveCriticalSection(&cancelConnLock);
+
+		return TRUE;
+	}
+	else
+		/* Return FALSE for any signals not being handled */
+		return FALSE;
+}
+
+void
+setup_cancel_handler(void)
+{
+	InitializeCriticalSection(&cancelConnLock);
+	SetConsoleCtrlHandler(consoleHandler, TRUE);
+}
+
+#endif /* WIN32 */
 
 
 /* ConnectionUp
@@ -327,20 +380,42 @@ CheckConnection(void)
 static void
 SetCancelConn(void)
 {
-	cancelConn = pset.db;
+#ifdef WIN32
+	EnterCriticalSection(&cancelConnLock);
+#endif
+
+	/* Free the old one if we have one */
+	if (cancelConn != NULL)
+		PQfreeCancel(cancelConn);
+
+	cancelConn = PQgetCancel(pset.db);
+
+#ifdef WIN32
+	LeaveCriticalSection(&cancelConnLock);
+#endif
 }
 
 
 /*
  * ResetCancelConn
  *
- * Set cancelConn to NULL.	I don't know what this means exactly, but it saves
- * having to export the variable.
+ * Free the current cancel connection, if any, and set to NULL.
  */
 void
 ResetCancelConn(void)
 {
+#ifdef WIN32
+	EnterCriticalSection(&cancelConnLock);
+#endif
+
+	if (cancelConn)
+		PQfreeCancel(cancelConn);
+
 	cancelConn = NULL;
+
+#ifdef WIN32
+	LeaveCriticalSection(&cancelConnLock);
+#endif
 }
 
 
diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h
index e9404ae256f..cf917cedca9 100644
--- a/src/bin/psql/common.h
+++ b/src/bin/psql/common.h
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2004, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/common.h,v 1.39 2004/08/29 04:13:02 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/common.h,v 1.40 2004/10/30 23:10:50 tgl Exp $
  */
 #ifndef COMMON_H
 #define COMMON_H
@@ -48,7 +48,9 @@ extern void ResetCancelConn(void);
 
 #ifndef WIN32
 extern void handle_sigint(SIGNAL_ARGS);
-#endif   /* not WIN32 */
+#else
+extern void setup_cancel_handler(void);
+#endif
 
 extern PGresult *PSQLexec(const char *query, bool start_xact);
 
diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c
index c555710a347..d9f3395b1a2 100644
--- a/src/bin/psql/mainloop.c
+++ b/src/bin/psql/mainloop.c
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2004, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/mainloop.c,v 1.64 2004/08/29 05:06:54 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/mainloop.c,v 1.65 2004/10/30 23:10:50 tgl Exp $
  */
 #include "postgres_fe.h"
 #include "mainloop.h"
@@ -121,7 +121,10 @@ MainLoop(FILE *source)
 		 * ready
 		 */
 		pqsignal(SIGINT, handle_sigint);		/* control-C => cancel */
-#endif   /* not WIN32 */
+
+#else /* WIN32 */
+		setup_cancel_handler();
+#endif
 
 		fflush(stdout);
 
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index e14a8bb58a9..363764fffe9 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -1,4 +1,4 @@
-# $PostgreSQL: pgsql/src/interfaces/libpq/exports.txt,v 1.2 2004/10/18 22:00:42 tgl Exp $
+# $PostgreSQL: pgsql/src/interfaces/libpq/exports.txt,v 1.3 2004/10/30 23:11:26 tgl Exp $
 # Functions to be exported by libpq DLLs
 PQconnectdb               1
 PQsetdbLogin              2
@@ -119,3 +119,6 @@ pg_valid_server_encoding  116
 pqsignal                  117
 PQprepare                 118
 PQsendPrepare             119
+PQgetCancel               120
+PQfreeCancel              121
+PQcancel                  122
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 3da9caecc6d..f44e1f606db 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/interfaces/libpq/fe-connect.c,v 1.288 2004/10/29 19:30:02 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/interfaces/libpq/fe-connect.c,v 1.289 2004/10/30 23:11:26 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -2189,19 +2189,55 @@ PQresetPoll(PGconn *conn)
 	return PGRES_POLLING_FAILED;
 }
 
+/*
+ * PQcancelGet: get a PGcancel structure corresponding to a connection.
+ *
+ * A copy is needed to be able to cancel a running query from a different
+ * thread. If the same structure is used all structure members would have
+ * to be individually locked (if the entire structure was locked, it would
+ * be impossible to cancel a synchronous query becuase the structure would
+ * have to stay locked for the duration of the query).
+ */
+PGcancel *
+PQgetCancel(PGconn *conn)
+{
+	PGcancel *cancel;
+
+	if (!conn)
+		return NULL;
+
+	if (conn->sock < 0)
+		return NULL;
+
+	cancel = malloc(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;
+
+	return cancel;
+}
+
+/* PQfreeCancel: free a cancel structure */
+void
+PQfreeCancel(PGcancel *cancel)
+{
+	if (cancel)
+		free(cancel);
+}
+
 
 /*
- * PQrequestCancel: attempt to request cancellation of the current operation.
+ * PQcancel and PQrequestCancel: attempt to request cancellation of the
+ * current operation.
  *
  * The return value is TRUE if the cancel request was successfully
- * dispatched, FALSE if not (in which case conn->errorMessage is set).
+ * dispatched, FALSE if not (in which case an error message is available).
  * Note: successful dispatch is no guarantee that there will be any effect at
  * the backend.  The application must read the operation result as usual.
  *
- * XXX it was a bad idea to have the error message returned in
- * conn->errorMessage, since it could overwrite a message already there.
- * Would be better to return it in a char array passed by the caller.
- *
  * CAUTION: we want this routine to be safely callable from a signal handler
  * (for example, an application might want to call it in a SIGINT handler).
  * This means we cannot use any C library routine that might be non-reentrant.
@@ -2210,59 +2246,40 @@ PQresetPoll(PGconn *conn)
  * error messages with strcpy/strcat is tedious but should be quite safe.
  * We also save/restore errno in case the signal handler support doesn't.
  *
- * NOTE: this routine must not generate any error message longer than
- * INITIAL_EXPBUFFER_SIZE (currently 256), since we dare not try to
- * expand conn->errorMessage!
+ * internal_cancel() is an internal helper function to make code-sharing
+ * between the two versions of the cancel function possible.
  */
-
-int
-PQrequestCancel(PGconn *conn)
+static int
+internal_cancel(SockAddr *raddr, int be_pid, int be_key,
+				char *errbuf, int errbufsize)
 {
 	int			save_errno = SOCK_ERRNO;
 	int			tmpsock = -1;
 	char		sebuf[256];
+	int         maxlen;
 	struct
 	{
 		uint32		packetlen;
 		CancelRequestPacket cp;
 	}			crp;
 
-	/* Check we have an open connection */
-	if (!conn)
-		return FALSE;
-
-	if (conn->sock < 0)
-	{
-		strcpy(conn->errorMessage.data,
-			   "PQrequestCancel() -- connection is not open\n");
-		conn->errorMessage.len = strlen(conn->errorMessage.data);
-#ifdef WIN32
-		WSASetLastError(save_errno);
-#else
-		errno = save_errno;
-#endif
-		return FALSE;
-	}
-
 	/*
-	 * We need to open a temporary connection to the postmaster. Use the
-	 * information saved by connectDB to do this with only kernel calls.
+	 * We need to open a temporary connection to the postmaster. Do
+	 * this with only kernel calls.
 	 */
-	if ((tmpsock = socket(conn->raddr.addr.ss_family, SOCK_STREAM, 0)) < 0)
+	if ((tmpsock = socket(raddr->addr.ss_family, SOCK_STREAM, 0)) < 0)
 	{
-		strcpy(conn->errorMessage.data,
-			   "PQrequestCancel() -- socket() failed: ");
+		StrNCpy(errbuf, "PQcancel() -- socket() failed: ", errbufsize);
 		goto cancel_errReturn;
 	}
 retry3:
-	if (connect(tmpsock, (struct sockaddr *) & conn->raddr.addr,
-				conn->raddr.salen) < 0)
+	if (connect(tmpsock, (struct sockaddr *) & raddr->addr,
+				raddr->salen) < 0)
 	{
 		if (SOCK_ERRNO == EINTR)
 			/* Interrupted system call - we'll just try again */
 			goto retry3;
-		strcpy(conn->errorMessage.data,
-			   "PQrequestCancel() -- connect() failed: ");
+		StrNCpy(errbuf, "PQcancel() -- connect() failed: ", errbufsize);
 		goto cancel_errReturn;
 	}
 
@@ -2274,8 +2291,8 @@ retry3:
 
 	crp.packetlen = htonl((uint32) sizeof(crp));
 	crp.cp.cancelRequestCode = (MsgType) htonl(CANCEL_REQUEST_CODE);
-	crp.cp.backendPID = htonl(conn->be_pid);
-	crp.cp.cancelAuthCode = htonl(conn->be_key);
+	crp.cp.backendPID = htonl(be_pid);
+	crp.cp.cancelAuthCode = htonl(be_key);
 
 retry4:
 	if (send(tmpsock, (char *) &crp, sizeof(crp), 0) != (int) sizeof(crp))
@@ -2283,8 +2300,7 @@ retry4:
 		if (SOCK_ERRNO == EINTR)
 			/* Interrupted system call - we'll just try again */
 			goto retry4;
-		strcpy(conn->errorMessage.data,
-			   "PQrequestCancel() -- send() failed: ");
+		StrNCpy(errbuf, "PQcancel() -- send() failed: ", errbufsize);
 		goto cancel_errReturn;
 	}
 
@@ -2316,21 +2332,90 @@ retry5:
 	return TRUE;
 
 cancel_errReturn:
-	strcat(conn->errorMessage.data, SOCK_STRERROR(SOCK_ERRNO, sebuf, sizeof(sebuf)));
-	strcat(conn->errorMessage.data, "\n");
-	conn->errorMessage.len = strlen(conn->errorMessage.data);
-	if (tmpsock >= 0)
+	/*
+	 * Make sure we don't overflow the error buffer. Leave space for
+	 * the \n at the end, and for the terminating zero.
+	 */
+	maxlen = errbufsize - strlen(errbuf) - 2;
+	if (maxlen >= 0)
 	{
+		strncat(errbuf, SOCK_STRERROR(SOCK_ERRNO, sebuf, sizeof(sebuf)),
+				maxlen);
+		strcat(errbuf, "\n");
+	}
+	if (tmpsock >= 0)
 		closesocket(tmpsock);
 #ifdef WIN32
-		WSASetLastError(save_errno);
+	WSASetLastError(save_errno);
 #else
-		errno = save_errno;
+	errno = save_errno;
 #endif
-	}
+
 	return FALSE;
 }
 
+/*
+ * PQcancel: request query cancel
+ *
+ * Returns TRUE if able to send the cancel request, FALSE if not.
+ *
+ * On failure, an error message is stored in *errbuf, which must be of size
+ * errbufsize (recommended size is 256 bytes).  *errbuf is not changed on
+ * success return.
+ */
+int
+PQcancel(PGcancel *cancel, char *errbuf, int errbufsize)
+{
+	if (!cancel)
+	{
+		StrNCpy(errbuf, "PQcancel() -- no cancel object supplied", errbufsize);
+		return FALSE;
+	}
+
+	return internal_cancel(&cancel->raddr, cancel->be_pid, cancel->be_key,
+						   errbuf, errbufsize);
+}
+
+/*
+ * PQrequestCancel: old, not thread-safe function for requesting query cancel
+ *
+ * Returns TRUE if able to send the cancel request, FALSE if not.
+ *
+ * On failure, the error message is saved in conn->errorMessage; this means
+ * that this can't be used when there might be other active operations on
+ * the connection object.
+ *
+ * NOTE: error messages will be cut off at the current size of the
+ * error message buffer, since we dare not try to expand conn->errorMessage!
+ */
+int
+PQrequestCancel(PGconn *conn)
+{
+	int r;
+
+	/* Check we have an open connection */
+	if (!conn)
+		return FALSE;
+
+	if (conn->sock < 0)
+	{
+		StrNCpy(conn->errorMessage.data,
+				"PQrequestCancel() -- connection is not open\n",
+				conn->errorMessage.maxlen);
+		conn->errorMessage.len = strlen(conn->errorMessage.data);
+
+		return FALSE;
+	}
+
+	r = internal_cancel(&conn->raddr, conn->be_pid, conn->be_key,
+						conn->errorMessage.data, conn->errorMessage.maxlen);
+
+	if (!r)
+		conn->errorMessage.len = strlen(conn->errorMessage.data);
+
+	return r;
+}
+
 
 /*
  * pqPacketSend() -- convenience routine to send a message to server.
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 1fba91bde4e..8d73fc8bf59 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -7,7 +7,7 @@
  * Portions Copyright (c) 1996-2004, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/interfaces/libpq/libpq-fe.h,v 1.112 2004/10/18 22:00:42 tgl Exp $
+ * $PostgreSQL: pgsql/src/interfaces/libpq/libpq-fe.h,v 1.113 2004/10/30 23:11:27 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -116,6 +116,12 @@ typedef struct pg_conn PGconn;
  */
 typedef struct pg_result PGresult;
 
+/* PGcancel encapsulates the information needed to cancel a running
+ * query on an existing connection.
+ * The contents of this struct are not supposed to be known to applications.
+ */
+typedef struct pg_cancel PGcancel;
+
 /* PGnotify represents the occurrence of a NOTIFY message.
  * Ideally this would be an opaque typedef, but it's so simple that it's
  * unlikely to change.
@@ -234,7 +240,16 @@ extern PostgresPollingStatusType PQresetPoll(PGconn *conn);
 /* Synchronous (blocking) */
 extern void PQreset(PGconn *conn);
 
+/* request a cancel structure */
+extern PGcancel *PQgetCancel(PGconn *conn);
+
+/* free a cancel structure */
+extern void PQfreeCancel(PGcancel *cancel);
+
 /* issue a cancel request */
+extern int	PQcancel(PGcancel *cancel, char *errbuf, int errbufsize);
+
+/* backwards compatible version of PQcancel; not thread-safe */
 extern int	PQrequestCancel(PGconn *conn);
 
 /* Accessor functions for PGconn objects */
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index d1fab1395b2..343834f0a20 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -12,7 +12,7 @@
  * Portions Copyright (c) 1996-2004, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/interfaces/libpq/libpq-int.h,v 1.95 2004/10/18 22:00:42 tgl Exp $
+ * $PostgreSQL: pgsql/src/interfaces/libpq/libpq-int.h,v 1.96 2004/10/30 23:11:27 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -346,6 +346,18 @@ struct pg_conn
 	PQExpBufferData workBuffer; /* expansible string */
 };
 
+/* PGcancel stores all data necessary to cancel a connection. A copy of this
+ * data is required to safely cancel a connection running on a different
+ * thread.
+ */
+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 */
+};
+
+
 /* String descriptions of the ExecStatusTypes.
  * direct use of this array is deprecated; call PQresStatus() instead.
  */