From 0256847b0645743ce9670c05834ea8fb62a3fc5d Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Mon, 17 Feb 2025 11:50:36 -0800
Subject: [PATCH v3 1/3] Introduce reserved background worker slots.

---
 src/backend/bootstrap/bootstrap.c           |   2 +
 src/backend/postmaster/bgworker.c           | 125 +++++++++++++++-----
 src/backend/postmaster/pmchild.c            |   3 +-
 src/backend/postmaster/postmaster.c         |   6 +
 src/backend/storage/lmgr/proc.c             |   2 +-
 src/backend/tcop/postgres.c                 |   4 +
 src/backend/utils/init/globals.c            |   1 +
 src/backend/utils/init/postinit.c           |   3 +-
 src/include/miscadmin.h                     |   1 +
 src/include/postmaster/bgworker.h           |   7 +-
 src/include/postmaster/bgworker_internals.h |  10 ++
 11 files changed, 129 insertions(+), 35 deletions(-)

diff --git a/src/backend/bootstrap/bootstrap.c b/src/backend/bootstrap/bootstrap.c
index 6db864892d0..4d886b22a97 100644
--- a/src/backend/bootstrap/bootstrap.c
+++ b/src/backend/bootstrap/bootstrap.c
@@ -31,6 +31,7 @@
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
 #include "pg_getopt.h"
+#include "postmaster/bgworker_internals.h"
 #include "postmaster/postmaster.h"
 #include "storage/bufpage.h"
 #include "storage/ipc.h"
@@ -322,6 +323,7 @@ BootstrapModeMain(int argc, char *argv[], bool check_only)
 	SetProcessingMode(BootstrapProcessing);
 	IgnoreSystemIndexes = true;
 
+	InitializeMaxBgWorkers();
 	InitializeMaxBackends();
 
 	/*
diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c
index b288915cec8..0f4389ad595 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -81,21 +81,36 @@ typedef struct BackgroundWorkerSlot
 } BackgroundWorkerSlot;
 
 /*
- * In order to limit the total number of parallel workers (according to
- * max_parallel_workers GUC), we maintain the number of active parallel
- * workers.  Since the postmaster cannot take locks, two variables are used for
- * this purpose: the number of registered parallel workers (modified by the
- * backends, protected by BackgroundWorkerLock) and the number of terminated
- * parallel workers (modified only by the postmaster, lockless).  The active
- * number of parallel workers is the number of registered workers minus the
- * terminated ones.  These counters can of course overflow, but it's not
- * important here since the subtraction will still give the right number.
+ * Struct holding information about background workers.
  */
 typedef struct BackgroundWorkerArray
 {
 	int			total_slots;
+
+	/*
+	 * In order to limit the total number of parallel workers (according to
+	 * max_parallel_workers GUC), we maintain the number of active parallel
+	 * workers.  Since the postmaster cannot take locks, two variables are
+	 * used for this purpose: the number of registered parallel workers
+	 * (modified by the backends, protected by BackgroundWorkerLock) and the
+	 * number of terminated parallel workers (modified only by the postmaster,
+	 * lockless).  The active number of parallel workers is the number of
+	 * registered workers minus the terminated ones.  These counters can of
+	 * course overflow, but it's not important here since the subtraction will
+	 * still give the right number.
+	 */
 	uint32		parallel_register_count;
 	uint32		parallel_terminate_count;
+
+	/*
+	 * Similar to parallel_register_count and parallel_terminate_count, but
+	 * these numbers tracks the total number of registration and termination.
+	 * The total active number of background workers is (register_count -
+	 * terminate_count).
+	 */
+	uint32		register_count;
+	uint32		terminate_count;
+
 	BackgroundWorkerSlot slot[FLEXIBLE_ARRAY_MEMBER];
 } BackgroundWorkerArray;
 
@@ -149,7 +164,7 @@ BackgroundWorkerShmemSize(void)
 
 	/* Array of workers is variably sized. */
 	size = offsetof(BackgroundWorkerArray, slot);
