From d88202f36d66d8bfe88fa34967ed23a2d2eade4a Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <daniel@yesql.se>
Date: Wed, 24 Jan 2018 12:22:20 +0100
Subject: [PATCH] Support optional message in backend cancel/terminate

This adds the ability for the caller of pg_terminate_backend() or
pg_cancel_backend() to include an optional message to the process
which is being signalled. The message will be appended to the error
message returned to the killed or cancelled process. The new syntax
overloaded the existing 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          |   3 +
 src/backend/tcop/postgres.c             |  38 ++++++-
 src/backend/utils/adt/misc.c            |  52 +++++++--
 src/backend/utils/init/postinit.c       |   2 +
 src/backend/utils/misc/Makefile         |   6 +-
 src/backend/utils/misc/backend_cancel.c | 195 ++++++++++++++++++++++++++++++++
 src/include/catalog/pg_proc.h           |   4 +
 src/include/utils/backend_cancel.h      |  25 ++++
 9 files changed, 312 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 487c7ff750..91f58eac55 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -18381,7 +18381,7 @@ SELECT set_config('log_statement_stats', 'off', false);
      <tbody>
       <row>
        <entry>
-        <literal><function>pg_cancel_backend(<parameter>pid</parameter> <type>int</type>)</function></literal>
+        <literal><function>pg_cancel_backend(<parameter>pid</parameter> <type>int</type> [, <parameter>message</parameter> <type>text</type>])</function></literal>
         </entry>
        <entry><type>boolean</type></entry>
        <entry>Cancel a backend's current query.  This is also allowed if the
@@ -18406,7 +18406,7 @@ SELECT set_config('log_statement_stats', 'off', false);
       </row>
       <row>
        <entry>
-        <literal><function>pg_terminate_backend(<parameter>pid</parameter> <type>int</type>)</function></literal>
+        <literal><function>pg_terminate_backend(<parameter>pid</parameter> <type>int</type> [, <parameter>message</parameter> <type>text</type>])</function></literal>
         </entry>
        <entry><type>boolean</type></entry>
        <entry>Terminate a backend.  This is also allowed if the calling role
@@ -18437,6 +18437,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 0c86a581c0..4e21d0f47b 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -44,6 +44,7 @@
 #include "storage/procsignal.h"
 #include "storage/sinvaladt.h"
 #include "storage/spin.h"
+#include "utils/backend_cancel.h"
 #include "utils/backend_random.h"
 #include "utils/snapmgr.h"
 
@@ -150,6 +151,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 +272,7 @@ CreateSharedMemoryAndSemaphores(bool makePrivate, int port)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	BackendRandomShmemInit();
+	BackendCancelShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index ddc3ec860a..165ea0f906 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"
@@ -2918,9 +2919,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)
 	{
@@ -3031,9 +3045,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 2e1e020c4b..30ac42fd99 100644
--- a/src/backend/utils/adt/misc.c
+++ b/src/backend/utils/adt/misc.c
@@ -42,6 +42,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"
 
@@ -217,7 +218,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);
 
@@ -249,6 +250,10 @@ pg_signal_backend(int pid, int sig)
 		!has_privs_of_role(GetUserId(), DEFAULT_ROLE_SIGNAL_BACKENDID))
 		return SIGNAL_BACKEND_NOPERMISSION;
 
