From f6692187a3c0f2e5d5472a92499c7443927a9a3b Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <daniel@yesql.se>
Date: Sun, 18 Jun 2017 09:00:52 +0200
Subject: [PATCH] Support optional message to pg_terminate_backend

This adds the ability for the caller of pg_terminate_backend() to
include an optional message to the process which is being killed.
The message will be appended to the error message returned to the
killed process. The new syntax is overloaded as:

    SELECT pg_terminate_backend(<pid> [, msg]);
    SELECT pg_cancel_backend(<pid> [, msg]);
---
 doc/src/sgml/func.sgml                   |   6 +-
 src/backend/storage/ipc/ipci.c           |   2 +
 src/backend/storage/lmgr/lwlocknames.txt |   1 +
 src/backend/tcop/postgres.c              |  38 +++++--
 src/backend/utils/adt/misc.c             |  50 +++++++--
 src/backend/utils/init/postinit.c        |   1 +
 src/backend/utils/misc/Makefile          |   6 +-
 src/backend/utils/misc/backend_cancel.c  | 167 +++++++++++++++++++++++++++++++
 src/include/catalog/pg_proc.h            |   4 +
 src/include/utils/backend_cancel.h       |  25 +++++
 10 files changed, 281 insertions(+), 19 deletions(-)
 create mode 100644 src/backend/utils/misc/backend_cancel.c
 create mode 100644 src/include/utils/backend_cancel.h

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index e073f7b57c..4e4d0f1ad6 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -18330,7 +18330,7 @@ SELECT set_config('log_statement_stats', 'off', false);
      <tbody>
       <row>
        <entry>
-        <literal><function>pg_cancel_backend(<parameter>pid</parameter> <type>int</>)</function></literal>
+        <literal><function>pg_cancel_backend(<parameter>pid</parameter> <type>int</> [, <parameter>message</parameter> <type>text</>])</function></literal>
         </entry>
        <entry><type>boolean</type></entry>
        <entry>Cancel a backend's current query.  This is also allowed if the
@@ -18355,7 +18355,7 @@ SELECT set_config('log_statement_stats', 'off', false);
       </row>
       <row>
        <entry>
-        <literal><function>pg_terminate_backend(<parameter>pid</parameter> <type>int</>)</function></literal>
+        <literal><function>pg_terminate_backend(<parameter>pid</parameter> <type>int</> [, <parameter>message</parameter> <type>text</>])</function></literal>
         </entry>
        <entry><type>boolean</type></entry>
        <entry>Terminate a backend.  This is also allowed if the calling role
@@ -18386,6 +18386,8 @@ SELECT set_config('log_statement_stats', 'off', false);
     The role of an active backend can be found from the
     <structfield>usename</structfield> column of the
     <structname>pg_stat_activity</structname> view.
+    If the optional <literal>message</literal> parameter is set, the text
+    will be appended to the error message returned to the signalled backend.
    </para>
 
    <para>
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 2d1ed143e0..63ac4f0ab3 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -150,6 +150,7 @@ CreateSharedMemoryAndSemaphores(bool makePrivate, int port)
 		size = add_size(size, SyncScanShmemSize());
 		size = add_size(size, AsyncShmemSize());
 		size = add_size(size, BackendRandomShmemSize());
+		size = add_size(size, CancelBackendMsgShmemSize());
 #ifdef EXEC_BACKEND
 		size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -270,6 +271,7 @@ CreateSharedMemoryAndSemaphores(bool makePrivate, int port)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	BackendRandomShmemInit();
+	BackendCancelShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/storage/lmgr/lwlocknames.txt b/src/backend/storage/lmgr/lwlocknames.txt
index e6025ecedb..772ff518ee 100644
--- a/src/backend/storage/lmgr/lwlocknames.txt
+++ b/src/backend/storage/lmgr/lwlocknames.txt
@@ -50,3 +50,4 @@ OldSnapshotTimeMapLock				42
 BackendRandomLock					43
 LogicalRepWorkerLock				44
 CLogTruncationLock					45
