From 76aee1354bae13a29c7ffb53be6736a66fb3024c Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Mon, 17 Feb 2025 11:52:09 -0800
Subject: [PATCH v3 2/3] Introduce bgtask, infrastructure to perform tasks in
 background.

---
 src/backend/postmaster/bgworker.c             |   4 +
 src/backend/storage/ipc/ipci.c                |   2 +
 .../utils/activity/wait_event_names.txt       |   1 +
 src/backend/utils/misc/Makefile               |   1 +
 src/backend/utils/misc/bgtask.c               | 241 ++++++++++++++++++
 src/backend/utils/misc/meson.build            |   1 +
 src/include/postmaster/bgworker_internals.h   |   8 +-
 src/include/storage/lwlocklist.h              |   1 +
 src/include/utils/bgtask.h                    |  38 +++
 9 files changed, 293 insertions(+), 4 deletions(-)
 create mode 100644 src/backend/utils/misc/bgtask.c
 create mode 100644 src/include/utils/bgtask.h

diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c
index 0f4389ad595..dc4af5ac904 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -29,6 +29,7 @@
 #include "storage/procsignal.h"
 #include "storage/shmem.h"
 #include "tcop/tcopprot.h"
+#include "utils/bgtask.h"
 #include "utils/ascii.h"
 #include "utils/memutils.h"
 #include "utils/ps_status.h"
@@ -147,6 +148,9 @@ static const struct
 	},
 	{
 		"TablesyncWorkerMain", TablesyncWorkerMain
+	},
+	{
+		"BgTaskWorkerMain", BgTaskWorkerMain
 	}
 };
 
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 174eed70367..7c3c74dcdaf 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -148,6 +148,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, WaitEventCustomShmemSize());
 	size = add_size(size, InjectionPointShmemSize());
 	size = add_size(size, SlotSyncShmemSize());
+	size = add_size(size, BgTaskShmemSize());
 
 	/* include additional requested shmem from preload libraries */
 	size = add_size(size, total_addin_request);
@@ -340,6 +341,7 @@ CreateOrAttachShmemStructs(void)
 	StatsShmemInit();
 	WaitEventCustomShmemInit();
 	InjectionPointShmemInit();