+	/* If the user supplied a message to the signalled backend */
+	if (msg != NULL)
+		SetBackendCancelMessage(pid, msg);
+
 	/*
 	 * Can the process we just validated above end, followed by the pid being
 	 * recycled for a new process, before reaching here?  Then we'd be trying
@@ -279,10 +284,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,
@@ -297,16 +302,32 @@ 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,
@@ -318,7 +339,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 f9b330998d..026419d23d 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -51,6 +51,7 @@
 #include "storage/smgr.h"
 #include "tcop/tcopprot.h"
 #include "utils/acl.h"
+#include "utils/backend_cancel.h"
 #include "utils/fmgroids.h"
 #include "utils/guc.h"
 #include "utils/memutils.h"
@@ -741,6 +742,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..66261893cb
--- /dev/null
+++ b/src/backend/utils/misc/backend_cancel.c
@@ -0,0 +1,195 @@
+/*-------------------------------------------------------------------------
+ *
+ * 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-2018, 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 "mb/pg_wchar.h"
+#include "miscadmin.h"
+#include "storage/ipc.h"
+#include "storage/spin.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 a per-slot spinlock.
+ */
+typedef struct
+{
+	pid_t	pid;
+	slock_t	mutex;
+	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;
+	int		i;
+
+	BackendCancelSlots = (BackendCancelShmemStruct *)
+		ShmemInitStruct("BackendCancelSlots", size, &found);
+
+	if (!found)
+	{
+		MemSet(BackendCancelSlots, 0, size);
+
+		for (i = 0; i < MaxBackends; i++)
+			SpinLockInit(&(BackendCancelSlots[i].mutex));
+	}
+}
+
+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 the message actually created. If the returned length
+ * is less than the length of the message parameter, truncation has occurred.
+ * If the backend isn't found, -1 is returned. If no message is passed, zero is
+ * returned. If two backends collide in setting a message, the existing message
+ * will be overwritten by the last one in.
+ */
+int
+SetBackendCancelMessage(pid_t backend, char *message)
+{
+	BackendCancelShmemStruct *slot;
+	int		i;
+	int		len;
+
+	if (!message)
+		return 0;
+
+	for (i = 0; i < MaxBackends; i++)
+	{
+		slot = &BackendCancelSlots[i];
+
+		if (slot->pid != 0 && slot->pid == backend)
+		{
+			SpinLockAcquire(&slot->mutex);
+			if (slot->pid != backend)
+			{
+				SpinLockRelease(&slot->mutex);
+				goto error;
+			}
+
+			len = pg_mbcliplen(message, strlen(message),
+							   sizeof(slot->message) - 1);
+			memset(slot->message, '\0', sizeof(slot->message));
+			memcpy(slot->message, message, len);
+			slot->len = len;
+			SpinLockRelease(&slot->mutex);
+
+			if (len != strlen(message))
+				ereport(NOTICE,
+						(errmsg("message is too long and has been truncated")));
+			return len;
+		}
+	}
+
+error:
+
+	elog(LOG, "Cancellation message requested for missing backend %d by %d",
+		 (int) backend, MyProcPid);
+
+	return -1;
+}
+
+bool
+HasCancelMessage(void)
+{
+	volatile BackendCancelShmemStruct *slot = MyCancelSlot;
+	bool 	has_message = false;
+
+	if (slot != NULL)
+	{
+		SpinLockAcquire(&slot->mutex);
+		has_message = (slot->len > 0);
+		SpinLockRelease(&slot->mutex);
+	}
+
+	return has_message;
+}
+
+/*
+ * Return the configured cancellation message and its length. If the returned
+ * length is greater than the size of the passed buffer, truncation has been
+ * performed. The message is cleared on reading.
+ */
+int
+GetCancelMessage(char **buffer, size_t buf_len)
+{
+	volatile BackendCancelShmemStruct *slot = MyCancelSlot;
+	int		msg_length = 0;
+
+	if (slot != NULL && slot->len > 0)
+	{
+		SpinLockAcquire(&slot->mutex);
+		strlcpy(*buffer, (const char *) slot->message, buf_len);
+		msg_length = slot->len;
+		slot->len = 0;
+		slot->message[0] = '\0';
+		SpinLockRelease(&slot->mutex);
+	}
+
+	return msg_length;
+}
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index f01648c961..9f6d4bda6b 100644
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -3254,8 +3254,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 = 3438 ( 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 = 3437 ( 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..061f0fab18
--- /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-2018, 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.14.1.145.gb3622a4ee