+BackendCancelLock					46
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index f99dd0a2d4..1c886426c0 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -69,6 +69,7 @@
 #include "tcop/pquery.h"
 #include "tcop/tcopprot.h"
 #include "tcop/utility.h"
+#include "utils/backend_cancel.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/ps_status.h"
@@ -2876,9 +2877,22 @@ ProcessInterrupts(void)
 					 errdetail_recovery_conflict()));
 		}
 		else
-			ereport(FATAL,
-					(errcode(ERRCODE_ADMIN_SHUTDOWN),
-			 errmsg("terminating connection due to administrator command")));
+		{
+			if (HasCancelMessage())
+			{
+				char   *buffer = palloc0(MAX_CANCEL_MSG);
+
+				GetCancelMessage(&buffer, MAX_CANCEL_MSG);
+				ereport(FATAL,
+						(errcode(ERRCODE_ADMIN_SHUTDOWN),
+				 errmsg("terminating connection due to administrator command: \"%s\"",
+						buffer)));
+			}
+			else
+				ereport(FATAL,
+						(errcode(ERRCODE_ADMIN_SHUTDOWN),
+				 errmsg("terminating connection due to administrator command")));
+		}
 	}
 	if (ClientConnectionLost)
 	{
@@ -2991,9 +3005,21 @@ ProcessInterrupts(void)
 		if (!DoingCommandRead)
 		{
 			LockErrorCleanup();
-			ereport(ERROR,
-					(errcode(ERRCODE_QUERY_CANCELED),
-					 errmsg("canceling statement due to user request")));
+
+			if (HasCancelMessage())
+			{
+				char   *buffer = palloc0(MAX_CANCEL_MSG);
+
+				GetCancelMessage(&buffer, MAX_CANCEL_MSG);
+				ereport(ERROR,
+						(errcode(ERRCODE_QUERY_CANCELED),
+						 errmsg("canceling statement due to user request: \"%s\"",
+								buffer)));
+			}
+			else
+				ereport(ERROR,
+						(errcode(ERRCODE_QUERY_CANCELED),
+						 errmsg("canceling statement due to user request")));
 		}
 	}
 
diff --git a/src/backend/utils/adt/misc.c b/src/backend/utils/adt/misc.c
index 9cc0b08e96..0abbe47c08 100644
--- a/src/backend/utils/adt/misc.c
+++ b/src/backend/utils/adt/misc.c
@@ -41,6 +41,7 @@
 #include "utils/ruleutils.h"
 #include "tcop/tcopprot.h"
 #include "utils/acl.h"
+#include "utils/backend_cancel.h"
 #include "utils/builtins.h"
 #include "utils/timestamp.h"
 
@@ -216,7 +217,7 @@ current_query(PG_FUNCTION_ARGS)
 #define SIGNAL_BACKEND_NOPERMISSION 2
 #define SIGNAL_BACKEND_NOSUPERUSER 3
 static int
