From eca26d5858fa0750427b15f4439d8936daff4d8b Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter_e@gmx.net>
Date: Tue, 14 Feb 2017 22:20:02 -0500
Subject: [PATCH] Add max_worker_processes_per_user setting

---
 doc/src/sgml/config.sgml          | 22 ++++++++++++++++++++++
 doc/src/sgml/ref/create_role.sgml |  3 ++-
 src/backend/postmaster/bgworker.c | 28 ++++++++++++++++++++++++++++
 src/backend/utils/init/globals.c  |  1 +
 src/backend/utils/misc/guc.c      | 12 ++++++++++++
 src/include/miscadmin.h           |  1 +
 6 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml
index dc63d7d5e4..ba74556444 100644
--- a/doc/src/sgml/config.sgml
+++ b/doc/src/sgml/config.sgml
@@ -2013,6 +2013,28 @@ <title>Asynchronous Behavior</title>
        </listitem>
       </varlistentry>
 
+      <varlistentry id="guc-max-worker-processes-per-user" xreflabel="max_worker_processes_per_user">
+       <term><varname>max_worker_processes_per_user</varname> (<type>integer</type>)
+       <indexterm>
+        <primary><varname>max_worker_processes_per_user</> configuration parameter</primary>
+       </indexterm>
+       </term>
+       <listitem>
+        <para>
+         Sets the maximum number of background processes allowed per user.  If
+         the setting is -1, then there is no limit.  That is also the default.
+        </para>
+
+        <para>
+         Only superusers can change this setting.
+         Unlike <varname>max_worker_processes</varname>, which controls the
+         overall instance limit, this setting can also be changed at run time
+         and can be set differently for different users by
+         using <literal>ALTER ROLE ... SET</literal>.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry id="guc-max-parallel-workers-per-gather" xreflabel="max_parallel_workers_per_gather">
        <term><varname>max_parallel_workers_per_gather</varname> (<type>integer</type>)
        <indexterm>
diff --git a/doc/src/sgml/ref/create_role.sgml b/doc/src/sgml/ref/create_role.sgml
index 2ae576ede6..4d0b8127f9 100644
--- a/doc/src/sgml/ref/create_role.sgml
+++ b/doc/src/sgml/ref/create_role.sgml
@@ -201,7 +201,8 @@ <title>Parameters</title>
         the role can make.  -1 (the default) means no limit. Note that only
         normal connections are counted towards this limit. Neither prepared
         transactions nor background worker connections are counted towards
-        this limit.
+        this limit (see <xref linkend="guc-max-worker-processes-per-user">
+        for the latter).
        </para>
       </listitem>
      </varlistentry>
diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c
index cd99b0b392..f1045ddcd2 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -77,6 +77,7 @@ typedef struct BackgroundWorkerSlot
 	bool		in_use;
 	bool		terminate;
 	pid_t		pid;			/* InvalidPid = not started yet; 0 = dead */
+	Oid			roleid;			/* user responsible for it */
 	uint64		generation;		/* incremented when slot is recycled */
 	BackgroundWorker worker;
 } BackgroundWorkerSlot;
@@ -905,6 +906,32 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 	}
 
 	/*
+	 * Check number of used slots for user
+	 */
+	if (max_worker_processes_per_user >= 0)
+	{
+		int		count = 0;
+
+		for (slotno = 0; slotno < BackgroundWorkerData->total_slots; ++slotno)
+		{
+			BackgroundWorkerSlot *slot = &BackgroundWorkerData->slot[slotno];
+
+			if (slot->in_use && slot->roleid == GetUserId())
+				count++;
+		}
+
+		if (count > max_worker_processes_per_user)
+		{
+			ereport(LOG,
+					(errcode(ERRCODE_CONFIGURATION_LIMIT_EXCEEDED),
+					 errmsg("too many worker processes for role \"%s\"",
+							GetUserNameFromId(GetUserId(), false))));
+			LWLockRelease(BackgroundWorkerLock);
+			return false;
+		}
+	}
+
+	/*
 	 * Look for an unused slot.  If we find one, grab it.
 	 */
 	for (slotno = 0; slotno < BackgroundWorkerData->total_slots; ++slotno)
@@ -915,6 +942,7 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 		{
 			memcpy(&slot->worker, worker, sizeof(BackgroundWorker));
 			slot->pid = InvalidPid;		/* indicates not started yet */
+			slot->roleid = GetUserId();
 			slot->generation++;
 			slot->terminate = false;
 			generation = slot->generation;
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index 08b6030a64..f92c8c196f 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -122,6 +122,7 @@ int			replacement_sort_tuples = 150000;
 int			NBuffers = 1000;
 int			MaxConnections = 90;
 int			max_worker_processes = 8;
+int			max_worker_processes_per_user = -1;
 int			max_parallel_workers = 8;
 int			MaxBackends = 0;
 
diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
index 0249721204..bbae5dc9f8 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -2480,6 +2480,18 @@ static struct config_int ConfigureNamesInt[] =
 	},
 
 	{
+		{"max_worker_processes_per_user",
+			PGC_SUSET,
+			RESOURCES_ASYNCHRONOUS,
+			gettext_noop("Maximum number of concurrent worker processes per user."),
+			NULL,
+		},
+		&max_worker_processes_per_user,
+		-1, -1, MAX_BACKENDS,
+		NULL, NULL, NULL
+	},
+
+	{
 		{"max_logical_replication_workers",
 			PGC_POSTMASTER,
 			RESOURCES_ASYNCHRONOUS,
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 4c607b299c..1c96f8876c 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -157,6 +157,7 @@ extern PGDLLIMPORT int NBuffers;
 extern int	MaxBackends;
 extern int	MaxConnections;
 extern int	max_worker_processes;
+extern int	max_worker_processes_per_user;
 extern int	max_parallel_workers;
 
 extern PGDLLIMPORT int MyProcPid;
-- 
2.11.1