-	size = add_size(size, mul_size(max_worker_processes,
+	size = add_size(size, mul_size(MaxBgWorkers,
 								   sizeof(BackgroundWorkerSlot)));
 
 	return size;
@@ -171,7 +186,9 @@ BackgroundWorkerShmemInit(void)
 		dlist_iter	iter;
 		int			slotno = 0;
 
-		BackgroundWorkerData->total_slots = max_worker_processes;
+		BackgroundWorkerData->total_slots = MaxBgWorkers;
+		BackgroundWorkerData->register_count = 0;
+		BackgroundWorkerData->terminate_count = 0;
 		BackgroundWorkerData->parallel_register_count = 0;
 		BackgroundWorkerData->parallel_terminate_count = 0;
 
@@ -187,7 +204,7 @@ BackgroundWorkerShmemInit(void)
 			RegisteredBgWorker *rw;
 
 			rw = dlist_container(RegisteredBgWorker, rw_lnode, iter.cur);
-			Assert(slotno < max_worker_processes);
+			Assert(slotno < MaxBgWorkers);
 			slot->in_use = true;
 			slot->terminate = false;
 			slot->pid = InvalidPid;
@@ -196,12 +213,13 @@ BackgroundWorkerShmemInit(void)
 			rw->rw_worker.bgw_notify_pid = 0;	/* might be reinit after crash */
 			memcpy(&slot->worker, &rw->rw_worker, sizeof(BackgroundWorker));
 			++slotno;
+			++BackgroundWorkerData->register_count;
 		}
 
 		/*
 		 * Mark any remaining slots as not in use.
 		 */
-		while (slotno < max_worker_processes)
+		while (slotno < MaxBgWorkers)
 		{
 			BackgroundWorkerSlot *slot = &BackgroundWorkerData->slot[slotno];
 
@@ -213,6 +231,15 @@ BackgroundWorkerShmemInit(void)
 		Assert(found);
 }
 
+/*
+ * Initialize MaxBgWorkers value form config options.
+ */
+void
+InitializeMaxBgWorkers(void)
+{
+	MaxBgWorkers = max_worker_processes + BGWORKER_NUM_RESERVED_WORKERS;
+}
+
 /*
  * Search the postmaster's backend-private list of RegisteredBgWorker objects
  * for the one that maps to the given slot number.
@@ -249,16 +276,16 @@ BackgroundWorkerStateChange(bool allow_new_workers)
 
 	/*
 	 * The total number of slots stored in shared memory should match our
-	 * notion of max_worker_processes.  If it does not, something is very
-	 * wrong.  Further down, we always refer to this value as
-	 * max_worker_processes, in case shared memory gets corrupted while we're
-	 * looping.
+	 * notion of MaxBgWorkers.  If it does not, something is very wrong.
+	 * Further down, we always refer to this value as max_worker_processes, in
+	 * case shared memory gets corrupted while we're looping.
 	 */
-	if (max_worker_processes != BackgroundWorkerData->total_slots)
+	if (MaxBgWorkers != BackgroundWorkerData->total_slots)
 	{
 		ereport(LOG,
-				(errmsg("inconsistent background worker state (\"max_worker_processes\"=%d, total slots=%d)",
+				(errmsg("inconsistent background worker state (\"max_worker_processes\"=%d, reserved slot=%d, total slots=%d)",
 						max_worker_processes,
+						BGWORKER_NUM_RESERVED_WORKERS,
 						BackgroundWorkerData->total_slots)));
 		return;
 	}
@@ -267,7 +294,7 @@ BackgroundWorkerStateChange(bool allow_new_workers)
 	 * Iterate through slots, looking for newly-registered workers or workers
 	 * who must die.
 	 */
-	for (slotno = 0; slotno < max_worker_processes; ++slotno)
+	for (slotno = 0; slotno < MaxBgWorkers; ++slotno)
 	{
 		BackgroundWorkerSlot *slot = &BackgroundWorkerData->slot[slotno];
 		RegisteredBgWorker *rw;
@@ -325,10 +352,11 @@ BackgroundWorkerStateChange(bool allow_new_workers)
 
 			/*
 			 * We need a memory barrier here to make sure that the load of
-			 * bgw_notify_pid and the update of parallel_terminate_count
-			 * complete before the store to in_use.
+			 * bgw_notify_pid and the update of terminate_count and
+			 * parallel_terminate_count complete before the store to in_use.
 			 */
 			notify_pid = slot->worker.bgw_notify_pid;
+			BackgroundWorkerData->terminate_count++;
 			if ((slot->worker.bgw_flags & BGWORKER_CLASS_PARALLEL) != 0)
 				BackgroundWorkerData->parallel_terminate_count++;
 			slot->pid = 0;
@@ -430,14 +458,16 @@ ForgetBackgroundWorker(RegisteredBgWorker *rw)
 {
 	BackgroundWorkerSlot *slot;
 
-	Assert(rw->rw_shmem_slot < max_worker_processes);
+	Assert(rw->rw_shmem_slot < MaxBgWorkers);
 	slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
 	Assert(slot->in_use);
 
 	/*
 	 * We need a memory barrier here to make sure that the update of
-	 * parallel_terminate_count completes before the store to in_use.
+	 * terminate_count and parallel_terminate_count completes before the store
+	 * to in_use.
 	 */
+	BackgroundWorkerData->terminate_count++;
 	if ((rw->rw_worker.bgw_flags & BGWORKER_CLASS_PARALLEL) != 0)
 		BackgroundWorkerData->parallel_terminate_count++;
 
@@ -462,7 +492,7 @@ ReportBackgroundWorkerPID(RegisteredBgWorker *rw)
 {
 	BackgroundWorkerSlot *slot;
 
-	Assert(rw->rw_shmem_slot < max_worker_processes);
+	Assert(rw->rw_shmem_slot < MaxBgWorkers);
 	slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
 	slot->pid = rw->rw_pid;
 
@@ -485,7 +515,7 @@ ReportBackgroundWorkerExit(RegisteredBgWorker *rw)
 	BackgroundWorkerSlot *slot;
 	int			notify_pid;
 
-	Assert(rw->rw_shmem_slot < max_worker_processes);
+	Assert(rw->rw_shmem_slot < MaxBgWorkers);
 	slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
 	slot->pid = rw->rw_pid;
 	notify_pid = rw->rw_worker.bgw_notify_pid;
@@ -548,7 +578,7 @@ ForgetUnstartedBackgroundWorkers(void)
 		BackgroundWorkerSlot *slot;
 
 		rw = dlist_container(RegisteredBgWorker, rw_lnode, iter.cur);
-		Assert(rw->rw_shmem_slot < max_worker_processes);
+		Assert(rw->rw_shmem_slot < MaxBgWorkers);
 		slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
 
 		/* If it's not yet started, and there's someone waiting ... */
@@ -939,6 +969,7 @@ void
 RegisterBackgroundWorker(BackgroundWorker *worker)
 {
 	RegisteredBgWorker *rw;
+	bool		use_reserved_slot;
 	static int	numworkers = 0;
 
 	/*
@@ -990,17 +1021,34 @@ RegisterBackgroundWorker(BackgroundWorker *worker)
 		return;
 	}
 
+	use_reserved_slot = (worker->bgw_flags & BGWORKER_CLASS_RESERVED) != 0;
+
 	/*
 	 * Enforce maximum number of workers.  Note this is overly restrictive: we
 	 * could allow more non-shmem-connected workers, because these don't count
 	 * towards the MAX_BACKENDS limit elsewhere.  For now, it doesn't seem
 	 * important to relax this restriction.
 	 */
-	if (++numworkers > max_worker_processes)
+	numworkers++;
+	if (!use_reserved_slot &&
+		(MaxBgWorkers - numworkers) <= BGWORKER_NUM_RESERVED_WORKERS)
+	{
+		ereport(LOG,
+				(errcode(ERRCODE_CONFIGURATION_LIMIT_EXCEEDED),
+				 errmsg("too many background workers max %d nworkers %d n_resreved %d",
+						MaxBgWorkers, numworkers, BGWORKER_NUM_RESERVED_WORKERS),
+				 errdetail_plural("Up to %d background worker can be registered with the current settings.",
+								  "Up to %d background workers can be registered with the current settings.",
+								  max_worker_processes,
+								  max_worker_processes),
+				 errhint("Consider increasing the configuration parameter \"%s\".", "max_worker_processes")));
+		return;
+	}
+	else if (numworkers > MaxBgWorkers)
 	{
 		ereport(LOG,
 				(errcode(ERRCODE_CONFIGURATION_LIMIT_EXCEEDED),
-				 errmsg("too many background workers"),
+				 errmsg("remaining background workers are reserved"),
 				 errdetail_plural("Up to %d background worker can be registered with the current settings.",
 								  "Up to %d background workers can be registered with the current settings.",
 								  max_worker_processes,
@@ -1048,6 +1096,7 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 	int			slotno;
 	bool		success = false;
 	bool		parallel;
+	bool		use_reserved_slot;
 	uint64		generation = 0;
 
 	/*
@@ -1065,6 +1114,7 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 		return false;
 
 	parallel = (worker->bgw_flags & BGWORKER_CLASS_PARALLEL) != 0;
+	use_reserved_slot = (worker->bgw_flags & BGWORKER_CLASS_RESERVED) != 0;
 
 	LWLockAcquire(BackgroundWorkerLock, LW_EXCLUSIVE);
 
@@ -1088,6 +1138,18 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 		return false;
 	}
 
+	if (!use_reserved_slot &&
+		(MaxBgWorkers -
+		 (BackgroundWorkerData->register_count - BackgroundWorkerData->terminate_count)) <=
+		BGWORKER_NUM_RESERVED_WORKERS)
+	{
+		Assert(BackgroundWorkerData->register_count -
+			   BackgroundWorkerData->terminate_count <=
+			   MAX_PARALLEL_WORKER_LIMIT);
+		LWLockRelease(BackgroundWorkerLock);
+		return false;
+	}
+
 	/*
 	 * Look for an unused slot.  If we find one, grab it.
 	 */
@@ -1102,6 +1164,7 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 			slot->generation++;
 			slot->terminate = false;
 			generation = slot->generation;
+			BackgroundWorkerData->register_count++;
 			if (parallel)
 				BackgroundWorkerData->parallel_register_count++;
 
@@ -1159,7 +1222,7 @@ GetBackgroundWorkerPid(BackgroundWorkerHandle *handle, pid_t *pidp)
 	BackgroundWorkerSlot *slot;
 	pid_t		pid;
 
-	Assert(handle->slot < max_worker_processes);
+	Assert(handle->slot < MaxBgWorkers);
 	slot = &BackgroundWorkerData->slot[handle->slot];
 
 	/*
@@ -1298,7 +1361,7 @@ TerminateBackgroundWorker(BackgroundWorkerHandle *handle)
 	BackgroundWorkerSlot *slot;
 	bool		signal_postmaster = false;
 
-	Assert(handle->slot < max_worker_processes);
+	Assert(handle->slot < MaxBgWorkers);
 	slot = &BackgroundWorkerData->slot[handle->slot];
 
 	/* Set terminate flag in shared memory, unless slot has been reused. */
diff --git a/src/backend/postmaster/pmchild.c b/src/backend/postmaster/pmchild.c
index 0d473226c3a..515a9fda3af 100644
--- a/src/backend/postmaster/pmchild.c
+++ b/src/backend/postmaster/pmchild.c
@@ -33,6 +33,7 @@
 
 #include "miscadmin.h"
 #include "postmaster/autovacuum.h"
+#include "postmaster/bgworker_internals.h"
 #include "postmaster/postmaster.h"
 #include "replication/walsender.h"
 #include "storage/pmsignal.h"
@@ -100,7 +101,7 @@ InitPostmasterChildSlots(void)
 	pmchild_pools[B_BACKEND].size = 2 * (MaxConnections + max_wal_senders);
 
 	pmchild_pools[B_AUTOVAC_WORKER].size = autovacuum_worker_slots;
-	pmchild_pools[B_BG_WORKER].size = max_worker_processes;
+	pmchild_pools[B_BG_WORKER].size = max_worker_processes + BGWORKER_NUM_RESERVED_WORKERS;
 
 	/*
 	 * There can be only one of each of these running at a time.  They each
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index bb22b13adef..108a1e5cfbe 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -910,6 +910,12 @@ PostmasterMain(int argc, char *argv[])
 	 */
 	LocalProcessControlFile(false);
 
+	/*
+	 * Initialize MaxBgWorker. This should be called after initializing GUC
+	 * before any chances to register background workers.
+	 */
+	InitializeMaxBgWorkers();
+
 	/*
 	 * Register the apply launcher.  It's probably a good idea to call this
 	 * before any modules had a chance to take the background worker slots.
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index 49204f91a20..33a5b062198 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -290,7 +290,7 @@ InitProcGlobal(void)
 			dlist_push_tail(&ProcGlobal->autovacFreeProcs, &proc->links);
 			proc->procgloballist = &ProcGlobal->autovacFreeProcs;
 		}
-		else if (i < MaxConnections + autovacuum_worker_slots + NUM_SPECIAL_WORKER_PROCS + max_worker_processes)
+		else if (i < MaxConnections + autovacuum_worker_slots + NUM_SPECIAL_WORKER_PROCS + MaxBgWorkers)
 		{
 			/* PGPROC for bgworker, add to bgworkerFreeProcs list */
 			dlist_push_tail(&ProcGlobal->bgworkerFreeProcs, &proc->links);
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 1149d89d7a1..7ae12517e8b 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -53,6 +53,7 @@
 #include "pg_getopt.h"
 #include "pg_trace.h"
 #include "pgstat.h"
+#include "postmaster/bgworker_internals.h"
 #include "postmaster/interrupt.h"
 #include "postmaster/postmaster.h"
 #include "replication/logicallauncher.h"
@@ -4064,6 +4065,9 @@ PostgresSingleUserMain(int argc, char *argv[],
 	/* read control file (error checking and contains config ) */
 	LocalProcessControlFile(false);
 
+	/* Initialize MaxBgWorkers */
+	InitializeMaxBgWorkers();
+
 	/*
 	 * process any libraries that should be preloaded at postmaster start
 	 */
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index b844f9fdaef..8bccc13684d 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -143,6 +143,7 @@ int			MaxConnections = 100;
 int			max_worker_processes = 8;
 int			max_parallel_workers = 8;
 int			MaxBackends = 0;
+int			MaxBgWorkers = 0;
 
 /* GUC parameters for vacuum */
 int			VacuumBufferUsageLimit = 2048;
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index 01bb6a410cb..d70a5be446a 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -545,10 +545,11 @@ void
 InitializeMaxBackends(void)
 {
 	Assert(MaxBackends == 0);
+	Assert(MaxBgWorkers > 0);
 
 	/* Note that this does not include "auxiliary" processes */
 	MaxBackends = MaxConnections + autovacuum_worker_slots +
-		max_worker_processes + max_wal_senders + NUM_SPECIAL_WORKER_PROCS;
+		MaxBgWorkers + max_wal_senders + NUM_SPECIAL_WORKER_PROCS;
 
 	if (MaxBackends > MAX_BACKENDS)
 		ereport(ERROR,
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index a2b63495eec..2885357cf4d 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -174,6 +174,7 @@ extern PGDLLIMPORT int data_directory_mode;
 
 extern PGDLLIMPORT int NBuffers;
 extern PGDLLIMPORT int MaxBackends;
+extern PGDLLIMPORT int MaxBgWorkers;
 extern PGDLLIMPORT int MaxConnections;
 extern PGDLLIMPORT int max_worker_processes;
 extern PGDLLIMPORT int max_parallel_workers;
diff --git a/src/include/postmaster/bgworker.h b/src/include/postmaster/bgworker.h
index 058667a47a0..9ff12ba1e57 100644
--- a/src/include/postmaster/bgworker.h
+++ b/src/include/postmaster/bgworker.h
@@ -66,8 +66,13 @@
  * background workers should not use this class.
  */
 #define BGWORKER_CLASS_PARALLEL					0x0010
-/* add additional bgworker classes here */
 
+/*
+ * This class is used for the worker to use a reserved background worker
+ * slots from the pool of BGWORKER_NUM_RESERVED_WORKERS workers.
+ */
+#define BGWORKER_CLASS_RESERVED					0x0020
+/* add additional bgworker classes here */
 
 typedef void (*bgworker_main_type) (Datum main_arg);
 
diff --git a/src/include/postmaster/bgworker_internals.h b/src/include/postmaster/bgworker_internals.h
index 092b1610663..c26ec6b2387 100644
--- a/src/include/postmaster/bgworker_internals.h
+++ b/src/include/postmaster/bgworker_internals.h
@@ -23,6 +23,15 @@
  */
 #define MAX_PARALLEL_WORKER_LIMIT 1024
 
+/*
+ * The number of background workers reserved for ones registered with
+ * BGWORKER_CLASS_RESERVED flag. We allocate total background worker
+ * slots for max_worker_process plus this number.
+ *
+ * XXX: we don't have any reserved slots for now.
+ */
+#define BGWORKER_NUM_RESERVED_WORKERS			0
+
 /*
  * List of background workers, private to postmaster.
  *
@@ -43,6 +52,7 @@ extern PGDLLIMPORT dlist_head BackgroundWorkerList;
 
 extern Size BackgroundWorkerShmemSize(void);
 extern void BackgroundWorkerShmemInit(void);
+extern void InitializeMaxBgWorkers(void);
 extern void BackgroundWorkerStateChange(bool allow_new_workers);
 extern void ForgetBackgroundWorker(RegisteredBgWorker *rw);
 extern void ReportBackgroundWorkerPID(RegisteredBgWorker *rw);
-- 
2.43.5