+	BgTaskShmemInit();
 }
 
 /*
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index e199f071628..fef6887eb43 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -346,6 +346,7 @@ WALSummarizer	"Waiting to read or update WAL summarization state."
 DSMRegistry	"Waiting to read or update the dynamic shared memory registry."
 InjectionPoint	"Waiting to read or update information related to injection points."
 SerialControl	"Waiting to read or update shared <filename>pg_serial</filename> state."
+BgTaskControl	"Waiting to read or update shared background task information."
 
 #
 # END OF PREDEFINED LWLOCKS (DO NOT CHANGE THIS LINE)
diff --git a/src/backend/utils/misc/Makefile b/src/backend/utils/misc/Makefile
index b362ae43771..b68f263a396 100644
--- a/src/backend/utils/misc/Makefile
+++ b/src/backend/utils/misc/Makefile
@@ -15,6 +15,7 @@ include $(top_builddir)/src/Makefile.global
 override CPPFLAGS := -I. -I$(srcdir) $(CPPFLAGS)
 
 OBJS = \
+	bgtask.o \
 	conffiles.o \
 	guc.o \
 	guc-file.o \
diff --git a/src/backend/utils/misc/bgtask.c b/src/backend/utils/misc/bgtask.c
new file mode 100644
index 00000000000..ac67968b9a1
--- /dev/null
+++ b/src/backend/utils/misc/bgtask.c
@@ -0,0 +1,241 @@
+/*--------------------------------------------------------------------
+ * bgtask.c
+ *		Infrastructure for executing internal task in background.
+ *
+ * The bgtask infrastructure can be used to offload internal tasks that are
+ * potentially costly to a background worker called "bgtask worker". The
+ * internal tasks that can be requested to bgtask workers are pre-defined
+ * (see BgTaskType and InternalBgTasks[]). Since the bgtask workers use
+ * reserved background worker slots it's guaranteed that they can be
+ * launched successfully.
+ *
+ * XXX: This files contains the minimal functionality to just request
+ * tasks to bgtask workers. Bgtask workers are never restarted on failure.
+ * These behavior may need to be changed.
+ *
+ * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/backend/utils/misc/bgtask.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "miscadmin.h"
+#include "storage/ipc.h"
+#include "storage/lwlock.h"
+#include "storage/shmem.h"
+#include "tcop/tcopprot.h"
+#include "utils/bgtask.h"
+
+/* Struct for single background task status */
+typedef struct BgTaskSlot
+{
+	Oid			dboid;
+	pid_t		worker_pid;
+
+	/*
+	 * True if this task is requested. This flag is set by the requester, and
+	 * cleared by the worker assigned to the task.
+	 */
+	bool		in_progress;
+}			BgTaskSlot;
+
+/* Control struct for internal background tasks */
+typedef struct BgTaskCtlData
+{
+	/* NUM_INTERNAL_BGTASKS */
+	int			total_tasks;
+
+	/*
+	 * Array of internal task status. The index of this array corresponds to
+	 * BgTaskType number.
+	 */
+	BgTaskSlot	task_slots[FLEXIBLE_ARRAY_MEMBER];
+}			BgTaskCtlData;
+static BgTaskCtlData * BgTaskCtl = NULL;
+
+static void dummy_bgtask_fn(void);
+
+static const struct
+{
+	const char *name;
+	bgtask_fn_type func;
+}			InternalBgTasks[] =
+
+{
+	[BGTASK_TYPE_DUMMY] = {
+		.name = "dummy task",
+		.func = &dummy_bgtask_fn,
+	},
+};
+
+static void BgTaskWorkerDetach(int code, Datum arg);
+
+Size
+BgTaskShmemSize(void)
+{
+	Size		size;
+
+	size = offsetof(BgTaskCtlData, task_slots);
+	size = add_size(size, mul_size(NUM_INTERNAL_BGTASKS,
+								   sizeof(BgTaskSlot)));
+
+	return size;
+}
+
+void
+BgTaskShmemInit(void)
+{
+	bool		found;
+
+	BgTaskCtl = ShmemInitStruct("Background Task Control Data",
+								BgTaskShmemSize(),
+								&found);
+
+	if (!found)
+	{
+		BgTaskCtl->total_tasks = NUM_INTERNAL_BGTASKS;
+		MemSet(BgTaskCtl->task_slots, 0,
+			   sizeof(BgTaskSlot) * BgTaskCtl->total_tasks);
+	}
+}
+
+/*
+ * Request the work on the given database to bgtask worker.
+ *
+ * Launching a bgtask worker must be done successfully because we allocates enough
+ * reserved background workers for bgtask workers.
+ */
+void
+BgTaskRequest(BgTaskType type, Oid dboid)
+{
+	BgTaskSlot *task;
+	BackgroundWorker bgw;
+	BackgroundWorkerHandle *bgw_handle;
+	pid_t		pid;
+
+	if (type > NUM_INTERNAL_BGTASKS)
+		elog(ERROR, "invalid task type %d", type);
+
+	LWLockAcquire(BgTaskControlLock, LW_EXCLUSIVE);
+
+	task = &(BgTaskCtl->task_slots[type]);
+
+	/* Return if the specified task is already in-progress */
+	if (task->in_progress)
+	{
+		LWLockRelease(BgTaskControlLock);
+		return;
+	}
+
+	/* Mark this task in-progress */
+	task->in_progress = true;
+	task->dboid = dboid;
+	LWLockRelease(BgTaskControlLock);
+
+	ereport(DEBUG1,
+			(errmsg("starting bgtask worker for task %d on database oid %u",
+					type, dboid)));
+
+	/* Register the new dynamic worker */
+	MemSet(&bgw, 0, sizeof(bgw));
+	bgw.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION |
+		BGWORKER_CLASS_RESERVED;
+	bgw.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	snprintf(bgw.bgw_library_name, MAXPGPATH, "postgres");
+	snprintf(bgw.bgw_function_name, BGW_MAXLEN, "BgTaskWorkerMain");
+	snprintf(bgw.bgw_type, BGW_MAXLEN, "background task worker");
+	bgw.bgw_restart_time = BGW_NEVER_RESTART;	/* XXX need to be configurable
+												 * by the requester? */
+	bgw.bgw_notify_pid = MyProcPid;
+	bgw.bgw_main_arg = Int32GetDatum(type);
+
+	if (RegisterDynamicBackgroundWorker(&bgw, &bgw_handle) &&
+		WaitForBackgroundWorkerStartup(bgw_handle, &pid))
+	{
+		/*
+		 * The dynamic worker registration must be done successfully since
+		 * bgtask workers use the reserved bgworker slots.
+		 */
+		elog(ERROR, "could not register a background task worker for %d", type);
+	}
+}
+
+void
+BgTaskWorkerMain(Datum main_arg)
+{
+	BgTaskType	type = DatumGetInt32(main_arg);
+	BgTaskSlot *task;
+	Oid			dboid;
+
+	Assert(type < NUM_INTERNAL_BGTASKS);
+
+	on_shmem_exit(BgTaskWorkerDetach, (Datum) main_arg);
+
+	LWLockAcquire(BgTaskControlLock, LW_EXCLUSIVE);
+	task = &(BgTaskCtl->task_slots[type]);
+
+	/* Attach to the task slot */
+	task->worker_pid = MyProcPid;
+	dboid = task->dboid;
+
+	LWLockRelease(BgTaskControlLock);
+
+	if (OidIsValid(dboid))
+		ereport(LOG,
+				(errmsg("bgtask worker for task \"%s\" on database %u has started ",
+						InternalBgTasks[type].name, dboid)));
+	else
+		ereport(LOG,
+				(errmsg("bgtask worker for task \"%s\" has started ",
+						InternalBgTasks[type].name)));
+
+	BackgroundWorkerInitializeConnectionByOid(dboid, InvalidOid, 0);
+
+	pqsignal(SIGHUP, SignalHandlerForConfigReload);
+	pqsignal(SIGTERM, die);
+	BackgroundWorkerUnblockSignals();
+
+	PG_TRY();
+	{
+		/* Execute the task */
+		(void) InternalBgTasks[type].func();
+	}
+	PG_FINALLY();
+	{
+		/* Ensure to mark this task as completed */
+		LWLockAcquire(BgTaskControlLock, LW_EXCLUSIVE);
+		task->in_progress = false;
+		LWLockRelease(BgTaskControlLock);
+	}
+	PG_END_TRY();
+
+	proc_exit(0);
+}
+
+/*
+ * Callback function for on_shmem_exit().
+ */
+static void
+BgTaskWorkerDetach(int code, Datum arg)
+{
+	BgTaskType	type = DatumGetInt32(arg);
+
+	LWLockAcquire(BgTaskControlLock, LW_EXCLUSIVE);
+	BgTaskCtl->task_slots[type].worker_pid = InvalidPid;
+	BgTaskCtl->task_slots[type].in_progress = false;
+	LWLockRelease(BgTaskControlLock);
+}
+
+static void
+dummy_bgtask_fn(void)
+{
+	elog(LOG, "executed dummy background task");
+}
diff --git a/src/backend/utils/misc/meson.build b/src/backend/utils/misc/meson.build
index 9e389a00d05..950f7cd9383 100644
--- a/src/backend/utils/misc/meson.build
+++ b/src/backend/utils/misc/meson.build
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2025, PostgreSQL Global Development Group
 
 backend_sources += files(
+  'bgtask.c',
   'conffiles.c',
   'guc.c',
   'guc_funcs.c',
diff --git a/src/include/postmaster/bgworker_internals.h b/src/include/postmaster/bgworker_internals.h
index c26ec6b2387..cb6e0e3e75e 100644
--- a/src/include/postmaster/bgworker_internals.h
+++ b/src/include/postmaster/bgworker_internals.h
@@ -15,6 +15,7 @@
 #include "datatype/timestamp.h"
 #include "lib/ilist.h"
 #include "postmaster/bgworker.h"
+#include "utils/bgtask.h"
 
 /* GUC options */
 
@@ -26,11 +27,10 @@
 /*
  * 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.
+ * slots for max_worker_process plus this number. Currently, we allocate
+ * reserved slots only for internal background task purposes.
  */
-#define BGWORKER_NUM_RESERVED_WORKERS			0
+#define BGWORKER_NUM_RESERVED_WORKERS			NUM_INTERNAL_BGTASKS
 
 /*
  * List of background workers, private to postmaster.
diff --git a/src/include/storage/lwlocklist.h b/src/include/storage/lwlocklist.h
index cf565452382..fef9f555310 100644
--- a/src/include/storage/lwlocklist.h
+++ b/src/include/storage/lwlocklist.h
@@ -83,3 +83,4 @@ PG_LWLOCK(49, WALSummarizer)
 PG_LWLOCK(50, DSMRegistry)
 PG_LWLOCK(51, InjectionPoint)
 PG_LWLOCK(52, SerialControl)
+PG_LWLOCK(53, BgTaskControl)
diff --git a/src/include/utils/bgtask.h b/src/include/utils/bgtask.h
new file mode 100644
index 00000000000..d9d89811c98
--- /dev/null
+++ b/src/include/utils/bgtask.h
@@ -0,0 +1,38 @@
+/*-------------------------------------------------------------------------
+ *
+ * bgtask.h
+ *	  Definitions of infrastructure for executing internal task in background.
+ *
+ *
+ * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/utils/bgtask.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef BGTASK_H
+#define BGTASK_H
+
+/*
+ * The type of internal background tasks.
+ *
+ * When adding a new task, please change InternalBgTask array in bgtask.c
+ * too.
+ */
+typedef enum
+{
+	BGTASK_TYPE_DUMMY = 0,
+}			BgTaskType;
+
+/* The number of internal background tasks */
+#define NUM_INTERNAL_BGTASKS	BGTASK_TYPE_DUMMY + 1
+
+typedef void (*bgtask_fn_type) (void);
+
+extern Size BgTaskShmemSize(void);
+extern void BgTaskShmemInit(void);
+extern void BgTaskRequest(BgTaskType type, Oid dboid);
+extern void BgTaskWorkerMain(Datum main_arg);
+
+#endif							/* BTASK_H */
-- 
2.43.5