-pg_signal_backend(int pid, int sig)
+pg_signal_backend(int pid, int sig, char *msg)
 {
 	PGPROC	   *proc = BackendPidGetProc(pid);
 
@@ -257,6 +258,9 @@ pg_signal_backend(int pid, int sig)
 	 * too unlikely to worry about.
 	 */
 
+	if (msg != NULL)
+		SetBackendCancelMessage(pid, msg);
+
 	/* If we have setsid(), signal the backend's whole process group */
 #ifdef HAVE_SETSID
 	if (kill(-pid, sig))
@@ -278,10 +282,10 @@ pg_signal_backend(int pid, int sig)
  *
  * Note that only superusers can signal superuser-owned processes.
  */
-Datum
-pg_cancel_backend(PG_FUNCTION_ARGS)
+static bool
+pg_cancel_backend_internal(pid_t pid, char *msg)
 {
-	int			r = pg_signal_backend(PG_GETARG_INT32(0), SIGINT);
+	int			r = pg_signal_backend(pid, SIGINT, msg);
 
 	if (r == SIGNAL_BACKEND_NOSUPERUSER)
 		ereport(ERROR,
@@ -296,16 +300,31 @@ pg_cancel_backend(PG_FUNCTION_ARGS)
 	PG_RETURN_BOOL(r == SIGNAL_BACKEND_SUCCESS);
 }
 
+Datum
+pg_cancel_backend(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(pg_cancel_backend_internal(PG_GETARG_INT32(0), NULL));
+}
+Datum
+pg_cancel_backend_msg(PG_FUNCTION_ARGS)
+{
+	pid_t		pid = PG_GETARG_INT32(0);
+	char 	   *msg = text_to_cstring(PG_GETARG_TEXT_PP(1));
+
+	PG_RETURN_BOOL(pg_cancel_backend_internal(pid, msg));
+}
+
+
 /*
  * Signal to terminate a backend process.  This is allowed if you are a member
  * of the role whose process is being terminated.
  *
  * Note that only superusers can signal superuser-owned processes.
  */
-Datum
-pg_terminate_backend(PG_FUNCTION_ARGS)
+static bool
+pg_terminate_backend_internal(pid_t pid, char *msg)
 {
-	int			r = pg_signal_backend(PG_GETARG_INT32(0), SIGTERM);
+	int r = pg_signal_backend(pid, SIGTERM, msg);
 
 	if (r == SIGNAL_BACKEND_NOSUPERUSER)
 		ereport(ERROR,
@@ -317,7 +336,22 @@ pg_terminate_backend(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
 				 (errmsg("must be a member of the role whose process is being terminated or member of pg_signal_backend"))));
 
-	PG_RETURN_BOOL(r == SIGNAL_BACKEND_SUCCESS);
+	return (r == SIGNAL_BACKEND_SUCCESS);
+}
+
+Datum
+pg_terminate_backend(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(pg_terminate_backend_internal(PG_GETARG_INT32(0), NULL));
+}
+
+Datum
+pg_terminate_backend_msg(PG_FUNCTION_ARGS)
+{
+	pid_t		pid = PG_GETARG_INT32(0);
+	char 	   *msg = text_to_cstring(PG_GETARG_TEXT_PP(1));
+
+	PG_RETURN_BOOL(pg_terminate_backend_internal(pid, msg));
 }
 
 /*
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index b8b4a06350..cbf6cb4f60 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -740,6 +740,7 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 		PerformAuthentication(MyProcPort);
 		InitializeSessionUserId(username, useroid);
 		am_superuser = superuser();
+		BackendCancelInit(MyBackendId);
 	}
 
 	/*
diff --git a/src/backend/utils/misc/Makefile b/src/backend/utils/misc/Makefile
index a53fcdf188..619c837e08 100644
--- a/src/backend/utils/misc/Makefile
+++ b/src/backend/utils/misc/Makefile
@@ -14,9 +14,9 @@ include $(top_builddir)/src/Makefile.global
 
 override CPPFLAGS := -I. -I$(srcdir) $(CPPFLAGS)
 
-OBJS = backend_random.o guc.o help_config.o pg_config.o pg_controldata.o \
-       pg_rusage.o ps_status.o queryenvironment.o rls.o sampling.o \
-       superuser.o timeout.o tzparser.o
+OBJS = backend_cancel.o backend_random.o guc.o help_config.o pg_config.o \
+       pg_controldata.o pg_rusage.o ps_status.o queryenvironment.o rls.o \
+	   sampling.o superuser.o timeout.o tzparser.o
 
 # This location might depend on the installation directories. Therefore
 # we can't substitute it into pg_config.h.
diff --git a/src/backend/utils/misc/backend_cancel.c b/src/backend/utils/misc/backend_cancel.c
new file mode 100644
index 0000000000..c1628fe574
--- /dev/null
+++ b/src/backend/utils/misc/backend_cancel.c
@@ -0,0 +1,167 @@
+/*-------------------------------------------------------------------------
+ *
+ * backend_cancel.c
+ *	  Backend cancellation messaging
+ *
+ *
+ * Module for supporting passing a user defined message to a cancelled,
+ * or terminated, backend from the user/administrator.
+ *
+ * Portions Copyright (c) 1996-2017, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/backend/utils/misc/backend_cancel.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "miscadmin.h"
+#include "storage/ipc.h"
+#include "storage/lwlock.h"
+#include "storage/shmem.h"
+#include "utils/backend_cancel.h"
+
+/*
+ * Each backend is registered per pid in the array which is indexed by
+ * Backend ID. Reading and writing the message is protected by the
+ * BackendCancelLock lwlock.
+ */
+typedef struct
+{
+	pid_t	pid;
+	char	message[MAX_CANCEL_MSG];
+	int		len;
+} BackendCancelShmemStruct;
+
+static BackendCancelShmemStruct	*BackendCancelSlots = NULL;
+static volatile BackendCancelShmemStruct *MyCancelSlot = NULL;
+
+static void CleanupCancelBackend(int status, Datum argument);
+
+Size
+CancelBackendMsgShmemSize(void)
+{
+	return MaxBackends * sizeof(BackendCancelShmemStruct);
+}
+
+void
+BackendCancelShmemInit(void)
+{
+	Size	size = CancelBackendMsgShmemSize();
+	bool	found;
+
+	BackendCancelSlots = (BackendCancelShmemStruct *)
+		ShmemInitStruct("BackendCancelSlots", size, &found);
+
+	if (!found)
+		MemSet(BackendCancelSlots, 0, size);
+}
+
+void
+BackendCancelInit(int backend_id)
+{
+	volatile BackendCancelShmemStruct *slot;
+
+	slot = &BackendCancelSlots[backend_id - 1];
+
+	slot->message[0] = '\0';
+	slot->len = 0;
+	slot->pid = MyProcPid;
+
+	MyCancelSlot = slot;
+
+	on_shmem_exit(CleanupCancelBackend, Int32GetDatum(backend_id));
+}
+
+static void
+CleanupCancelBackend(int status, Datum argument)
+{
+	int backend_id = DatumGetInt32(argument);
+	volatile BackendCancelShmemStruct *slot;
+
+	slot = &BackendCancelSlots[backend_id - 1];
+
+	Assert(slot == MyCancelSlot);
+
+	MyCancelSlot = NULL;
+
+	if (slot->len > 0)
+		slot->message[0] = '\0';
+
+	slot->len = 0;
+	slot->pid = 0;
+}
+
+/*
+ * Sets a cancellation message for the backend with the specified pid and
+ * returns the length of message actually created. If the returned length
+ * is equal to the length of the message parameter, truncation may have
+ * occurred. If the backend wasn't found and no message was set, -1 is
+ * returned.
+ */
+int
+SetBackendCancelMessage(pid_t backend, char *message)
+{
+	BackendCancelShmemStruct *slot;
+	int		i;
+
+	if (!message)
+		return 0;
+
+
+	for (i = 0; i < MaxBackends; i++)
+	{
+		slot = &BackendCancelSlots[i];
+
+		if (slot->pid != 0 && slot->pid == backend)
+		{
+			LWLockAcquire(BackendCancelLock, LW_EXCLUSIVE);
+
+			strlcpy(slot->message, message, sizeof(slot->message));
+			slot->len = strlen(message);
+
+			LWLockRelease(BackendCancelLock);
+			return slot->len;
+		}
+	}
+
+	elog(LOG, "Cancellation message requested for missing backend %d by %d",
+		 (int) backend, MyProcPid);
+
+	return -1;
+}
+
+bool
+HasCancelMessage(void)
+{
+	volatile BackendCancelShmemStruct *slot = MyCancelSlot;
+
+	return (slot != NULL && slot->len > 0);
+}
+
+/*
+ * Return the configured cancellation message and its length. If the
+ * returned length is greater than, or equal to, the size of the passed
+ * buffer, truncation may have been performed. The message is cleared
+ * on reading.
+ */
+int
+GetCancelMessage(char **buffer, size_t buf_len)
+{
+	volatile BackendCancelShmemStruct *slot = MyCancelSlot;
+
+	if (slot != NULL && slot->len > 0)
+	{
+		LWLockAcquire(BackendCancelLock, LW_EXCLUSIVE);
+		strlcpy(*buffer, (const char *) slot->message, buf_len);
+		slot->len = 0;
+		slot->message[0] = '\0';
+		LWLockRelease(BackendCancelLock);
+
+		return slot->len;
+	}
+
+	return 0;
+}
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index 6c44def6e6..edd4a6cdb0 100644
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -3203,8 +3203,12 @@ DESCR("is schema another session's temp schema?");
 
 DATA(insert OID = 2171 ( pg_cancel_backend		PGNSP PGUID 12 1 0 0 0 f f f f t f v s 1 0 16 "23" _null_ _null_ _null_ _null_ _null_ pg_cancel_backend _null_ _null_ _null_ ));
 DESCR("cancel a server process' current query");
+DATA(insert OID = 772 ( pg_cancel_backend		PGNSP PGUID 12 1 0 0 0 f f f f t f v s 2 0 16 "23 25" _null_ _null_ _null_ _null_ _null_ pg_cancel_backend_msg _null_ _null_ _null_ ));
+DESCR("cancel a server process' current query");
 DATA(insert OID = 2096 ( pg_terminate_backend		PGNSP PGUID 12 1 0 0 0 f f f f t f v s 1 0 16 "23" _null_ _null_ _null_ _null_ _null_ pg_terminate_backend _null_ _null_ _null_ ));
 DESCR("terminate a server process");
+DATA(insert OID = 972 ( pg_terminate_backend		PGNSP PGUID 12 1 0 0 0 f f f f t f v s 2 0 16 "23 25" _null_ _null_ _null_ _null_ _null_ pg_terminate_backend_msg _null_ _null_ _null_ ));
+DESCR("terminate a server process");
 DATA(insert OID = 2172 ( pg_start_backup		PGNSP PGUID 12 1 0 0 0 f f f f t f v r 3 0 3220 "25 16 16" _null_ _null_ _null_ _null_ _null_ pg_start_backup _null_ _null_ _null_ ));
 DESCR("prepare for taking an online backup");
 DATA(insert OID = 2173 ( pg_stop_backup			PGNSP PGUID 12 1 0 0 0 f f f f t f v r 0 0 3220 "" _null_ _null_ _null_ _null_ _null_ pg_stop_backup _null_ _null_ _null_ ));
diff --git a/src/include/utils/backend_cancel.h b/src/include/utils/backend_cancel.h
new file mode 100644
index 0000000000..7f210553d6
--- /dev/null
+++ b/src/include/utils/backend_cancel.h
@@ -0,0 +1,25 @@
+/*-------------------------------------------------------------------------
+ *
+ * backend_cancel.h
+ *		Declarations for backend cancellation messaging
+ *
+ * Portions Copyright (c) 1996-2017, PostgreSQL Global Development Group
+ *
+ *	  src/include/utils/backend_cancel.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef BACKEND_CANCEL_H
+#define BACKEND_CANCEL_H
+
+#define MAX_CANCEL_MSG 128
+
+extern Size CancelBackendMsgShmemSize(void);
+extern void BackendCancelShmemInit(void);
+extern void BackendCancelInit(int backend_id);
+
+extern int SetBackendCancelMessage(pid_t backend, char *message);
+extern bool HasCancelMessage(void);
+extern int GetCancelMessage(char **msg, size_t len);
+
+#endif /* BACKEND_CANCEL_H */
-- 
2.13.0.rc0.45.ge2cb6ab.dirty

