From e99d8ccf9dfd007c054b20e8d10ffb35cfb722b7 Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter_e@gmx.net>
Date: Wed, 23 Aug 2017 19:10:21 -0400
Subject: [PATCH v2] Add max_worker_processes_per_user setting

---
 doc/src/sgml/config.sgml          | 28 ++++++++++++++++++++++++++++
 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, 72 insertions(+), 1 deletion(-)

diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml
index 2b6255ed95..24f1d13f8a 100644
--- a/doc/src/sgml/config.sgml
+++ b/doc/src/sgml/config.sgml
@@ -2041,6 +2041,34 @@ <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>
+
+        <para>
+         This setting is checked at the time the worker process is started
+         using the setting in the session that causes it to start.  Changes to
+         the setting will not affect already running workers.
+        </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 36772b678a..a7c547d907 100644
--- a/doc/src/sgml/ref/create_role.sgml
+++ b/doc/src/sgml/ref/create_role.sgml
@@ -204,7 +204,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 28af6f0f07..443dcaa4f3 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -79,6 +79,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;
@@ -976,6 +977,32 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 		return false;
 	}
 
+	/*
+	 * 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.
 	 */
@@ -987,6 +1014,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 7c09498dc0..fa1b689ae7 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -123,6 +123,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 246fea8693..23dc33bd61 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -2506,6 +2506,18 @@ static struct config_int ConfigureNamesInt[] =
 		check_max_worker_processes, NULL, NULL
 	},
 
+	{
+		{"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,
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index dad98de98d..a9f21cc34b 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -158,6 +158,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.14.1