Support to define custom wait events for extensions

From 206ce9e1d74d1449b50cfc765936172de98687e8 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Tue, 20 Jun 2023 17:59:34 +0900
Subject: [PATCH 2/2] Add test codes for custom wait events

---
 src/test/modules/Makefile                     |   1 +
 .../test_custom_wait_events/.gitignore        |   2 +
 .../modules/test_custom_wait_events/Makefile  |  23 +++
 .../test_custom_wait_events/meson.build       |  33 ++++
 .../test_custom_wait_events/t/001_basic.pl    |  34 ++++
 .../test_custom_wait_events--1.0.sql          |  14 ++
 .../test_custom_wait_events.c                 | 182 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 8 files changed, 292 insertions(+)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..716e17f5a2
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,2 @@
+# Generated subdirectories
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..da2b720f8f
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) 2022-2023, PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..59ba14919b
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,34 @@
+# Copyright (c) 2021-2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+$node->init;
+$node->append_conf(
+	'postgresql.conf', 
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+
+# setup
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT wait_worker_launch();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT 
+		(SELECT count(*) FROM pg_stat_activity
+	 		WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) > 0;])
+) or die "Timed out waiting the custom wait event to be showed up";
+
+$node->stop;
+done_testing();
\ No newline at end of file
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..1fc2130f84
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,14 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- wait_worker_launch()
+--
+-- Launch a background worker to wait forever with a custom wait event.
+--
+CREATE FUNCTION wait_worker_launch()
+RETURNS int
+AS 'MODULE_PATHNAME', 'wait_worker_launch'
+LANGUAGE C STRICT PARALLEL SAFE;
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..ba3fe45ad6
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,182 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events
+ *
+ * This code initializes a custom wait event in shmem_request_hook() and
+ * provide a function to launch a background worker waiting forever
+ * with the custom wait event.
+ * 
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(wait_worker_launch);
+
+PGDLLEXPORT void worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct state
+{
+	uint32			wait_event_info;  /* a allocated wait event */
+}			state;
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static char * worker_database = "postgres";   /* connected by a background worker */
+static state * ss = NULL;
+
+static void shmem_request(void);
+static void shmem_startup(void);
+static Size memsize(void);
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = shmem_startup;
+}
+
+static void
+shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(memsize());
+}
+
+static Size
+memsize(void)
+{
+	return sizeof(state);
+}
+
+/*
+ * Allocate and register a new wait event.
+ *
+ * This routine allocate a wait event id when it's called first time.
+ * After that, the name of the wait event is associated with the id. It's
+ * registered in the lookup table of each processes dynamically.
+ */
+static void
+shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(state),
+						 &found);
+
+	/* Allocate a new wait event. */
+	if (!found)
+		ss->wait_event_info = ExtensionWaitEventNewTranche();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	ExtensionWaitEventRegisterTranche(ss->wait_event_info, "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the custom wait event
+ */
+void
+worker_main(Datum main_arg)
+{
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the custom wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 ss->wait_event_info);  /* the custom wait event */
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically launch a background worker waiting forever
+ */
+Datum
+wait_worker_launch(PG_FUNCTION_ARGS)
+{
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.25.1

From 8a8093c7cdfff785bdde1d8e830367e5783ab37b Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Fri, 16 Jun 2023 11:53:29 +0900
Subject: [PATCH 1/2] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to allocate
new wait events dynamically.

The APIs are
* ExtensionWaitEventNewTranche()
* ExtensionWaitEventRegisterTranche()

They similar to existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extension calls ExtensionWaitEventNewTranche()
just once to obtain a new wait event tranche; The ID is
allocated from a shared counter. Next, each individual
process using the tranche should call
ExtensionWaitEventRegisterTranche() to associate that
wait event with a name.
---
 src/backend/storage/ipc/ipci.c          |   3 +
 src/backend/storage/ipc/shmem.c         |   3 +-
 src/backend/utils/activity/wait_event.c | 138 +++++++++++++++++++++++-
 src/include/utils/wait_event.h          |  29 +++++
 4 files changed, 171 insertions(+), 2 deletions(-)

diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 8f1ded7338..ed05121fa3 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -294,6 +296,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 5465fa1964..02c72ebbb1 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -85,7 +85,8 @@ static void *ShmemBase;			/* start address of shared memory */
 
 static void *ShmemEnd;			/* end+1 address of shared memory */
 
-slock_t    *ShmemLock;			/* spinlock for shared memory and LWLock
+slock_t    *ShmemLock;			/* spinlock for shared memory, LWLock
+								 * allocation, and named extension wait event
 								 * allocation */
 
 static HTAB *ShmemIndex = NULL; /* primary index hashtable for shmem */
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 7940d64639..724a720f2c 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,11 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
+/* We use the ShmemLock spinlock to protect ExtensionWaitEventCounter */
+extern slock_t *ShmemLock;
+
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
@@ -38,6 +45,23 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+/* dynamic allocation counter in shared memory */
+static uint16 *ExtensionWaitEventCounter;
+
+/* first event ID for named extension wait event tranche */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event ID known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **ExtensionWaitEventTrancheNames = NULL;
+static int	ExtensionWaitEventTrancheNamesAllocated = 0;
+
+static const char *GetExtensionWaitEventIdentifier(uint16 eventId);
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -165,7 +189,7 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				break;
 			}
 		case PG_WAIT_EXTENSION:
-			event_name = "Extension";
+			event_name = GetExtensionWaitEventIdentifier(eventId);
 			break;
 		case PG_WAIT_IPC:
 			{
@@ -762,3 +786,115 @@ pgstat_get_wait_io(WaitEventIO w)
 
 	return event_name;
 }
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventShmemSize(void)
+{
+	return sizeof(*ExtensionWaitEventCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventShmemInit(void)
+{
+	if (!IsUnderPostmaster)
+	{
+		Size		space = WaitEventShmemSize();
+
+		/* Initialize the dynamic-allocation counter */
+		ExtensionWaitEventCounter = (uint16 *) ShmemAlloc(space);
+		*ExtensionWaitEventCounter = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event info.
+ */
+uint32
+ExtensionWaitEventNewTranche(void)
+{
+	uint16			eventId;
+
+	SpinLockAcquire(ShmemLock);
+	eventId = (*ExtensionWaitEventCounter)++;
+	SpinLockRelease(ShmemLock);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic tranche name in the lookup table of the current process.
+ *
+ * This routine will save a pointer to the wait event tranche name passed
+ * as an argument, so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+ExtensionWaitEventRegisterTranche(uint32 wait_event_info, const char *wait_event_name)
+{
+	uint32		classId;
+	uint16		eventId;
+
+	classId = wait_event_info & 0xFF000000;
+	eventId = wait_event_info & 0x0000FFFF;
+	Assert(classId == PG_WAIT_EXTENSION);
+	
+	/* This should only be called for user-defined tranches. */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return;
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= ExtensionWaitEventTrancheNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (ExtensionWaitEventTrancheNames == NULL)
+			ExtensionWaitEventTrancheNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			ExtensionWaitEventTrancheNames =
+				repalloc0_array(ExtensionWaitEventTrancheNames, const char *, ExtensionWaitEventTrancheNamesAllocated, newalloc);
+		ExtensionWaitEventTrancheNamesAllocated = newalloc;
+	}
+
+	ExtensionWaitEventTrancheNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an Extension wait event ID.
+ */
+static const char *
+GetExtensionWaitEventIdentifier(uint16 eventId)
+{
+	/* Build-in tranche? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an extension tranche, so look in ExtensionWaitEventTrancheNames[].
+	 * However, it's possible that the tranche has never been registered by
+	 * calling ExtensionWaitEventRegisterTranche() in the current process, in
+	 * which case give up and return "Extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= ExtensionWaitEventTrancheNamesAllocated ||
+		ExtensionWaitEventTrancheNames[eventId] == NULL)
+		return "Extension";
+
+	return ExtensionWaitEventTrancheNames[eventId];
+}
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 518d3b0a1f..789fcbf52a 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -70,6 +70,35 @@ typedef enum
 	WAIT_EVENT_WAL_SENDER_WRITE_DATA,
 } WaitEventClient;
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define custom wait events.  First, call
+ * ExtensionWaitEventNewTranche() just once to obtain a new wait event tranche. 
+ * The ID is allocated from a shared counter. Next, each individual process using
+ * the tranche should call ExtensionWaitEventRegisterTranche() to associate
+ * that wait event with a name.
+ *
+ * It may seem strange that each process using the tranche must register it
+ * separately, but dynamic shared memory segments aren't guaranteed to be
+ * mapped at the same address in all coordinating backends, so storing the
+ * registration in the main shared memory segment wouldn't work for that case.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+}			WaitEventExtension;
+
+extern void WaitEventShmemInit(void);
+extern Size WaitEventShmemSize(void);
+
+extern uint32	ExtensionWaitEventNewTranche(void);
+extern void ExtensionWaitEventRegisterTranche(uint32 wait_event_info, const char *tranche_name);
+
 /* ----------
  * Wait Events - IPC
  *
-- 
2.25.1

From 901b75d31070ef0029557db6981c98e06f5c16c3 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Tue, 20 Jun 2023 17:59:34 +0900
Subject: [PATCH 2/2] Add test codes for custom wait events

---
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../test_custom_wait_events/.gitignore        |   2 +
 .../modules/test_custom_wait_events/Makefile  |  23 +++
 .../test_custom_wait_events/meson.build       |  33 ++++
 .../test_custom_wait_events/t/001_basic.pl    |  34 ++++
 .../test_custom_wait_events--1.0.sql          |  14 ++
 .../test_custom_wait_events.c                 | 182 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 9 files changed, 293 insertions(+)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 17d369e378..83a1d8fd19 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -14,6 +14,7 @@ subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_wait_events')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..716e17f5a2
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,2 @@
+# Generated subdirectories
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..8ad073e577
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..2809afbd47
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,34 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+$node->init;
+$node->append_conf(
+	'postgresql.conf', 
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+
+# setup
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT wait_worker_launch();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT 
+		(SELECT count(*) FROM pg_stat_activity
+	 		WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) > 0;])
+) or die "Timed out waiting the custom wait event to be showed up";
+
+$node->stop;
+done_testing();
\ No newline at end of file
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..1fc2130f84
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,14 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- wait_worker_launch()
+--
+-- Launch a background worker to wait forever with a custom wait event.
+--
+CREATE FUNCTION wait_worker_launch()
+RETURNS int
+AS 'MODULE_PATHNAME', 'wait_worker_launch'
+LANGUAGE C STRICT PARALLEL SAFE;
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..ba3fe45ad6
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,182 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events
+ *
+ * This code initializes a custom wait event in shmem_request_hook() and
+ * provide a function to launch a background worker waiting forever
+ * with the custom wait event.
+ * 
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(wait_worker_launch);
+
+PGDLLEXPORT void worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct state
+{
+	uint32			wait_event_info;  /* a allocated wait event */
+}			state;
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static char * worker_database = "postgres";   /* connected by a background worker */
+static state * ss = NULL;
+
+static void shmem_request(void);
+static void shmem_startup(void);
+static Size memsize(void);
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = shmem_startup;
+}
+
+static void
+shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(memsize());
+}
+
+static Size
+memsize(void)
+{
+	return sizeof(state);
+}
+
+/*
+ * Allocate and register a new wait event.
+ *
+ * This routine allocate a wait event id when it's called first time.
+ * After that, the name of the wait event is associated with the id. It's
+ * registered in the lookup table of each processes dynamically.
+ */
+static void
+shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(state),
+						 &found);
+
+	/* Allocate a new wait event. */
+	if (!found)
+		ss->wait_event_info = ExtensionWaitEventNewTranche();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	ExtensionWaitEventRegisterTranche(ss->wait_event_info, "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the custom wait event
+ */
+void
+worker_main(Datum main_arg)
+{
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the custom wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 ss->wait_event_info);  /* the custom wait event */
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically launch a background worker waiting forever
+ */
+Datum
+wait_worker_launch(PG_FUNCTION_ARGS)
+{
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.25.1

From bf06b8100cb747031959fe81a2d19baabc4838cf Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Fri, 16 Jun 2023 11:53:29 +0900
Subject: [PATCH 1/2] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to allocate
new wait events dynamically.

The APIs are
* ExtensionWaitEventNewTranche()
* ExtensionWaitEventRegisterTranche()

They similar to existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extension calls ExtensionWaitEventNewTranche()
just once to obtain a new wait event tranche; The ID is
allocated from a shared counter. Next, each individual
process using the tranche should call
ExtensionWaitEventRegisterTranche() to associate that
wait event with a name.
---
 src/backend/storage/ipc/ipci.c          |   3 +
 src/backend/storage/ipc/shmem.c         |   3 +-
 src/backend/utils/activity/wait_event.c | 138 +++++++++++++++++++++++-
 src/include/utils/wait_event.h          |  29 +++++
 4 files changed, 171 insertions(+), 2 deletions(-)

diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 8f1ded7338..ed05121fa3 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -294,6 +296,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 5465fa1964..02c72ebbb1 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -85,7 +85,8 @@ static void *ShmemBase;			/* start address of shared memory */
 
 static void *ShmemEnd;			/* end+1 address of shared memory */
 
-slock_t    *ShmemLock;			/* spinlock for shared memory and LWLock
+slock_t    *ShmemLock;			/* spinlock for shared memory, LWLock
+								 * allocation, and named extension wait event
 								 * allocation */
 
 static HTAB *ShmemIndex = NULL; /* primary index hashtable for shmem */
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 7940d64639..724a720f2c 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,11 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
+/* We use the ShmemLock spinlock to protect ExtensionWaitEventCounter */
+extern slock_t *ShmemLock;
+
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
@@ -38,6 +45,23 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+/* dynamic allocation counter in shared memory */
+static uint16 *ExtensionWaitEventCounter;
+
+/* first event ID for named extension wait event tranche */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event ID known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **ExtensionWaitEventTrancheNames = NULL;
+static int	ExtensionWaitEventTrancheNamesAllocated = 0;
+
+static const char *GetExtensionWaitEventIdentifier(uint16 eventId);
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -165,7 +189,7 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				break;
 			}
 		case PG_WAIT_EXTENSION:
-			event_name = "Extension";
+			event_name = GetExtensionWaitEventIdentifier(eventId);
 			break;
 		case PG_WAIT_IPC:
 			{
@@ -762,3 +786,115 @@ pgstat_get_wait_io(WaitEventIO w)
 
 	return event_name;
 }
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventShmemSize(void)
+{
+	return sizeof(*ExtensionWaitEventCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventShmemInit(void)
+{
+	if (!IsUnderPostmaster)
+	{
+		Size		space = WaitEventShmemSize();
+
+		/* Initialize the dynamic-allocation counter */
+		ExtensionWaitEventCounter = (uint16 *) ShmemAlloc(space);
+		*ExtensionWaitEventCounter = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event info.
+ */
+uint32
+ExtensionWaitEventNewTranche(void)
+{
+	uint16			eventId;
+
+	SpinLockAcquire(ShmemLock);
+	eventId = (*ExtensionWaitEventCounter)++;
+	SpinLockRelease(ShmemLock);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic tranche name in the lookup table of the current process.
+ *
+ * This routine will save a pointer to the wait event tranche name passed
+ * as an argument, so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+ExtensionWaitEventRegisterTranche(uint32 wait_event_info, const char *wait_event_name)
+{
+	uint32		classId;
+	uint16		eventId;
+
+	classId = wait_event_info & 0xFF000000;
+	eventId = wait_event_info & 0x0000FFFF;
+	Assert(classId == PG_WAIT_EXTENSION);
+	
+	/* This should only be called for user-defined tranches. */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return;
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= ExtensionWaitEventTrancheNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (ExtensionWaitEventTrancheNames == NULL)
+			ExtensionWaitEventTrancheNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			ExtensionWaitEventTrancheNames =
+				repalloc0_array(ExtensionWaitEventTrancheNames, const char *, ExtensionWaitEventTrancheNamesAllocated, newalloc);
+		ExtensionWaitEventTrancheNamesAllocated = newalloc;
+	}
+
+	ExtensionWaitEventTrancheNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an Extension wait event ID.
+ */
+static const char *
+GetExtensionWaitEventIdentifier(uint16 eventId)
+{
+	/* Build-in tranche? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an extension tranche, so look in ExtensionWaitEventTrancheNames[].
+	 * However, it's possible that the tranche has never been registered by
+	 * calling ExtensionWaitEventRegisterTranche() in the current process, in
+	 * which case give up and return "Extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= ExtensionWaitEventTrancheNamesAllocated ||
+		ExtensionWaitEventTrancheNames[eventId] == NULL)
+		return "Extension";
+
+	return ExtensionWaitEventTrancheNames[eventId];
+}
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 518d3b0a1f..789fcbf52a 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -70,6 +70,35 @@ typedef enum
 	WAIT_EVENT_WAL_SENDER_WRITE_DATA,
 } WaitEventClient;
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define custom wait events.  First, call
+ * ExtensionWaitEventNewTranche() just once to obtain a new wait event tranche. 
+ * The ID is allocated from a shared counter. Next, each individual process using
+ * the tranche should call ExtensionWaitEventRegisterTranche() to associate
+ * that wait event with a name.
+ *
+ * It may seem strange that each process using the tranche must register it
+ * separately, but dynamic shared memory segments aren't guaranteed to be
+ * mapped at the same address in all coordinating backends, so storing the
+ * registration in the main shared memory segment wouldn't work for that case.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+}			WaitEventExtension;
+
+extern void WaitEventShmemInit(void);
+extern Size WaitEventShmemSize(void);
+
+extern uint32	ExtensionWaitEventNewTranche(void);
+extern void ExtensionWaitEventRegisterTranche(uint32 wait_event_info, const char *tranche_name);
+
 /* ----------
  * Wait Events - IPC
  *
-- 
2.25.1

From 1bb78fa2cbe6b030cea7a570bec88bd4d68f314a Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Fri, 23 Jun 2023 17:38:38 +0900
Subject: [PATCH 2/2] Add test codes for custom wait events

---
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../test_custom_wait_events/.gitignore        |   2 +
 .../modules/test_custom_wait_events/Makefile  |  23 +++
 .../test_custom_wait_events/meson.build       |  33 ++++
 .../test_custom_wait_events/t/001_basic.pl    |  34 ++++
 .../test_custom_wait_events--1.0.sql          |  14 ++
 .../test_custom_wait_events.c                 | 182 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 9 files changed, 293 insertions(+)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 17d369e378..83a1d8fd19 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -14,6 +14,7 @@ subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_wait_events')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..716e17f5a2
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,2 @@
+# Generated subdirectories
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..8ad073e577
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..b43c82a713
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,34 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+
+# setup
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT wait_worker_launch();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) > 0;])
+) or die "Timed out waiting the custom wait event to be showed up";
+
+$node->stop;
+done_testing();
\ No newline at end of file
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..1fc2130f84
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,14 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- wait_worker_launch()
+--
+-- Launch a background worker to wait forever with a custom wait event.
+--
+CREATE FUNCTION wait_worker_launch()
+RETURNS int
+AS 'MODULE_PATHNAME', 'wait_worker_launch'
+LANGUAGE C STRICT PARALLEL SAFE;
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..6a47665524
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,182 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events
+ *
+ * This code initializes a custom wait event in shmem_request_hook() and
+ * provide a function to launch a background worker waiting forever
+ * with the custom wait event.
+ *
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(wait_worker_launch);
+
+PGDLLEXPORT void worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct state
+{
+	uint32			wait_event_info;  /* a allocated wait event */
+}			state;
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static char * worker_database = "postgres";   /* connected by a background worker */
+static state * ss = NULL;
+
+static void shmem_request(void);
+static void shmem_startup(void);
+static Size memsize(void);
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = shmem_startup;
+}
+
+static void
+shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(memsize());
+}
+
+static Size
+memsize(void)
+{
+	return sizeof(state);
+}
+
+/*
+ * Allocate and register a new wait event.
+ *
+ * This routine allocate a wait event id when it's called first time.
+ * After that, the name of the wait event is associated with the id. It's
+ * registered in the lookup table of each processes dynamically.
+ */
+static void
+shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(state),
+						 &found);
+
+	/* Allocate a new wait event. */
+	if (!found)
+		ss->wait_event_info = ExtensionWaitEventNewTranche();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	ExtensionWaitEventRegisterTranche(ss->wait_event_info, "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the custom wait event
+ */
+void
+worker_main(Datum main_arg)
+{
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the custom wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 ss->wait_event_info);  /* the custom wait event */
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically launch a background worker waiting forever
+ */
+Datum
+wait_worker_launch(PG_FUNCTION_ARGS)
+{
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.25.1

From 8cf1f026fc04633412f03ace94f24dd4692d7310 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Fri, 16 Jun 2023 11:53:29 +0900
Subject: [PATCH 1/2] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to allocate
new wait events dynamically.

The APIs are
* ExtensionWaitEventNewTranche()
* ExtensionWaitEventRegisterTranche()

They similar to existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extension calls ExtensionWaitEventNewTranche()
just once to obtain a new wait event tranche; The ID is
allocated from a shared counter. Next, each individual
process using the tranche should call
ExtensionWaitEventRegisterTranche() to associate that
wait event with a name.
---
 src/backend/storage/ipc/ipci.c          |   3 +
 src/backend/storage/ipc/shmem.c         |   3 +-
 src/backend/utils/activity/wait_event.c | 138 +++++++++++++++++++++++-
 src/include/utils/wait_event.h          |  29 +++++
 4 files changed, 171 insertions(+), 2 deletions(-)

diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 8f1ded7338..ed05121fa3 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -294,6 +296,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 5465fa1964..02c72ebbb1 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -85,7 +85,8 @@ static void *ShmemBase;			/* start address of shared memory */
 
 static void *ShmemEnd;			/* end+1 address of shared memory */
 
-slock_t    *ShmemLock;			/* spinlock for shared memory and LWLock
+slock_t    *ShmemLock;			/* spinlock for shared memory, LWLock
+								 * allocation, and named extension wait event
 								 * allocation */
 
 static HTAB *ShmemIndex = NULL; /* primary index hashtable for shmem */
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 7940d64639..078bb221b8 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,11 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
+/* We use the ShmemLock spinlock to protect ExtensionWaitEventCounter */
+extern slock_t *ShmemLock;
+
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
@@ -38,6 +45,23 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+/* dynamic allocation counter in shared memory */
+static uint16 *ExtensionWaitEventCounter;
+
+/* first event ID for named extension wait event tranche */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event ID known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **ExtensionWaitEventTrancheNames = NULL;
+static int	ExtensionWaitEventTrancheNamesAllocated = 0;
+
+static const char *GetExtensionWaitEventIdentifier(uint16 eventId);
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -165,7 +189,7 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				break;
 			}
 		case PG_WAIT_EXTENSION:
-			event_name = "Extension";
+			event_name = GetExtensionWaitEventIdentifier(eventId);
 			break;
 		case PG_WAIT_IPC:
 			{
@@ -762,3 +786,115 @@ pgstat_get_wait_io(WaitEventIO w)
 
 	return event_name;
 }
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventShmemSize(void)
+{
+	return sizeof(*ExtensionWaitEventCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventShmemInit(void)
+{
+	if (!IsUnderPostmaster)
+	{
+		Size		space = WaitEventShmemSize();
+
+		/* Initialize the dynamic-allocation counter */
+		ExtensionWaitEventCounter = (uint16 *) ShmemAlloc(space);
+		*ExtensionWaitEventCounter = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event info.
+ */
+uint32
+ExtensionWaitEventNewTranche(void)
+{
+	uint16			eventId;
+
+	SpinLockAcquire(ShmemLock);
+	eventId = (*ExtensionWaitEventCounter)++;
+	SpinLockRelease(ShmemLock);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic tranche name in the lookup table of the current process.
+ *
+ * This routine will save a pointer to the wait event tranche name passed
+ * as an argument, so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+ExtensionWaitEventRegisterTranche(uint32 wait_event_info, const char *wait_event_name)
+{
+	uint16		eventId;
+
+	/* Check wait event class. */
+	Assert((wait_event_info & 0xFF000000) == PG_WAIT_EXTENSION);
+
+	eventId = wait_event_info & 0x0000FFFF;
+
+	/* This should only be called for user-defined tranches. */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return;
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= ExtensionWaitEventTrancheNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (ExtensionWaitEventTrancheNames == NULL)
+			ExtensionWaitEventTrancheNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			ExtensionWaitEventTrancheNames =
+				repalloc0_array(ExtensionWaitEventTrancheNames, const char *, ExtensionWaitEventTrancheNamesAllocated, newalloc);
+		ExtensionWaitEventTrancheNamesAllocated = newalloc;
+	}
+
+	ExtensionWaitEventTrancheNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an Extension wait event ID.
+ */
+static const char *
+GetExtensionWaitEventIdentifier(uint16 eventId)
+{
+	/* Build-in tranche? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an extension tranche, so look in ExtensionWaitEventTrancheNames[].
+	 * However, it's possible that the tranche has never been registered by
+	 * calling ExtensionWaitEventRegisterTranche() in the current process, in
+	 * which case give up and return "Extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= ExtensionWaitEventTrancheNamesAllocated ||
+		ExtensionWaitEventTrancheNames[eventId] == NULL)
+		return "Extension";
+
+	return ExtensionWaitEventTrancheNames[eventId];
+}
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 518d3b0a1f..9507c4fd71 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -70,6 +70,35 @@ typedef enum
 	WAIT_EVENT_WAL_SENDER_WRITE_DATA,
 } WaitEventClient;
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define custom wait events.  First, call
+ * ExtensionWaitEventNewTranche() just once to obtain a new wait event tranche.
+ * The ID is allocated from a shared counter. Next, each individual process using
+ * the tranche should call ExtensionWaitEventRegisterTranche() to associate
+ * that wait event with a name.
+ *
+ * It may seem strange that each process using the tranche must register it
+ * separately, but dynamic shared memory segments aren't guaranteed to be
+ * mapped at the same address in all coordinating backends, so storing the
+ * registration in the main shared memory segment wouldn't work for that case.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+}			WaitEventExtension;
+
+extern void WaitEventShmemInit(void);
+extern Size WaitEventShmemSize(void);
+
+extern uint32	ExtensionWaitEventNewTranche(void);
+extern void ExtensionWaitEventRegisterTranche(uint32 wait_event_info, const char *tranche_name);
+
 /* ----------
  * Wait Events - IPC
  *
-- 
2.25.1

From 630f78822534171bb8839c79d0fdca3bb8268a95 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Tue, 11 Jul 2023 16:41:35 +0900
Subject: [PATCH v5] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to allocate new wait events
dynamically.

The APIs are
* ExtensionWaitEventNewTranche()
* ExtensionWaitEventRegisterTranche()

These are similar to the existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extension should call ExtensionWaitEventNewTranche() to get one
or more wait event "tranches", which are ID is allocated from a shared
counter.  Next, each individual process can use the tranche numbers with
ExtensionWaitEventRegisterTranche() to associate that a wait event
string to the associated name.

Note that this includes a test module, which perhaps should be split
later into its own commit.
---
 src/include/utils/wait_event.h                |  31 +++
 src/backend/storage/ipc/ipci.c                |   3 +
 src/backend/storage/ipc/shmem.c               |   3 +-
 .../activity/generate-wait_event_types.pl     |  12 +-
 src/backend/utils/activity/wait_event.c       | 149 +++++++++++++-
 .../utils/activity/wait_event_names.txt       |   2 +-
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../test_custom_wait_events/.gitignore        |   2 +
 .../modules/test_custom_wait_events/Makefile  |  23 +++
 .../test_custom_wait_events/meson.build       |  33 +++
 .../test_custom_wait_events/t/001_basic.pl    |  34 ++++
 .../test_custom_wait_events--1.0.sql          |  14 ++
 .../test_custom_wait_events.c                 | 188 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 15 files changed, 484 insertions(+), 15 deletions(-)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..dfcaf8c506 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -67,6 +67,37 @@ pgstat_report_wait_start(uint32 wait_event_info)
 	*(volatile uint32 *) my_wait_event_info = wait_event_info;
 }
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define custom wait events.  First, call
+ * WaitEventExtensionNewTranche() just once to obtain a new wait event
+ * tranche.  The ID is allocated from a shared counter.  Next, each
+ * individual process using the tranche should call
+ * WaitEventExtensionRegisterTranche() to associate that wait event with
+ * a name.
+ *
+ * It may seem strange that each process using the tranche must register it
+ * separately, but dynamic shared memory segments aren't guaranteed to be
+ * mapped at the same address in all coordinating backends, so storing the
+ * registration in the main shared memory segment wouldn't work for that case.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNewTranche(void);
+extern void WaitEventExtensionRegisterTranche(uint32 wait_event_info,
+											  const char *tranche_name);
+
 /* ----------
  * pgstat_report_wait_end() -
  *
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 5465fa1964..02c72ebbb1 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -85,7 +85,8 @@ static void *ShmemBase;			/* start address of shared memory */
 
 static void *ShmemEnd;			/* end+1 address of shared memory */
 
-slock_t    *ShmemLock;			/* spinlock for shared memory and LWLock
+slock_t    *ShmemLock;			/* spinlock for shared memory, LWLock
+								 * allocation, and named extension wait event
 								 * allocation */
 
 static HTAB *ShmemIndex = NULL; /* primary index hashtable for shmem */
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index 6d1a2af42a..2cad28fbfe 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -85,10 +85,11 @@ foreach my $line (@lines_sorted)
 	my $trimmedwaiteventname = $waiteventenumname;
 	$trimmedwaiteventname =~ s/^WAIT_EVENT_//;
 
-	# An exception is required for LWLock and Lock as these don't require
-	# any C and header files generated.
+	# An exception is required for Extension, LWLock and Lock as these don't
+	# require any C and header files generated.
 	die "wait event names must start with 'WAIT_EVENT_'"
 	  if ( $trimmedwaiteventname eq $waiteventenumname
+		&& $waiteventenumname !~ /^Extension/
 		&& $waiteventenumname !~ /^LWLock/
 		&& $waiteventenumname !~ /^Lock/);
 	$continue = ",\n";
@@ -138,11 +139,12 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
 		  if ( $waitclass =~ /^WaitEventLWLock$/
-			|| $waitclass =~ /^WaitEventLock$/);
+			|| $waitclass =~ /^WaitEventLock$/
+			|| $waitclass =~ /^WaitEventExtension$/);
 
 		my $last = $waitclass;
 		$last =~ s/^WaitEvent//;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 4aad11c111..9d4652fcc9 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,21 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
+/* We use the ShmemLock spinlock to protect WaitEventExtensionCounter */
+extern slock_t *ShmemLock;
+
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -40,6 +46,137 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+/* dynamic allocation counter in shared memory */
+static uint16 *WaitEventExtensionCounter;
+
+/* first event ID for named extension wait event tranche */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event ID known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionTrancheNames = NULL;
+static int	WaitEventExtensionTrancheNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(*WaitEventExtensionCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	if (!IsUnderPostmaster)
+	{
+		Size		space = WaitEventExtensionShmemSize();
+
+		/* Initialize the dynamic-allocation counter */
+		WaitEventExtensionCounter = (uint16 *) ShmemAlloc(space);
+		*WaitEventExtensionCounter = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event info.
+ */
+uint32
+WaitEventExtensionNewTranche(void)
+{
+	uint16		eventId;
+
+	SpinLockAcquire(ShmemLock);
+	eventId = (*WaitEventExtensionCounter)++;
+	SpinLockRelease(ShmemLock);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic tranche name in the lookup table of the current process.
+ *
+ * This routine will save a pointer to the wait event tranche name passed
+ * as an argument, so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterTranche(uint32 wait_event_info,
+								  const char *wait_event_name)
+{
+	uint16		eventId;
+
+	/* Check wait event class. */
+	Assert((wait_event_info & 0xFF000000) == PG_WAIT_EXTENSION);
+
+	eventId = wait_event_info & 0x0000FFFF;
+
+	/* This should only be called for user-defined tranches. */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return;
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionTrancheNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionTrancheNames == NULL)
+			WaitEventExtensionTrancheNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionTrancheNames =
+				repalloc0_array(WaitEventExtensionTrancheNames, const char *, WaitEventExtensionTrancheNamesAllocated, newalloc);
+		WaitEventExtensionTrancheNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionTrancheNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an Extension wait event ID.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Build-in tranche? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an extension tranche, so look in WaitEventExtensionTrancheNames[].
+	 * However, it's possible that the tranche has never been registered by
+	 * calling WaitEventExtensionRegisterTranche() in the current process, in
+	 * which case give up and return "Extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionTrancheNamesAllocated ||
+		WaitEventExtensionTrancheNames[eventId] == NULL)
+		return "Extension";
+
+	return WaitEventExtensionTrancheNames[eventId];
+}
+
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -149,6 +286,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -170,13 +310,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 7af1a61f95..d3cddde95f 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	"BufferPin"	"Waiting to acquire an exclusive pin on a buff
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	"Extension"	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	"Extension"	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 17d369e378..83a1d8fd19 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -14,6 +14,7 @@ subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_wait_events')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..716e17f5a2
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,2 @@
+# Generated subdirectories
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..8ad073e577
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..339a9f66ba
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,34 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+
+# setup
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT wait_worker_launch();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) > 0;])
+) or die "Timed out waiting the custom wait event to be showed up";
+
+$node->stop;
+done_testing();
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..1fc2130f84
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,14 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- wait_worker_launch()
+--
+-- Launch a background worker to wait forever with a custom wait event.
+--
+CREATE FUNCTION wait_worker_launch()
+RETURNS int
+AS 'MODULE_PATHNAME', 'wait_worker_launch'
+LANGUAGE C STRICT PARALLEL SAFE;
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..f6d3739411
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,188 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events
+ *
+ * This code initializes a custom wait event in shmem_request_hook() and
+ * provide a function to launch a background worker waiting forever
+ * with the custom wait event.
+ *
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(wait_worker_launch);
+
+PGDLLEXPORT void custom_we_worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct CustomWaitEventState
+{
+	uint32			wait_event_info;  /* a allocated wait event */
+}			CustomWaitEventState;
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static char *worker_database = "postgres";   /* connected by a background worker */
+static CustomWaitEventState *ss = NULL;
+
+static void custom_we_shmem_request(void);
+static void custom_we_shmem_startup(void);
+static Size custom_we_memsize(void);
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = custom_we_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = custom_we_shmem_startup;
+}
+
+static void
+custom_we_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(custom_we_memsize());
+}
+
+static Size
+custom_we_memsize(void)
+{
+	return sizeof(CustomWaitEventState);
+}
+
+/*
+ * Allocate and register a new wait event.
+ *
+ * This routine allocate a wait event id when it's called first time.
+ * After that, the name of the wait event is associated with the id. It's
+ * registered in the lookup table of each processes dynamically.
+ */
+static void
+custom_we_shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(CustomWaitEventState),
+						 &found);
+
+	/* Allocate a new wait event. */
+	if (!found)
+		ss->wait_event_info = WaitEventExtensionNewTranche();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterTranche(ss->wait_event_info,
+									  "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the custom wait event
+ */
+void
+custom_we_worker_main(Datum main_arg)
+{
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the custom wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 ss->wait_event_info);  /* the custom wait event */
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically launch a background worker waiting forever
+ */
+Datum
+wait_worker_launch(PG_FUNCTION_ARGS)
+{
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	if (ss == NULL)
+		ereport(ERROR,
+				(errmsg("cannot use \"%s\" if \"%s\" has not been loaded with shared_preload_libraries",
+						"wait_worker_launch", "test_custom_wait_events")));
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "custom_we_worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.40.1

From 5198fc6302a3bf4232252b23b45d39d987e736bc Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Wed, 12 Jul 2023 16:28:27 +0900
Subject: [PATCH] change the die condition in generate-wait_event_types.pl

---
 src/backend/utils/activity/generate-wait_event_types.pl | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index 6d1a2af42a..fbe011039a 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -89,9 +89,8 @@ foreach my $line (@lines_sorted)
 	# any C and header files generated.
 	die "wait event names must start with 'WAIT_EVENT_'"
 	  if ( $trimmedwaiteventname eq $waiteventenumname
-		&& $waiteventenumname !~ /^LWLock/
-		&& $waiteventenumname !~ /^Lock/);
-	$continue = ",\n";
+		&& $waitclassname !~ /^WaitEventLWLock$/
+		&& $waitclassname !~ /^WaitEventLock$/);
 	push(@{ $hashwe{$waitclassname} }, @waiteventlist);
 }
 
-- 
2.25.1

From 9c5c1f1f4f0be87cc06c7127396e7d228b665b8a Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Wed, 19 Jul 2023 12:28:12 +0900
Subject: [PATCH] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to define new wait events
for extensions dynamically.

The APIs are
* WaitEventExtensionNew()
* WaitEventExtensionRegisterName()

These are similar to the existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extensions should call WaitEventExtensionNew() to get one
or more new wait event, which IDs are allocated from a shared
counter.  Next, each individual process can use the wait event with
WaitEventExtensionRegisterName() to associate that a wait event
string to the associated name.

Note that this includes a test module, which perhaps should be split
later into its own commit.
---
 doc/src/sgml/monitoring.sgml                  |  10 +-
 doc/src/sgml/xfunc.sgml                       |  23 ++
 src/backend/storage/ipc/ipci.c                |   3 +
 .../activity/generate-wait_event_types.pl     |   7 +-
 src/backend/utils/activity/wait_event.c       | 171 +++++++++++-
 .../utils/activity/wait_event_names.txt       |   2 +-
 src/include/utils/wait_event.h                |  25 ++
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../test_custom_wait_events/.gitignore        |   4 +
 .../modules/test_custom_wait_events/Makefile  |  23 ++
 .../test_custom_wait_events/meson.build       |  33 +++
 .../test_custom_wait_events/t/001_basic.pl    | 137 ++++++++++
 .../test_custom_wait_events--1.0.sql          |  39 +++
 .../test_custom_wait_events.c                 | 253 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 16 files changed, 718 insertions(+), 17 deletions(-)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index 588b720f57..3aee5c243f 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1117,12 +1117,12 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
 
    <note>
     <para>
-     Extensions can add <literal>LWLock</literal> types to the list shown in
-     <xref linkend="wait-event-lwlock-table"/>.  In some cases, the name
+     Extensions can add <literal>Extension</literal> and <literal>LWLock</literal> types
+     to the list shown in <xref linkend="wait-event-extension-table"/> and
+     <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
      assigned by an extension will not be available in all server processes;
-     so an <literal>LWLock</literal> wait event might be reported as
-     just <quote><literal>extension</literal></quote> rather than the
-     extension-assigned name.
+     so an wait event might be reported as just <quote><literal>extension</literal></quote>
+     rather than the extension-assigned name.
     </para>
    </note>
  </sect2>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9620ea9ae3..bc210cd03b 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3453,6 +3453,29 @@ if (!ptr)
     </para>
    </sect2>
 
+   <sect2 id="xfunc-addin-wait-events">
+    <title>Custom Wait Events for Add-ins</title>
+
+    <para>
+    Add-ins can define custom wait events that the wait event type is
+    <literal>Extension</literal>.
+    </para>
+    <para>
+    First, add-ins should get new one or more wait events by calling:
+<programlisting>
+    uint32 WaitEventExtensionNew(void)
+</programlisting>
+    Next, each individual process can use them to associate that
+    a wait event string to the associated name by calling:
+<programlisting>
+    void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name);
+</programlisting>
+    An example can be found in
+    <filename>src/test/modules/test_custom_wait_events/test_custom_wait_events.c</filename>
+    in the PostgreSQL source tree.
+    </para>
+   </sect2>
+
    <sect2 id="extend-cpp">
     <title>Using C++ for Extensibility</title>
 
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index f63c991051..56335e8730 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -133,10 +133,11 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
-		  if ( $waitclass eq 'WaitEventLWLock'
+		  if ( $waitclass eq 'WaitEventExtension'
+			|| $waitclass eq 'WaitEventLWLock'
 			|| $waitclass eq 'WaitEventLock');
 
 		my $last = $waitclass;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 4aad11c111..03acc10407 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -40,6 +43,162 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+
+/* dynamic allocation counter for custom wait events for extensions */
+typedef struct WaitEventExtensionCounter
+{
+	int nextId;		/* next ID to assign */
+	slock_t	mutex;  /* protects the counter only */
+} WaitEventExtensionCounter;
+
+/* pointer to the shared memory */
+static WaitEventExtensionCounter *waitEventExtensionCounter;
+
+/* first event ID of custom wait events for extensions */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event IDs known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionNames = NULL;
+static int	WaitEventExtensionNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(WaitEventExtensionCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	bool	found;
+
+	if (!IsUnderPostmaster)
+	{
+		/* Allocate space in shared memory. */
+		waitEventExtensionCounter = (WaitEventExtensionCounter *)
+			ShmemInitStruct("waitEventExtensionCounter", WaitEventExtensionShmemSize(), &found);
+		if (found)
+			return;
+
+		/* Initialize the dynamic-allocation counter and the spinlock. */
+		waitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&waitEventExtensionCounter->mutex);
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event.
+ */
+uint32
+WaitEventExtensionNew(void)
+{
+	uint16		eventId;
+
+	SpinLockAcquire(&waitEventExtensionCounter->mutex);
+
+	/* Check for the counter overflow. */
+	if (waitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	{
+		SpinLockRelease(&waitEventExtensionCounter->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
+	}
+
+	eventId = waitEventExtensionCounter->nextId++;
+
+	SpinLockRelease(&waitEventExtensionCounter->mutex);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic wait event name for extension in the lookup table
+ * of the current process.
+ *
+ * This routine will save a pointer to the wait event name passed as an argument,
+ * so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterName(uint32 wait_event_info,
+								const char *wait_event_name)
+{
+	uint16		eventId;
+
+	eventId = wait_event_info & 0x0000FFFF;
+
+	/* Check the wait event class. */
+	Assert((wait_event_info & 0xFF000000) == PG_WAIT_EXTENSION);
+
+	/* This should only be called for user-defined wait event. */
+	Assert(eventId >= NUM_BUILTIN_WAIT_EVENT_EXTENSION);
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionNames == NULL)
+			WaitEventExtensionNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionNames =
+				repalloc0_array(WaitEventExtensionNames, const char *, WaitEventExtensionNamesAllocated, newalloc);
+		WaitEventExtensionNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an wait event ID for extension.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Built-in event? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an user-defined wait event, so look in WaitEventExtensionNames[].
+	 * However, it's possible that the name has never been registered by
+	 * calling WaitEventExtensionRegisterName() in the current process, in
+	 * which case give up and return "extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionNamesAllocated ||
+		WaitEventExtensionNames[eventId] == NULL)
+		return "extension";
+
+	return WaitEventExtensionNames[eventId];
+}
+
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -149,6 +308,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -170,13 +332,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 3fabad96d9..2ea4789b00 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	BufferPin	"Waiting to acquire an exclusive pin on a buffer
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	Extension	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	Extension	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..ae3dac88a4 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -67,6 +67,31 @@ pgstat_report_wait_start(uint32 wait_event_info)
 	*(volatile uint32 *) my_wait_event_info = wait_event_info;
 }
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define their wait events. First, extensions should call
+ * WaitEventExtensionNew() to get one or more wait events, which IDs are
+ * allocated from a shared counter.  Next, each individual process can use
+ * them with WaitEventExtensionRegisterName() to associate that a wait
+ * event string to the associated name.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNew(void);
+extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
+								const char *wait_event_name);
+
 /* ----------
  * pgstat_report_wait_end() -
  *
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 17d369e378..83a1d8fd19 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -14,6 +14,7 @@ subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_wait_events')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..8ad073e577
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..6817eaa2b0
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,137 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+# Check that a new wait event can be defined in initialized phase.
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT start_bgworker_waiting_forever();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) = 1;]),
+	"registering custom wait event in initialized phase doen't work"
+);
+
+$node->stop;
+$node->clean_node;
+
+
+# Check that multiple new wait events can be defined dynamically.
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+# Start new background workers. They wait with new wait events dinamically defined.
+my $wait_event_info_1 = $node->safe_psql('postgres', 'SELECT get_new_wait_event_info();');
+$node->safe_psql('postgres', "SELECT start_bgworker_waiting_forever($wait_event_info_1);");
+my $wait_event_info_2 = $node->safe_psql('postgres', 'SELECT get_new_wait_event_info();');
+$node->safe_psql('postgres', "SELECT start_bgworker_waiting_forever($wait_event_info_2);");
+
+# Since the new wait event names are not registered in current process now,
+# 'wait_event' must be 'extension' which means that it's not registered.
+ok($node->poll_query_until(
+	'postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'extension'
+		) = 2;]),
+	"unregisterd wait event names are wrong"
+);
+
+# Register wait event names and check that they will be visible.
+my $new_wait_event_name_1 = 'new_dynamic_wait_event_1';
+my $new_wait_event_name_2 = 'new_dynamic_wait_event_2';
+my $session = $node->background_psql('postgres');
+$session->query_safe(
+	"SELECT register_wait_event_name($wait_event_info_1, '$new_wait_event_name_1');");
+$session->query_safe(
+	"SELECT register_wait_event_name($wait_event_info_2, '$new_wait_event_name_2');");
+my $res = $session->query_safe(
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND (wait_event = '$new_wait_event_name_1')
+		) = 1;
+]);
+like($res, qr/^t$/m, "$new_wait_event_name_1 dynamically registered wasn't found");
+$res = $session->query_safe(
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND (wait_event = '$new_wait_event_name_2')
+		) = 1;
+]);
+like($res, qr/^t$/m, "$new_wait_event_name_2 dynamically registered wasn't found");
+$session->quit;
+
+$node->stop;
+$node->clean_node;
+
+
+# Check the boundary condition. An error should happen if the number of wait
+# events for extension exceeds the limits.
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+# skip to the maximum number of wait event for extension. "$uint16_max" can
+# be allocated since the first wait event is reserved by core.
+my $uint16_max=65535;
+$node->safe_psql('postgres',
+	"SELECT COUNT(get_new_wait_event_info()) FROM generate_series(1, $uint16_max - 1)");
+
+# can use the maximum one
+my $wait_event_info = $node->safe_psql('postgres', 'SELECT get_new_wait_event_info();');
+$node->safe_psql('postgres', "SELECT start_bgworker_waiting_forever($wait_event_info);");
+ok($node->poll_query_until(
+	'postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'extension'
+		) = 1;]),
+	"unregisterd wait event name is wrong"
+);
+
+# but, next call should output an error.
+my ($ret, $stdout, $stderr) =
+	$node->psql('postgres', 'SELECT get_new_wait_event_info();');
+like(
+	$stderr,
+	qr/ERROR:  too many wait events for extensions/,
+	'expected error when the number of wait events for extensions exceeds the limits'
+);
+
+$node->stop;
+$node->clean_node;
+
+done_testing();
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..993289b1ff
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,39 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- get_new_wait_event_info()
+--
+-- Dynamically get a new wait event information.
+--
+CREATE FUNCTION get_new_wait_event_info()
+RETURNS bigint
+AS 'MODULE_PATHNAME', 'get_new_wait_event_info'
+LANGUAGE C STRICT PARALLEL SAFE;
+
+--
+-- register_wait_event_name()
+--
+-- Dynamically register a wait event name.
+--
+CREATE FUNCTION register_wait_event_name(
+    wait_event_info bigint, wait_event_name text
+)
+RETURNS void
+AS 'MODULE_PATHNAME', 'register_wait_event_name'
+LANGUAGE C STRICT PARALLEL SAFE;
+
+--
+-- start_bgworker_waiting_forever()
+--
+-- Dynamically start a background worker waiting forever
+-- with the specified custom wait event.
+--
+CREATE FUNCTION start_bgworker_waiting_forever(
+    wait_event_info bigint DEFAULT NULL
+)
+RETURNS int4
+AS 'MODULE_PATHNAME', 'start_bgworker_waiting_forever'
+LANGUAGE C PARALLEL SAFE;
\ No newline at end of file
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..7527d25f8f
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,253 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events for extensions
+ *
+ * This code has some functions to define and register new custom wait
+ * events, and to launch a background worker waiting forever with the
+ * specified wait event.
+ *
+ * You can define them not only in initialized phases but also dynamically
+ * via calling user functions. A new wait event "custom_wait_event" is
+ * defined and registered in shmem_startup_hook(). It will be saved in
+ * shared memory. Users can also define and register new custom wait events
+ * via get_new_wait_event_info() and register_wait_event_name() dynamically.
+ *
+ * You can use the wait event for a background worker launched in
+ * start_bgworker_waiting_forever(). The reason why it starts background
+ * workers is for tap tests. It enables to check the wait events on
+ * pg_stat_activity from other processes.
+ *
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/builtins.h"
+#include "utils/memutils.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(get_new_wait_event_info);
+PG_FUNCTION_INFO_V1(register_wait_event_name);
+PG_FUNCTION_INFO_V1(start_bgworker_waiting_forever);
+
+PGDLLEXPORT void custom_we_worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct CustomWaitEventState
+{
+	uint32			wait_event_info; 		 /* the wait event defined in initialized phase */
+} CustomWaitEventState;
+static CustomWaitEventState *ss = NULL;      /* the pointer to the shared memory */
+
+static char *wait_event_name;                /* the wait event name dynamically registered  */
+static char *worker_database = "postgres";   /* connected by a background worker */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void custom_we_shmem_request(void);
+static void custom_we_shmem_startup(void);
+static Size custom_we_memsize(void);
+
+void
+_PG_init(void)
+{
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = custom_we_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = custom_we_shmem_startup;
+}
+
+static void
+custom_we_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(custom_we_memsize());
+}
+
+static Size
+custom_we_memsize(void)
+{
+	return sizeof(CustomWaitEventState);
+}
+
+/*
+ * Define and register a new wait event.
+ *
+ * This routine defines a new wait event when it's called first time.
+ * After that, the name of the wait event is associated with it. It's
+ * registered in the lookup table of each processes.
+ */
+static void
+custom_we_shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(CustomWaitEventState),
+						 &found);
+
+	/* Define a new wait event. */
+	if (!found)
+		ss->wait_event_info = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(ss->wait_event_info,
+									  "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the wait event
+ * specified in first argument.
+ */
+void
+custom_we_worker_main(Datum main_arg)
+{
+	uint32 wait_event_info;
+	wait_event_info = DatumGetUInt32(main_arg);
+
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the specified wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 wait_event_info);
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically get and return a new wait event.
+ */
+Datum
+get_new_wait_event_info(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_UINT32(WaitEventExtensionNew());
+}
+
+/*
+ * Dynamically register a wait event name in the lookup table of
+ * current process.
+ */
+Datum
+register_wait_event_name(PG_FUNCTION_ARGS)
+{
+	uint32 wait_event_info;
+	char *wait_event_name_arg;
+	int size;
+
+	wait_event_info = PG_GETARG_UINT32(0);
+
+	/*
+	 * Copy to the name in a backend-lifetime context since
+	 * WaitEventExtensionRegisterName() doesn't copy the name.
+	 */
+	wait_event_name_arg = text_to_cstring(PG_GETARG_TEXT_PP(1));
+	size = strlen(wait_event_name_arg) + 1;
+	Assert(size <= NAMEDATALEN);
+	wait_event_name = (char *) MemoryContextAlloc(TopMemoryContext, size);
+	strlcpy(wait_event_name, wait_event_name_arg, size);
+
+	WaitEventExtensionRegisterName(wait_event_info, wait_event_name);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * Dynamically start a background worker waiting forever
+ * with the custom wait event specified in the first
+ * argument. If the argument is NULL, it will wait with
+ * the wait event defined in initialized phase.
+ */
+Datum
+start_bgworker_waiting_forever(PG_FUNCTION_ARGS)
+{
+	uint32 wait_event_info;
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	if (PG_ARGISNULL(0))
+	{
+		if (ss == NULL)
+			ereport(ERROR,
+					(errmsg("cannot use \"%s\" if \"%s\" has not been loaded with shared_preload_libraries",
+							"wait_worker_launch", "test_custom_wait_events")));
+
+		wait_event_info = ss->wait_event_info;
+	} else
+		wait_event_info = PG_GETARG_UINT32(0);
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "custom_we_worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_main_arg = wait_event_info;
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.25.1

From e5d63913cbba9f20098252cc1e415eda4d32f5ad Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Wed, 19 Jul 2023 12:28:12 +0900
Subject: [PATCH] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to define new wait events
for extensions dynamically.

The APIs are
* WaitEventExtensionNew()
* WaitEventExtensionRegisterName()

These are similar to the existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extensions should call WaitEventExtensionNew() to get one
or more new wait event, which IDs are allocated from a shared
counter.  Next, each individual process can use the wait event with
WaitEventExtensionRegisterName() to associate that a wait event
string to the associated name.

Note that this includes a test module, which perhaps should be split
later into its own commit.
---
 doc/src/sgml/monitoring.sgml                  |  10 +-
 doc/src/sgml/xfunc.sgml                       |  23 ++
 src/backend/storage/ipc/ipci.c                |   3 +
 .../activity/generate-wait_event_types.pl     |   7 +-
 src/backend/utils/activity/wait_event.c       | 171 +++++++++++-
 .../utils/activity/wait_event_names.txt       |   2 +-
 src/include/utils/wait_event.h                |  25 ++
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 .../test_custom_wait_events/.gitignore        |   4 +
 .../modules/test_custom_wait_events/Makefile  |  23 ++
 .../test_custom_wait_events/meson.build       |  33 +++
 .../test_custom_wait_events/t/001_basic.pl    | 137 ++++++++++
 .../test_custom_wait_events--1.0.sql          |  39 +++
 .../test_custom_wait_events.c                 | 253 ++++++++++++++++++
 .../test_custom_wait_events.control           |   3 +
 16 files changed, 718 insertions(+), 17 deletions(-)
 create mode 100644 src/test/modules/test_custom_wait_events/.gitignore
 create mode 100644 src/test/modules/test_custom_wait_events/Makefile
 create mode 100644 src/test/modules/test_custom_wait_events/meson.build
 create mode 100644 src/test/modules/test_custom_wait_events/t/001_basic.pl
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.c
 create mode 100644 src/test/modules/test_custom_wait_events/test_custom_wait_events.control

diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index 588b720f57..3aee5c243f 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1117,12 +1117,12 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
 
    <note>
     <para>
-     Extensions can add <literal>LWLock</literal> types to the list shown in
-     <xref linkend="wait-event-lwlock-table"/>.  In some cases, the name
+     Extensions can add <literal>Extension</literal> and <literal>LWLock</literal> types
+     to the list shown in <xref linkend="wait-event-extension-table"/> and
+     <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
      assigned by an extension will not be available in all server processes;
-     so an <literal>LWLock</literal> wait event might be reported as
-     just <quote><literal>extension</literal></quote> rather than the
-     extension-assigned name.
+     so an wait event might be reported as just <quote><literal>extension</literal></quote>
+     rather than the extension-assigned name.
     </para>
    </note>
  </sect2>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9620ea9ae3..bc210cd03b 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3453,6 +3453,29 @@ if (!ptr)
     </para>
    </sect2>
 
+   <sect2 id="xfunc-addin-wait-events">
+    <title>Custom Wait Events for Add-ins</title>
+
+    <para>
+    Add-ins can define custom wait events that the wait event type is
+    <literal>Extension</literal>.
+    </para>
+    <para>
+    First, add-ins should get new one or more wait events by calling:
+<programlisting>
+    uint32 WaitEventExtensionNew(void)
+</programlisting>
+    Next, each individual process can use them to associate that
+    a wait event string to the associated name by calling:
+<programlisting>
+    void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name);
+</programlisting>
+    An example can be found in
+    <filename>src/test/modules/test_custom_wait_events/test_custom_wait_events.c</filename>
+    in the PostgreSQL source tree.
+    </para>
+   </sect2>
+
    <sect2 id="extend-cpp">
     <title>Using C++ for Extensibility</title>
 
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index f63c991051..56335e8730 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -133,10 +133,11 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
-		  if ( $waitclass eq 'WaitEventLWLock'
+		  if ( $waitclass eq 'WaitEventExtension'
+			|| $waitclass eq 'WaitEventLWLock'
 			|| $waitclass eq 'WaitEventLock');
 
 		my $last = $waitclass;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 4aad11c111..03acc10407 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -40,6 +43,162 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+
+/* dynamic allocation counter for custom wait events for extensions */
+typedef struct WaitEventExtensionCounter
+{
+	int nextId;		/* next ID to assign */
+	slock_t	mutex;  /* protects the counter only */
+} WaitEventExtensionCounter;
+
+/* pointer to the shared memory */
+static WaitEventExtensionCounter *waitEventExtensionCounter;
+
+/* first event ID of custom wait events for extensions */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event IDs known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionNames = NULL;
+static int	WaitEventExtensionNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(WaitEventExtensionCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	bool	found;
+
+	if (!IsUnderPostmaster)
+	{
+		/* Allocate space in shared memory. */
+		waitEventExtensionCounter = (WaitEventExtensionCounter *)
+			ShmemInitStruct("waitEventExtensionCounter", WaitEventExtensionShmemSize(), &found);
+		if (found)
+			return;
+
+		/* Initialize the dynamic-allocation counter and the spinlock. */
+		waitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&waitEventExtensionCounter->mutex);
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event.
+ */
+uint32
+WaitEventExtensionNew(void)
+{
+	uint16		eventId;
+
+	SpinLockAcquire(&waitEventExtensionCounter->mutex);
+
+	/* Check for the counter overflow. */
+	if (waitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	{
+		SpinLockRelease(&waitEventExtensionCounter->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
+	}
+
+	eventId = waitEventExtensionCounter->nextId++;
+
+	SpinLockRelease(&waitEventExtensionCounter->mutex);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic wait event name for extension in the lookup table
+ * of the current process.
+ *
+ * This routine will save a pointer to the wait event name passed as an argument,
+ * so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterName(uint32 wait_event_info,
+								const char *wait_event_name)
+{
+	uint16		eventId;
+
+	eventId = wait_event_info & 0x0000FFFF;
+
+	/* Check the wait event class. */
+	Assert((wait_event_info & 0xFF000000) == PG_WAIT_EXTENSION);
+
+	/* This should only be called for user-defined wait event. */
+	Assert(eventId >= NUM_BUILTIN_WAIT_EVENT_EXTENSION);
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionNames == NULL)
+			WaitEventExtensionNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionNames =
+				repalloc0_array(WaitEventExtensionNames, const char *, WaitEventExtensionNamesAllocated, newalloc);
+		WaitEventExtensionNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an wait event ID for extension.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Built-in event? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an user-defined wait event, so look in WaitEventExtensionNames[].
+	 * However, it's possible that the name has never been registered by
+	 * calling WaitEventExtensionRegisterName() in the current process, in
+	 * which case give up and return "extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionNamesAllocated ||
+		WaitEventExtensionNames[eventId] == NULL)
+		return "extension";
+
+	return WaitEventExtensionNames[eventId];
+}
+
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -149,6 +308,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -170,13 +332,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 3fabad96d9..2ea4789b00 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	BufferPin	"Waiting to acquire an exclusive pin on a buffer
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	Extension	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	Extension	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..ae3dac88a4 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -67,6 +67,31 @@ pgstat_report_wait_start(uint32 wait_event_info)
 	*(volatile uint32 *) my_wait_event_info = wait_event_info;
 }
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define their wait events. First, extensions should call
+ * WaitEventExtensionNew() to get one or more wait events, which IDs are
+ * allocated from a shared counter.  Next, each individual process can use
+ * them with WaitEventExtensionRegisterName() to associate that a wait
+ * event string to the associated name.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNew(void);
+extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
+								const char *wait_event_name);
+
 /* ----------
  * pgstat_report_wait_end() -
  *
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..84312b7e57 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -17,6 +17,7 @@ SUBDIRS = \
 		  test_bloomfilter \
 		  test_copy_callbacks \
 		  test_custom_rmgrs \
+		  test_custom_wait_events \
 		  test_ddl_deparse \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 17d369e378..83a1d8fd19 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -14,6 +14,7 @@ subdir('ssl_passphrase_callback')
 subdir('test_bloomfilter')
 subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
+subdir('test_custom_wait_events')
 subdir('test_ddl_deparse')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_custom_wait_events/.gitignore b/src/test/modules/test_custom_wait_events/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_custom_wait_events/Makefile b/src/test/modules/test_custom_wait_events/Makefile
new file mode 100644
index 0000000000..dda365d523
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_custom_wait_events/Makefile
+
+MODULE_big = test_custom_wait_events
+OBJS = \
+	$(WIN32RES) \
+	test_custom_wait_events.o
+PGFILEDESC = "test_custom_wait_events - test custom wait events"
+
+EXTENSION = test_custom_wait_events
+DATA = test_custom_wait_events--1.0.sql
+
+TAP_TESTS = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_wait_events
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_wait_events/meson.build b/src/test/modules/test_custom_wait_events/meson.build
new file mode 100644
index 0000000000..8ad073e577
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) PostgreSQL Global Development Group
+
+test_custom_wait_events = files(
+  'test_custom_wait_events.c',
+)
+
+if host_system == 'windows'
+  test_custom_wait_events += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_custom_wait_events',
+    '--FILEDESC', 'test_custom_wait_events - test custom wait events',])
+endif
+
+test_custom_wait_events = shared_module('test_custom_wait_events',
+  test_custom_wait_events,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_custom_wait_events
+
+test_install_data += files(
+  'test_custom_wait_events.control',
+  'test_custom_wait_events--1.0.sql',
+)
+
+tests += {
+  'name': 'test_custom_wait_events',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'tests': [
+      't/001_basic.pl',
+    ],
+  },
+}
diff --git a/src/test/modules/test_custom_wait_events/t/001_basic.pl b/src/test/modules/test_custom_wait_events/t/001_basic.pl
new file mode 100644
index 0000000000..6817eaa2b0
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/t/001_basic.pl
@@ -0,0 +1,137 @@
+# Copyright (c) 2023, PostgreSQL Global Development Group
+
+use strict;
+use warnings;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+# Check that a new wait event can be defined in initialized phase.
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+	"shared_preload_libraries = 'test_custom_wait_events'"
+);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+$node->safe_psql('postgres', 'SELECT start_bgworker_waiting_forever();');
+
+ok($node->poll_query_until('postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'custom_wait_event'
+		) = 1;]),
+	"registering custom wait event in initialized phase doen't work"
+);
+
+$node->stop;
+$node->clean_node;
+
+
+# Check that multiple new wait events can be defined dynamically.
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+# Start new background workers. They wait with new wait events dinamically defined.
+my $wait_event_info_1 = $node->safe_psql('postgres', 'SELECT get_new_wait_event_info();');
+$node->safe_psql('postgres', "SELECT start_bgworker_waiting_forever($wait_event_info_1);");
+my $wait_event_info_2 = $node->safe_psql('postgres', 'SELECT get_new_wait_event_info();');
+$node->safe_psql('postgres', "SELECT start_bgworker_waiting_forever($wait_event_info_2);");
+
+# Since the new wait event names are not registered in current process now,
+# 'wait_event' must be 'extension' which means that it's not registered.
+ok($node->poll_query_until(
+	'postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'extension'
+		) = 2;]),
+	"unregisterd wait event names are wrong"
+);
+
+# Register wait event names and check that they will be visible.
+my $new_wait_event_name_1 = 'new_dynamic_wait_event_1';
+my $new_wait_event_name_2 = 'new_dynamic_wait_event_2';
+my $session = $node->background_psql('postgres');
+$session->query_safe(
+	"SELECT register_wait_event_name($wait_event_info_1, '$new_wait_event_name_1');");
+$session->query_safe(
+	"SELECT register_wait_event_name($wait_event_info_2, '$new_wait_event_name_2');");
+my $res = $session->query_safe(
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND (wait_event = '$new_wait_event_name_1')
+		) = 1;
+]);
+like($res, qr/^t$/m, "$new_wait_event_name_1 dynamically registered wasn't found");
+$res = $session->query_safe(
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND (wait_event = '$new_wait_event_name_2')
+		) = 1;
+]);
+like($res, qr/^t$/m, "$new_wait_event_name_2 dynamically registered wasn't found");
+$session->quit;
+
+$node->stop;
+$node->clean_node;
+
+
+# Check the boundary condition. An error should happen if the number of wait
+# events for extension exceeds the limits.
+$node->init;
+$node->append_conf(
+	'postgresql.conf',
+);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_custom_wait_events;');
+
+# skip to the maximum number of wait event for extension. "$uint16_max" can
+# be allocated since the first wait event is reserved by core.
+my $uint16_max=65535;
+$node->safe_psql('postgres',
+	"SELECT COUNT(get_new_wait_event_info()) FROM generate_series(1, $uint16_max - 1)");
+
+# can use the maximum one
+my $wait_event_info = $node->safe_psql('postgres', 'SELECT get_new_wait_event_info();');
+$node->safe_psql('postgres', "SELECT start_bgworker_waiting_forever($wait_event_info);");
+ok($node->poll_query_until(
+	'postgres',
+	qq[SELECT
+		(SELECT count(*) FROM pg_stat_activity
+			WHERE backend_type = 'test_custom_wait_events'
+			  AND wait_event_type = 'Extension'
+			  AND wait_event = 'extension'
+		) = 1;]),
+	"unregisterd wait event name is wrong"
+);
+
+# but, next call should output an error.
+my ($ret, $stdout, $stderr) =
+	$node->psql('postgres', 'SELECT get_new_wait_event_info();');
+like(
+	$stderr,
+	qr/ERROR:  too many wait events for extensions/,
+	'expected error when the number of wait events for extensions exceeds the limits'
+);
+
+$node->stop;
+$node->clean_node;
+
+done_testing();
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
new file mode 100644
index 0000000000..df916d4d5e
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql
@@ -0,0 +1,39 @@
+/* src/test/modules/test_custom_wait_events/test_custom_wait_events--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_custom_wait_events" to load this file. \quit
+
+--
+-- get_new_wait_event_info()
+--
+-- Dynamically get a new wait event information.
+--
+CREATE FUNCTION get_new_wait_event_info()
+RETURNS int4
+AS 'MODULE_PATHNAME', 'get_new_wait_event_info'
+LANGUAGE C STRICT PARALLEL SAFE;
+
+--
+-- register_wait_event_name()
+--
+-- Dynamically register a wait event name.
+--
+CREATE FUNCTION register_wait_event_name(
+    wait_event_info int4, wait_event_name text
+)
+RETURNS void
+AS 'MODULE_PATHNAME', 'register_wait_event_name'
+LANGUAGE C STRICT PARALLEL SAFE;
+
+--
+-- start_bgworker_waiting_forever()
+--
+-- Dynamically start a background worker waiting forever
+-- with the specified custom wait event.
+--
+CREATE FUNCTION start_bgworker_waiting_forever(
+    wait_event_info int4 DEFAULT NULL
+)
+RETURNS int4
+AS 'MODULE_PATHNAME', 'start_bgworker_waiting_forever'
+LANGUAGE C PARALLEL SAFE;
\ No newline at end of file
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.c b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
new file mode 100644
index 0000000000..7527d25f8f
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.c
@@ -0,0 +1,253 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_custom_wait_events.c
+ *		Code for testing custom wait events for extensions
+ *
+ * This code has some functions to define and register new custom wait
+ * events, and to launch a background worker waiting forever with the
+ * specified wait event.
+ *
+ * You can define them not only in initialized phases but also dynamically
+ * via calling user functions. A new wait event "custom_wait_event" is
+ * defined and registered in shmem_startup_hook(). It will be saved in
+ * shared memory. Users can also define and register new custom wait events
+ * via get_new_wait_event_info() and register_wait_event_name() dynamically.
+ *
+ * You can use the wait event for a background worker launched in
+ * start_bgworker_waiting_forever(). The reason why it starts background
+ * workers is for tap tests. It enables to check the wait events on
+ * pg_stat_activity from other processes.
+ *
+ * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_custom_wait_events/test_custom_wait_events.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
+#include "storage/ipc.h"
+#include "storage/latch.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
+#include "utils/builtins.h"
+#include "utils/memutils.h"
+#include "utils/wait_event.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(get_new_wait_event_info);
+PG_FUNCTION_INFO_V1(register_wait_event_name);
+PG_FUNCTION_INFO_V1(start_bgworker_waiting_forever);
+
+PGDLLEXPORT void custom_we_worker_main(Datum main_arg);
+
+/* GUC variables */
+typedef struct CustomWaitEventState
+{
+	uint32			wait_event_info; 		 /* the wait event defined in initialized phase */
+} CustomWaitEventState;
+static CustomWaitEventState *ss = NULL;      /* the pointer to the shared memory */
+
+static char *wait_event_name;                /* the wait event name dynamically registered  */
+static char *worker_database = "postgres";   /* connected by a background worker */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void custom_we_shmem_request(void);
+static void custom_we_shmem_startup(void);
+static Size custom_we_memsize(void);
+
+void
+_PG_init(void)
+{
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = custom_we_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = custom_we_shmem_startup;
+}
+
+static void
+custom_we_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(custom_we_memsize());
+}
+
+static Size
+custom_we_memsize(void)
+{
+	return sizeof(CustomWaitEventState);
+}
+
+/*
+ * Define and register a new wait event.
+ *
+ * This routine defines a new wait event when it's called first time.
+ * After that, the name of the wait event is associated with it. It's
+ * registered in the lookup table of each processes.
+ */
+static void
+custom_we_shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	ss = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	ss = ShmemInitStruct("custom_wait_event",
+						 sizeof(CustomWaitEventState),
+						 &found);
+
+	/* Define a new wait event. */
+	if (!found)
+		ss->wait_event_info = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(ss->wait_event_info,
+									  "custom_wait_event");
+
+	return;
+}
+
+/*
+ * Connect to the database and wait forever with the wait event
+ * specified in first argument.
+ */
+void
+custom_we_worker_main(Datum main_arg)
+{
+	uint32 wait_event_info;
+	wait_event_info = DatumGetUInt32(main_arg);
+
+	pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
+	BackgroundWorkerUnblockSignals();
+
+	/* Connect to our database to show up a wait event in pg_stat_activity */
+	BackgroundWorkerInitializeConnection(worker_database, NULL, 0);
+
+	/* Wait forever with the specified wait event */
+	while (!ShutdownRequestPending)
+	{
+		(void) WaitLatch(MyLatch,
+						 WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+						 -1L,
+						 wait_event_info);
+		ResetLatch(MyLatch);
+	}
+}
+
+/*
+ * Dynamically get and return a new wait event.
+ */
+Datum
+get_new_wait_event_info(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_UINT32(WaitEventExtensionNew());
+}
+
+/*
+ * Dynamically register a wait event name in the lookup table of
+ * current process.
+ */
+Datum
+register_wait_event_name(PG_FUNCTION_ARGS)
+{
+	uint32 wait_event_info;
+	char *wait_event_name_arg;
+	int size;
+
+	wait_event_info = PG_GETARG_UINT32(0);
+
+	/*
+	 * Copy to the name in a backend-lifetime context since
+	 * WaitEventExtensionRegisterName() doesn't copy the name.
+	 */
+	wait_event_name_arg = text_to_cstring(PG_GETARG_TEXT_PP(1));
+	size = strlen(wait_event_name_arg) + 1;
+	Assert(size <= NAMEDATALEN);
+	wait_event_name = (char *) MemoryContextAlloc(TopMemoryContext, size);
+	strlcpy(wait_event_name, wait_event_name_arg, size);
+
+	WaitEventExtensionRegisterName(wait_event_info, wait_event_name);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * Dynamically start a background worker waiting forever
+ * with the custom wait event specified in the first
+ * argument. If the argument is NULL, it will wait with
+ * the wait event defined in initialized phase.
+ */
+Datum
+start_bgworker_waiting_forever(PG_FUNCTION_ARGS)
+{
+	uint32 wait_event_info;
+	BackgroundWorker worker;
+	BackgroundWorkerHandle *handle;
+	BgwHandleStatus status;
+	pid_t		pid;
+
+	if (PG_ARGISNULL(0))
+	{
+		if (ss == NULL)
+			ereport(ERROR,
+					(errmsg("cannot use \"%s\" if \"%s\" has not been loaded with shared_preload_libraries",
+							"wait_worker_launch", "test_custom_wait_events")));
+
+		wait_event_info = ss->wait_event_info;
+	} else
+		wait_event_info = PG_GETARG_UINT32(0);
+
+	memset(&worker, 0, sizeof(worker));
+	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+		BGWORKER_BACKEND_DATABASE_CONNECTION;
+	worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+	worker.bgw_restart_time = BGW_NEVER_RESTART;
+	sprintf(worker.bgw_function_name, "custom_we_worker_main");
+	sprintf(worker.bgw_library_name, "test_custom_wait_events");
+	snprintf(worker.bgw_name, BGW_MAXLEN, "test_custom_wait_events worker");
+	snprintf(worker.bgw_type, BGW_MAXLEN, "test_custom_wait_events");
+	worker.bgw_main_arg = wait_event_info;
+	worker.bgw_notify_pid = MyProcPid;
+
+	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
+		PG_RETURN_NULL();
+
+	status = WaitForBackgroundWorkerStartup(handle, &pid);
+
+	if (status == BGWH_STOPPED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("could not start background process"),
+				 errhint("More details may be available in the server log.")));
+	if (status == BGWH_POSTMASTER_DIED)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
+				 errmsg("cannot start background processes without postmaster"),
+				 errhint("Kill all remaining database processes and restart the database.")));
+	Assert(status == BGWH_STARTED);
+
+	PG_RETURN_INT32(pid);
+}
diff --git a/src/test/modules/test_custom_wait_events/test_custom_wait_events.control b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
new file mode 100644
index 0000000000..8dd875a252
--- /dev/null
+++ b/src/test/modules/test_custom_wait_events/test_custom_wait_events.control
@@ -0,0 +1,3 @@
+comment = 'Test code for custom wait events'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_wait_events'
-- 
2.25.1

From deb7cab66671e0277cc34236c7515f4ea18fac65 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda.us@hco.ntt.co.jp>
Date: Wed, 19 Jul 2023 12:28:12 +0900
Subject: [PATCH v8] Support custom wait events for extensions.

To support custom wait events, it add 2 APIs to define new wait events
for extensions dynamically.

The APIs are
* WaitEventExtensionNew()
* WaitEventExtensionRegisterName()

These are similar to the existing LWLockNewTrancheId() and
LWLockRegisterTranche().

First, extensions should call WaitEventExtensionNew() to get one
or more new wait event, which IDs are allocated from a shared
counter.  Next, each individual process can use the wait event with
WaitEventExtensionRegisterName() to associate that a wait event
string to the associated name.

Note that this includes an example of how to use this new facility in
worker_spi.
---
 src/include/utils/wait_event.h                |  25 +++
 src/backend/storage/ipc/ipci.c                |   3 +
 .../activity/generate-wait_event_types.pl     |   7 +-
 src/backend/utils/activity/wait_event.c       | 171 +++++++++++++++++-
 .../utils/activity/wait_event_names.txt       |   2 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  31 +++-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 +
 src/test/modules/worker_spi/worker_spi.c      | 105 ++++++++++-
 doc/src/sgml/monitoring.sgml                  |  13 +-
 doc/src/sgml/xfunc.sgml                       |  34 ++++
 10 files changed, 370 insertions(+), 26 deletions(-)

diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..d405398b75 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -67,6 +67,31 @@ pgstat_report_wait_start(uint32 wait_event_info)
 	*(volatile uint32 *) my_wait_event_info = wait_event_info;
 }
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define their wait events. First, extensions should call
+ * WaitEventExtensionNew() to get one or more wait events, which IDs are
+ * allocated from a shared counter.  Next, each individual process can use
+ * them with WaitEventExtensionRegisterName() to associate that a wait
+ * event string to the associated name.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNew(void);
+extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
+										   const char *wait_event_name);
+
 /* ----------
  * pgstat_report_wait_end() -
  *
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index f63c991051..56335e8730 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -133,10 +133,11 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
-		  if ( $waitclass eq 'WaitEventLWLock'
+		  if ( $waitclass eq 'WaitEventExtension'
+			|| $waitclass eq 'WaitEventLWLock'
 			|| $waitclass eq 'WaitEventLock');
 
 		my $last = $waitclass;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 4aad11c111..9e7e65bb01 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -40,6 +43,162 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+
+/* dynamic allocation counter for custom wait events for extensions */
+typedef struct WaitEventExtensionCounter
+{
+	int			nextId;			/* next ID to assign */
+	slock_t		mutex;			/* protects the counter only */
+}			WaitEventExtensionCounter;
+
+/* pointer to the shared memory */
+static WaitEventExtensionCounter * waitEventExtensionCounter;
+
+/* first event ID of custom wait events for extensions */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event IDs known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionNames = NULL;
+static int	WaitEventExtensionNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(WaitEventExtensionCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	bool		found;
+
+	if (!IsUnderPostmaster)
+	{
+		/* Allocate space in shared memory. */
+		waitEventExtensionCounter = (WaitEventExtensionCounter *)
+			ShmemInitStruct("waitEventExtensionCounter", WaitEventExtensionShmemSize(), &found);
+		if (found)
+			return;
+
+		/* Initialize the dynamic-allocation counter and the spinlock. */
+		waitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&waitEventExtensionCounter->mutex);
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event.
+ */
+uint32
+WaitEventExtensionNew(void)
+{
+	uint16		eventId;
+
+	SpinLockAcquire(&waitEventExtensionCounter->mutex);
+
+	/* Check for the counter overflow. */
+	if (waitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	{
+		SpinLockRelease(&waitEventExtensionCounter->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
+	}
+
+	eventId = waitEventExtensionCounter->nextId++;
+
+	SpinLockRelease(&waitEventExtensionCounter->mutex);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic wait event name for extension in the lookup table
+ * of the current process.
+ *
+ * This routine will save a pointer to the wait event name passed as an argument,
+ * so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterName(uint32 wait_event_info,
+							   const char *wait_event_name)
+{
+	uint16		eventId;
+
+	eventId = wait_event_info & 0x0000FFFF;
+
+	/* Check the wait event class. */
+	Assert((wait_event_info & 0xFF000000) == PG_WAIT_EXTENSION);
+
+	/* This should only be called for user-defined wait event. */
+	Assert(eventId >= NUM_BUILTIN_WAIT_EVENT_EXTENSION);
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionNamesAllocated)
+	{
+		int			newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionNames == NULL)
+			WaitEventExtensionNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionNames =
+				repalloc0_array(WaitEventExtensionNames, const char *, WaitEventExtensionNamesAllocated, newalloc);
+		WaitEventExtensionNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an wait event ID for extension.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Built-in event? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It's an user-defined wait event, so look in WaitEventExtensionNames[].
+	 * However, it's possible that the name has never been registered by
+	 * calling WaitEventExtensionRegisterName() in the current process, in
+	 * which case give up and return "extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionNamesAllocated ||
+		WaitEventExtensionNames[eventId] == NULL)
+		return "extension";
+
+	return WaitEventExtensionNames[eventId];
+}
+
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -149,6 +308,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -170,13 +332,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 3fabad96d9..2ea4789b00 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	BufferPin	"Waiting to acquire an exclusive pin on a buffer
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	Extension	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	Extension	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index c293871313..f914791a64 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,6 +39,25 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
+# Check that the wait event used by the dynamic bgworker.  For a session
+# not without the state in shared memory known, the default of "Extension"
+# is the value waited on.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	'extension');
+is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
+# If the shared memory state is loaded, the wait event is a custom one.
+# The expected result is a special pattern here with a newline being first
+# as the initialization of the shared memory state is enforced.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	qq[
+worker_spi_main]);
+is($result, 1,
+	'dynamic bgworker has reported "worker_spi_main" as wait event');
+
 note "testing bgworkers loaded with shared_preload_libraries";
 
 # Create the database first so as the workers can connect to it when
@@ -58,9 +77,9 @@ $node->restart;
 # Check that bgworkers have been registered and launched.
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname) FROM pg_stat_activity
-            WHERE backend_type = 'worker_spi' GROUP BY datname;],
-		'mydb|3'),
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
+            WHERE backend_type = 'worker_spi' GROUP BY datname, wait_event;],
+		'mydb|3|worker_spi_main'),
 	'bgworkers all launched'
 ) or die "Timed out while waiting for bgworkers to be launched";
 
@@ -70,10 +89,10 @@ my $worker1_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(1);');
 my $worker2_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(2);');
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname)  FROM pg_stat_activity
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
             WHERE backend_type = 'worker_spi dynamic' AND
-            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname;],
-		'mydb|2'),
+            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, wait_event;],
+		'mydb|2|worker_spi_main'),
 	'dynamic bgworkers all launched'
 ) or die "Timed out while waiting for dynamic bgworkers to be launched";
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index e9d5b07373..f13f7e0f98 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,3 +7,8 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
+
+CREATE FUNCTION worker_spi_init()
+RETURNS VOID STRICT
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 903dcddef9..48b5aad022 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,22 +44,97 @@
 
 PG_MODULE_MAGIC;
 
+PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
+/* Shared memory state */
+typedef struct worker_spi_state
+{
+	/* the wait event defined in initialized phase */
+	uint32		wait_event;
+}			worker_spi_state;
+
+static worker_spi_state * wsstate = NULL;	/* pointer to shared memory */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void worker_spi_shmem_request(void);
+static void worker_spi_shmem_startup(void);
+static void worker_spi_shmem_init(void);
+static Size worker_spi_memsize(void);
+
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database = NULL;
 
-
 typedef struct worktable
 {
 	const char *schema;
 	const char *name;
 } worktable;
 
+static void
+worker_spi_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(worker_spi_memsize());
+}
+
+static void
+worker_spi_shmem_startup(void)
+{
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	worker_spi_shmem_init();
+}
+
+static Size
+worker_spi_memsize(void)
+{
+	return MAXALIGN(sizeof(worker_spi_state));
+}
+
+/*
+ * Initialize the shared memory state of worker_spi.
+ *
+ * This routine allocates a new wait event when called the first time.
+ * On follow-up calls, the name of the wait event associated with the
+ * existing shared memory state is registered.
+ */
+static void
+worker_spi_shmem_init(void)
+{
+	bool		found;
+
+	wsstate = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	wsstate = ShmemInitStruct("custom_wait_event",
+							  sizeof(worker_spi_state),
+							  &found);
+
+	/* Define a new wait event */
+	if (!found)
+		wsstate->wait_event = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
+
+	return;
+}
+
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -149,6 +224,9 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -199,7 +277,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 WAIT_EVENT_EXTENSION);
+						 wsstate->wait_event);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -328,6 +406,11 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = worker_spi_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = worker_spi_shmem_startup;
+
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -351,6 +434,21 @@ _PG_init(void)
 	}
 }
 
+/*
+ * Wrapper to initialize a session with the shared memory state
+ * used by this module.  This is a convenience routine to be able to
+ * see the custom wait event stored in shared memory without loading
+ * through shared_preload_libraries.
+ */
+Datum
+worker_spi_init(PG_FUNCTION_ARGS)
+{
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
+	PG_RETURN_VOID();
+}
+
 /*
  * Dynamically launch an SPI worker.
  */
@@ -363,6 +461,9 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index 588b720f57..f8c0220d47 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1117,12 +1117,13 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
 
    <note>
     <para>
-     Extensions can add <literal>LWLock</literal> types to the list shown in
-     <xref linkend="wait-event-lwlock-table"/>.  In some cases, the name
-     assigned by an extension will not be available in all server processes;
-     so an <literal>LWLock</literal> wait event might be reported as
-     just <quote><literal>extension</literal></quote> rather than the
-     extension-assigned name.
+     Extensions can add <literal>Extension</literal> and <literal>LWLock</literal> types
+     to the list shown in <xref linkend="wait-event-extension-table"/> and
+     <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
+     assigned by an extension will not be available in all server processes
+     if the extension's library is not loaded; so a custom wait event might
+     be reported as just <quote><literal>extension</literal></quote>
+     rather than the assigned, custom, name.
     </para>
    </note>
  </sect2>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9620ea9ae3..97344b219a 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3453,6 +3453,40 @@ if (!ptr)
     </para>
    </sect2>
 
+   <sect2 id="xfunc-addin-wait-events">
+    <title>Shared Memory and Custom Wait Events</title>
+
+    <para>
+     Add-ins can define Custom Wait Events under the wait event type
+     <literal>Extension</literal>. The add-in's shared library must be
+     preloaded by specifying it in <literal>shared_preload_libraries</literal>,
+     and register a <literal>shmem_request_hook</literal> and a
+     <literal>shmem_startup_hook</literal> in its
+     <function>_PG_init</function> function.
+     <literal>shmem_request_hook</literal> can request a shared memory size
+     to be later used at startup by calling from the
+     <literal>shmem_request_hook</literal>:
+<programlisting>
+void RequestAddinShmemSpace(int size)
+</programlisting>
+    </para>
+    <para>
+     Custom Wait Events are allocated in shared memory by calling:
+<programlisting>
+uint32 WaitEventExtensionNew(void)
+</programlisting>
+     in the <literal>shmem_startup_hook</literal> while holding the LWLock
+     <function>AddinShmemInitLock</function> to avoid any race conditions.
+     Next, each process needs to associate the wait event allocated previously
+     to a string, which is something done by calling:
+<programlisting>
+void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name);
+</programlisting>
+     An example can be found in <filename>src/test/modules/worker_spi</filename>
+     in the PostgreSQL source tree.
+    </para>
+   </sect2>
+
    <sect2 id="extend-cpp">
     <title>Using C++ for Extensibility</title>
 
-- 
2.40.1

From 766d1db13035eefd87a02adb88853f575abe75cb Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 28 Jul 2023 10:01:43 +0900
Subject: [PATCH v9] Support custom wait events for wait event type "Extension"

Two backend routines are added to allow extension to allocate and define
custom wait events, all of these being allocated as of type "Extension":
* WaitEventExtensionNew(), that allocates a wait event ID computed from
a counter in shared memory.
* WaitEventExtensionRegisterName(), to associate a custom string to the
wait event registered.

Note that this includes an example of how to use this new facility in
worker_spi, with tests for various scenarios.

The use of these routines should be extended across more in-core
modules, but this is left as work for future patches.
---
 src/include/utils/wait_event.h                |  26 +++
 src/backend/storage/ipc/ipci.c                |   3 +
 .../activity/generate-wait_event_types.pl     |   7 +-
 src/backend/utils/activity/wait_event.c       | 171 +++++++++++++++++-
 .../utils/activity/wait_event_names.txt       |   2 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  31 +++-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 +
 src/test/modules/worker_spi/worker_spi.c      | 104 ++++++++++-
 doc/src/sgml/monitoring.sgml                  |  13 +-
 doc/src/sgml/xfunc.sgml                       |  33 ++++
 10 files changed, 369 insertions(+), 26 deletions(-)

diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..3e38215c2d 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -38,6 +38,32 @@ extern void pgstat_reset_wait_event_storage(void);
 extern PGDLLIMPORT uint32 *my_wait_event_info;
 
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define their own wait events in this categiry.  First,
+ * they should call WaitEventExtensionNew() to get one or more wait event
+ * IDs that are allocated from a shared counter.  These can be used directly
+ * with pgstat_report_wait_start() or equivalent.  Next, each individual
+ * process should call WaitEventExtensionRegisterName() to associate a wait
+ * event string to the number allocated previously.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNew(void);
+extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
+										   const char *wait_event_name);
+
 /* ----------
  * pgstat_report_wait_start() -
  *
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index f63c991051..56335e8730 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -133,10 +133,11 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
-		  if ( $waitclass eq 'WaitEventLWLock'
+		  if ( $waitclass eq 'WaitEventExtension'
+			|| $waitclass eq 'WaitEventLWLock'
 			|| $waitclass eq 'WaitEventLock');
 
 		my $last = $waitclass;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 4aad11c111..7754ba0e07 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -40,6 +43,162 @@ static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
 
+
+/* dynamic allocation counter for custom wait events in extensions */
+typedef struct WaitEventExtensionCounter
+{
+	int			nextId;			/* next ID to assign */
+	slock_t		mutex;			/* protects the counter */
+} WaitEventExtensionCounter;
+
+/* pointer to the shared memory */
+static WaitEventExtensionCounter *waitEventExtensionCounter;
+
+/* first event ID of custom wait events for extensions */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event IDs known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionNames = NULL;
+static int	WaitEventExtensionNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(WaitEventExtensionCounter);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	bool		found;
+
+	waitEventExtensionCounter = (WaitEventExtensionCounter *)
+		ShmemInitStruct("waitEventExtensionCounter",
+						WaitEventExtensionShmemSize(), &found);
+
+	if (!found)
+	{
+		/* initialize the allocation counter and its spinlock. */
+		waitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&waitEventExtensionCounter->mutex);
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event.
+ */
+uint32
+WaitEventExtensionNew(void)
+{
+	uint16		eventId;
+
+	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+
+	SpinLockAcquire(&waitEventExtensionCounter->mutex);
+
+	if (waitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	{
+		SpinLockRelease(&waitEventExtensionCounter->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
+	}
+
+	eventId = waitEventExtensionCounter->nextId++;
+
+	SpinLockRelease(&waitEventExtensionCounter->mutex);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic wait event name for extension in the lookup table
+ * of the current process.
+ *
+ * This routine will save a pointer to the wait event name passed as an argument,
+ * so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterName(uint32 wait_event_info,
+							   const char *wait_event_name)
+{
+	uint16		eventId;
+
+	eventId = wait_event_info & 0x0000FFFF;
+
+	/* Check the wait event class. */
+	Assert((wait_event_info & 0xFF000000) == PG_WAIT_EXTENSION);
+
+	/* This should only be called for user-defined wait event. */
+	Assert(eventId >= NUM_BUILTIN_WAIT_EVENT_EXTENSION);
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionNamesAllocated)
+	{
+		uint32		newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionNames == NULL)
+			WaitEventExtensionNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionNames =
+				repalloc0_array(WaitEventExtensionNames, const char *,
+								WaitEventExtensionNamesAllocated, newalloc);
+		WaitEventExtensionNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an wait event ID for extension.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Built-in event? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
+	 * However, it is possible that the name has never been registered by
+	 * calling WaitEventExtensionRegisterName() in the current process, in
+	 * which case give up and return an unknown state.
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionNamesAllocated ||
+		WaitEventExtensionNames[eventId] == NULL)
+		return "???";
+
+	return WaitEventExtensionNames[eventId];
+}
+
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -149,6 +308,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -170,13 +332,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 3fabad96d9..2ea4789b00 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	BufferPin	"Waiting to acquire an exclusive pin on a buffer
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	Extension	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	Extension	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index c293871313..44764d4bf4 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,6 +39,25 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
+# Check the wait event used by the dynamic bgworker.  For a session
+# without the state in shared memory known, the default of "???" is the
+# value waited on.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	'???');
+is($result, 1, 'dynamic bgworker has reported "???" as wait event');
+# If the shared memory state is loaded, the wait event is the custom one.
+# The expected result is a special pattern here with a newline coming from the
+# first query where the shared memory state is set.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	qq[
+worker_spi_main]);
+is($result, 1,
+	'dynamic bgworker has reported "worker_spi_main" as wait event');
+
 note "testing bgworkers loaded with shared_preload_libraries";
 
 # Create the database first so as the workers can connect to it when
@@ -58,9 +77,9 @@ $node->restart;
 # Check that bgworkers have been registered and launched.
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname) FROM pg_stat_activity
-            WHERE backend_type = 'worker_spi' GROUP BY datname;],
-		'mydb|3'),
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
+            WHERE backend_type = 'worker_spi' GROUP BY datname, wait_event;],
+		'mydb|3|worker_spi_main'),
 	'bgworkers all launched'
 ) or die "Timed out while waiting for bgworkers to be launched";
 
@@ -70,10 +89,10 @@ my $worker1_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(1);');
 my $worker2_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(2);');
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname)  FROM pg_stat_activity
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
             WHERE backend_type = 'worker_spi dynamic' AND
-            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname;],
-		'mydb|2'),
+            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, wait_event;],
+		'mydb|2|worker_spi_main'),
 	'dynamic bgworkers all launched'
 ) or die "Timed out while waiting for dynamic bgworkers to be launched";
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index e9d5b07373..f13f7e0f98 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,3 +7,8 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
+
+CREATE FUNCTION worker_spi_init()
+RETURNS VOID STRICT
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 903dcddef9..8e234f912c 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,22 +44,96 @@
 
 PG_MODULE_MAGIC;
 
+PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
+/* Shared memory state */
+typedef struct worker_spi_state
+{
+	/* the wait event defined during initialization phase */
+	uint32		wait_event;
+} worker_spi_state;
+
+static worker_spi_state *wsstate = NULL;	/* pointer to shared memory */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void worker_spi_shmem_request(void);
+static void worker_spi_shmem_startup(void);
+static void worker_spi_shmem_init(void);
+static Size worker_spi_memsize(void);
+
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database = NULL;
 
-
 typedef struct worktable
 {
 	const char *schema;
 	const char *name;
 } worktable;
 
+static void
+worker_spi_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(worker_spi_memsize());
+}
+
+static void
+worker_spi_shmem_startup(void)
+{
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	worker_spi_shmem_init();
+}
+
+static Size
+worker_spi_memsize(void)
+{
+	return MAXALIGN(sizeof(worker_spi_state));
+}
+
+/*
+ * Initialize the shared memory state of worker_spi.
+ *
+ * This routine allocates a new wait event when called the first time.
+ * On follow-up calls, the name of the wait event associated with the
+ * existing shared memory state is registered.
+ */
+static void
+worker_spi_shmem_init(void)
+{
+	bool		found;
+
+	wsstate = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	wsstate = ShmemInitStruct("custom_wait_event",
+							  sizeof(worker_spi_state),
+							  &found);
+
+	/* Define a new wait event */
+	if (!found)
+		wsstate->wait_event = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
+	return;
+}
+
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -149,6 +223,9 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -199,7 +276,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 WAIT_EVENT_EXTENSION);
+						 wsstate->wait_event);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -328,6 +405,11 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = worker_spi_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = worker_spi_shmem_startup;
+
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -351,6 +433,21 @@ _PG_init(void)
 	}
 }
 
+/*
+ * Wrapper to initialize a session with the shared memory state
+ * used by this module.  This is a convenience routine to be able to
+ * see the custom wait event stored in shared memory without loading
+ * through shared_preload_libraries.
+ */
+Datum
+worker_spi_init(PG_FUNCTION_ARGS)
+{
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
+	PG_RETURN_VOID();
+}
+
 /*
  * Dynamically launch an SPI worker.
  */
@@ -363,6 +460,9 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index 588b720f57..023761dacb 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1117,12 +1117,13 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
 
    <note>
     <para>
-     Extensions can add <literal>LWLock</literal> types to the list shown in
-     <xref linkend="wait-event-lwlock-table"/>.  In some cases, the name
-     assigned by an extension will not be available in all server processes;
-     so an <literal>LWLock</literal> wait event might be reported as
-     just <quote><literal>extension</literal></quote> rather than the
-     extension-assigned name.
+     Extensions can add <literal>Extension</literal> and <literal>LWLock</literal> types
+     to the list shown in <xref linkend="wait-event-extension-table"/> and
+     <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
+     assigned by an extension will not be available in all server processes
+     if the extension's library is not loaded; so a custom wait event might
+     be reported as just <quote><literal>???</literal></quote>
+     rather than the custom name assigned.
     </para>
    </note>
  </sect2>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9620ea9ae3..315afef2a5 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3453,6 +3453,39 @@ if (!ptr)
     </para>
    </sect2>
 
+   <sect2 id="xfunc-addin-wait-events">
+    <title>Shared Memory and Custom Wait Events</title>
+
+    <para>
+     Add-ins can define custom wait events under the wait event type
+     <literal>Extension</literal>. The add-in's shared library must be
+     preloaded by specifying it in <literal>shared_preload_libraries</literal>,
+     and register a <literal>shmem_request_hook</literal> and a
+     <literal>shmem_startup_hook</literal> in its
+     <function>_PG_init</function> function.
+     <literal>shmem_request_hook</literal> can request a shared memory size
+     to be later used at startup by calling:
+<programlisting>
+void RequestAddinShmemSpace(int size)
+</programlisting>
+    </para>
+    <para>
+     <literal>shmem_startup_hook</literal> can allocate in shared memory
+     custom wait events by calling while holding the LWLock
+     <function>AddinShmemInitLock</function> to avoid any race conditions:
+<programlisting>
+uint32 WaitEventExtensionNew(void)
+</programlisting>
+     Next, each process needs to associate the wait event allocated previously
+     to a user-facing custom string, which is something done by calling:
+<programlisting>
+void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+</programlisting>
+     An example can be found in <filename>src/test/modules/worker_spi</filename>
+     in the PostgreSQL source tree.
+    </para>
+   </sect2>
+
    <sect2 id="extend-cpp">
     <title>Using C++ for Extensibility</title>
 
-- 
2.40.1

From 7ab64c42814d36800f53d532a567ec55a5ff7174 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 31 Jul 2023 09:51:02 +0900
Subject: [PATCH v10] Support custom wait events for wait event type
 "Extension"

Two backend routines are added to allow extension to allocate and define
custom wait events, all of these being allocated as of type "Extension":
* WaitEventExtensionNew(), that allocates a wait event ID computed from
a counter in shared memory.
* WaitEventExtensionRegisterName(), to associate a custom string to the
wait event registered.

Note that this includes an example of how to use this new facility in
worker_spi, with tests for various scenarios.

The use of these routines should be extended across more in-core
modules, but this is left as work for future patches.
---
 src/include/utils/wait_event.h                |  26 +++
 src/backend/storage/ipc/ipci.c                |   3 +
 .../activity/generate-wait_event_types.pl     |   7 +-
 src/backend/utils/activity/wait_event.c       | 186 ++++++++++++++++--
 .../utils/activity/wait_event_names.txt       |   2 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  31 ++-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 +
 src/test/modules/worker_spi/worker_spi.c      | 104 +++++++++-
 doc/src/sgml/monitoring.sgml                  |  11 +-
 doc/src/sgml/xfunc.sgml                       |  45 +++++
 src/tools/pgindent/typedefs.list              |   2 +
 11 files changed, 395 insertions(+), 27 deletions(-)

diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..aad8bc08fa 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -38,6 +38,32 @@ extern void pgstat_reset_wait_event_storage(void);
 extern PGDLLIMPORT uint32 *my_wait_event_info;
 
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define their own wait events in this category.  First,
+ * they should call WaitEventExtensionNew() to get one or more wait event
+ * IDs that are allocated from a shared counter.  These can be used directly
+ * with pgstat_report_wait_start() or equivalent.  Next, each individual
+ * process should call WaitEventExtensionRegisterName() to associate a wait
+ * event string to the number allocated previously.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNew(void);
+extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
+										   const char *wait_event_name);
+
 /* ----------
  * pgstat_report_wait_start() -
  *
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index f63c991051..56335e8730 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -133,10 +133,11 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
-		  if ( $waitclass eq 'WaitEventLWLock'
+		  if ( $waitclass eq 'WaitEventExtension'
+			|| $waitclass eq 'WaitEventLWLock'
 			|| $waitclass eq 'WaitEventLock');
 
 		my $last = $waitclass;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 4aad11c111..eaf3d0515a 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -39,6 +42,171 @@ static const char *pgstat_get_wait_io(WaitEventIO w);
 static uint32 local_my_wait_event_info;
 uint32	   *my_wait_event_info = &local_my_wait_event_info;
 
+#define WAIT_EVENT_ID_MASK		0x0000FFFF
+#define WAIT_EVENT_CLASS_MASK	0xFF000000
+
+/* dynamic allocation counter for custom wait events in extensions */
+typedef struct WaitEventExtensionCounterData
+{
+	int			nextId;			/* next ID to assign */
+	slock_t		mutex;			/* protects the counter */
+} WaitEventExtensionCounterData;
+
+/* pointer to the shared memory */
+static WaitEventExtensionCounterData *WaitEventExtensionCounter;
+
+/* first event ID of custom wait events for extensions */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event IDs known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionNames = NULL;
+static int	WaitEventExtensionNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(WaitEventExtensionCounterData);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	bool		found;
+
+	WaitEventExtensionCounter = (WaitEventExtensionCounterData *)
+		ShmemInitStruct("WaitEventExtensionCounterData",
+						WaitEventExtensionShmemSize(), &found);
+
+	if (!found)
+	{
+		/* initialize the allocation counter and its spinlock. */
+		WaitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&WaitEventExtensionCounter->mutex);
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event.
+ */
+uint32
+WaitEventExtensionNew(void)
+{
+	uint16		eventId;
+
+	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+
+	SpinLockAcquire(&WaitEventExtensionCounter->mutex);
+
+	if (WaitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	{
+		SpinLockRelease(&WaitEventExtensionCounter->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
+	}
+
+	eventId = WaitEventExtensionCounter->nextId++;
+
+	SpinLockRelease(&WaitEventExtensionCounter->mutex);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic wait event name for extension in the lookup table
+ * of the current process.
+ *
+ * This routine will save a pointer to the wait event name passed as an argument,
+ * so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterName(uint32 wait_event_info,
+							   const char *wait_event_name)
+{
+	uint32		classId;
+	uint16		eventId;
+
+	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
+	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
+
+	/* Check the wait event class. */
+	if (classId != PG_WAIT_EXTENSION)
+		ereport(ERROR,
+				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("invalid wait event class %u", classId));
+
+	/* This should only be called for user-defined wait event. */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		ereport(ERROR,
+				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("invalid wait event ID %u", eventId));
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionNamesAllocated)
+	{
+		uint32		newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionNames == NULL)
+			WaitEventExtensionNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionNames =
+				repalloc0_array(WaitEventExtensionNames, const char *,
+								WaitEventExtensionNamesAllocated, newalloc);
+		WaitEventExtensionNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an wait event ID for extension.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Built-in event? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
+	 * However, it is possible that the name has never been registered by
+	 * calling WaitEventExtensionRegisterName() in the current process, in
+	 * which case give up and return an unknown state.
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionNamesAllocated ||
+		WaitEventExtensionNames[eventId] == NULL)
+		return "extension";
+
+	return WaitEventExtensionNames[eventId];
+}
+
 
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
@@ -82,7 +250,7 @@ pgstat_get_wait_event_type(uint32 wait_event_info)
 	if (wait_event_info == 0)
 		return NULL;
 
-	classId = wait_event_info & 0xFF000000;
+	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
 
 	switch (classId)
 	{
@@ -138,8 +306,8 @@ pgstat_get_wait_event(uint32 wait_event_info)
 	if (wait_event_info == 0)
 		return NULL;
 
-	classId = wait_event_info & 0xFF000000;
-	eventId = wait_event_info & 0x0000FFFF;
+	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
+	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
 
 	switch (classId)
 	{
@@ -149,6 +317,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -170,13 +341,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 3fabad96d9..2ea4789b00 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	BufferPin	"Waiting to acquire an exclusive pin on a buffer
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	Extension	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	Extension	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index 74e109f9a1..0798653d00 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,6 +39,25 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
+# Check the wait event used by the dynamic bgworker.  For a session without
+# the state in shared memory known, the default of "extension" is the value
+# waited on.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	'extension');
+is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
+# If the shared memory state is loaded, the wait event is the custom one.
+# The expected result is a special pattern here with a newline coming from the
+# first query where the shared memory state is set.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	qq[
+worker_spi_main]);
+is($result, 1,
+	'dynamic bgworker has reported "worker_spi_main" as wait event');
+
 note "testing bgworkers loaded with shared_preload_libraries";
 
 # Create the database first so as the workers can connect to it when
@@ -58,9 +77,9 @@ $node->restart;
 # Check that bgworkers have been registered and launched.
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname) FROM pg_stat_activity
-            WHERE backend_type = 'worker_spi' GROUP BY datname;],
-		'mydb|3'),
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
+            WHERE backend_type = 'worker_spi' GROUP BY datname, wait_event;],
+		'mydb|3|worker_spi_main'),
 	'bgworkers all launched'
 ) or die "Timed out while waiting for bgworkers to be launched";
 
@@ -72,10 +91,10 @@ my $worker2_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(11);');
 
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname)  FROM pg_stat_activity
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
             WHERE backend_type = 'worker_spi dynamic' AND
-            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname;],
-		'mydb|2'),
+            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, wait_event;],
+		'mydb|2|worker_spi_main'),
 	'dynamic bgworkers all launched'
 ) or die "Timed out while waiting for dynamic bgworkers to be launched";
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index e9d5b07373..f13f7e0f98 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,3 +7,8 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
+
+CREATE FUNCTION worker_spi_init()
+RETURNS VOID STRICT
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 903dcddef9..e030d4b1de 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,22 +44,96 @@
 
 PG_MODULE_MAGIC;
 
+PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
+/* Shared memory state */
+typedef struct worker_spi_state
+{
+	/* the wait event defined during initialization phase */
+	uint32		wait_event;
+} worker_spi_state;
+
+static worker_spi_state *wsstate = NULL;	/* pointer to shared memory */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void worker_spi_shmem_request(void);
+static void worker_spi_shmem_startup(void);
+static void worker_spi_shmem_init(void);
+static Size worker_spi_memsize(void);
+
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database = NULL;
 
-
 typedef struct worktable
 {
 	const char *schema;
 	const char *name;
 } worktable;
 
+static void
+worker_spi_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(worker_spi_memsize());
+}
+
+static void
+worker_spi_shmem_startup(void)
+{
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	worker_spi_shmem_init();
+}
+
+static Size
+worker_spi_memsize(void)
+{
+	return MAXALIGN(sizeof(worker_spi_state));
+}
+
+/*
+ * Initialize the shared memory state of worker_spi.
+ *
+ * This routine allocates a new wait event when called the first time.
+ * On follow-up calls, the name of the wait event associated with the
+ * existing shared memory state is registered.
+ */
+static void
+worker_spi_shmem_init(void)
+{
+	bool		found;
+
+	wsstate = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	wsstate = ShmemInitStruct("worker_spi State",
+							  sizeof(worker_spi_state),
+							  &found);
+
+	/* Define a new wait event */
+	if (!found)
+		wsstate->wait_event = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
+	return;
+}
+
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -149,6 +223,9 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -199,7 +276,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 WAIT_EVENT_EXTENSION);
+						 wsstate->wait_event);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -328,6 +405,11 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = worker_spi_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = worker_spi_shmem_startup;
+
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -351,6 +433,21 @@ _PG_init(void)
 	}
 }
 
+/*
+ * Wrapper to initialize a session with the shared memory state
+ * used by this module.  This is a convenience routine to be able to
+ * see the custom wait event stored in shared memory without loading
+ * through shared_preload_libraries.
+ */
+Datum
+worker_spi_init(PG_FUNCTION_ARGS)
+{
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
+	PG_RETURN_VOID();
+}
+
 /*
  * Dynamically launch an SPI worker.
  */
@@ -363,6 +460,9 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index 588b720f57..c97008eeb3 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1117,11 +1117,14 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
 
    <note>
     <para>
-     Extensions can add <literal>LWLock</literal> types to the list shown in
-     <xref linkend="wait-event-lwlock-table"/>.  In some cases, the name
+     Extensions can add <literal>Extension</literal> and
+     <literal>LWLock</literal> types
+     to the list shown in <xref linkend="wait-event-extension-table"/> and
+     <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
      assigned by an extension will not be available in all server processes;
-     so an <literal>LWLock</literal> wait event might be reported as
-     just <quote><literal>extension</literal></quote> rather than the
+     so an <literal>LWLock</literal> or <literal>Extension</literal> wait
+     event might be reported as just
+     <quote><literal>extension</literal></quote> rather than the
      extension-assigned name.
     </para>
    </note>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9620ea9ae3..d6345a775b 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3453,6 +3453,51 @@ if (!ptr)
     </para>
    </sect2>
 
+   <sect2 id="xfunc-addin-wait-events">
+    <title>Shared Memory and Custom Wait Events</title>
+
+    <para>
+     Add-ins can define custom wait events under the wait event type
+     <literal>Extension</literal>. The add-in's shared library must be
+     preloaded by specifying it in <literal>shared_preload_libraries</literal>,
+     and register a <literal>shmem_request_hook</literal> and a
+     <literal>shmem_startup_hook</literal> in its
+     <function>_PG_init</function> function.
+     <literal>shmem_request_hook</literal> can request a shared memory size
+     to be later used at startup by calling:
+<programlisting>
+void RequestAddinShmemSpace(int size)
+</programlisting>
+    </para>
+    <para>
+     <literal>shmem_startup_hook</literal> can allocate in shared memory
+     custom wait events by calling while holding the LWLock
+     <function>AddinShmemInitLock</function> to avoid any race conditions:
+<programlisting>
+uint32 WaitEventExtensionNew(void)
+</programlisting>
+     Next, each process needs to associate the wait event allocated previously
+     to a user-facing custom string, which is something done by calling:
+<programlisting>
+void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+</programlisting>
+     An example can be found in <filename>src/test/modules/worker_spi</filename>
+     in the PostgreSQL source tree.
+    </para>
+    <para>
+     Custom wait events can be viewed in
+     <link linkend="monitoring-pg-stat-activity-view"><structname>pg_stat_activity</structname></link>:
+<screen>
+=# SELECT wait_event_type, wait_event FROM pg_stat_activity
+     WHERE backend_type ~ 'worker_spi';
+ wait_event_type |   wait_event
+-----------------+-----------------
+ Extension       | worker_spi_main
+(1 row)
+</screen>
+    </para>
+   </sect2>
+
    <sect2 id="extend-cpp">
     <title>Using C++ for Extensibility</title>
 
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 11d47294cf..ab97f1794a 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2992,6 +2992,7 @@ WaitEventActivity
 WaitEventBufferPin
 WaitEventClient
 WaitEventExtension
+WaitEventExtensionCounterData
 WaitEventIO
 WaitEventIPC
 WaitEventSet
@@ -3862,6 +3863,7 @@ wchar2mb_with_len_converter
 wchar_t
 win32_deadchild_waitinfo
 wint_t
+worker_spi_state
 worker_state
 worktable
 wrap
-- 
2.40.1

From 23b6c4c36fbf23ca9c6fb4f42a3f943be44689a6 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 31 Jul 2023 16:04:20 +0900
Subject: [PATCH v11] Support custom wait events for wait event type
 "Extension"

Two backend routines are added to allow extension to allocate and define
custom wait events, all of these being allocated as of type "Extension":
* WaitEventExtensionNew(), that allocates a wait event ID computed from
a counter in shared memory.
* WaitEventExtensionRegisterName(), to associate a custom string to the
wait event registered.

Note that this includes an example of how to use this new facility in
worker_spi, with tests for various scenarios.

The use of these routines should be extended across more in-core
modules, but this is left as work for future patches.
---
 src/include/utils/wait_event.h                |  26 +++
 src/backend/storage/ipc/ipci.c                |   3 +
 .../activity/generate-wait_event_types.pl     |   7 +-
 src/backend/utils/activity/wait_event.c       | 178 +++++++++++++++++-
 .../utils/activity/wait_event_names.txt       |   2 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  34 +++-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 +
 src/test/modules/worker_spi/worker_spi.c      | 104 +++++++++-
 doc/src/sgml/monitoring.sgml                  |  11 +-
 doc/src/sgml/xfunc.sgml                       |  45 +++++
 src/tools/pgindent/typedefs.list              |   2 +
 11 files changed, 393 insertions(+), 24 deletions(-)

diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index 4517425f84..aad8bc08fa 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -38,6 +38,32 @@ extern void pgstat_reset_wait_event_storage(void);
 extern PGDLLIMPORT uint32 *my_wait_event_info;
 
 
+/* ----------
+ * Wait Events - Extension
+ *
+ * Use this category when the server process is waiting for some condition
+ * defined by an extension module.
+ *
+ * Extensions can define their own wait events in this category.  First,
+ * they should call WaitEventExtensionNew() to get one or more wait event
+ * IDs that are allocated from a shared counter.  These can be used directly
+ * with pgstat_report_wait_start() or equivalent.  Next, each individual
+ * process should call WaitEventExtensionRegisterName() to associate a wait
+ * event string to the number allocated previously.
+ */
+typedef enum
+{
+	WAIT_EVENT_EXTENSION = PG_WAIT_EXTENSION,
+	WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED
+} WaitEventExtension;
+
+extern void WaitEventExtensionShmemInit(void);
+extern Size WaitEventExtensionShmemSize(void);
+
+extern uint32 WaitEventExtensionNew(void);
+extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
+										   const char *wait_event_name);
+
 /* ----------
  * pgstat_report_wait_start() -
  *
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index cc387c00a1..5551afffc0 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -49,6 +49,7 @@
 #include "storage/spin.h"
 #include "utils/guc.h"
 #include "utils/snapmgr.h"
+#include "utils/wait_event.h"
 
 /* GUCs */
 int			shared_memory_type = DEFAULT_SHARED_MEMORY_TYPE;
@@ -142,6 +143,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, SyncScanShmemSize());
 	size = add_size(size, AsyncShmemSize());
 	size = add_size(size, StatsShmemSize());
+	size = add_size(size, WaitEventExtensionShmemSize());
 #ifdef EXEC_BACKEND
 	size = add_size(size, ShmemBackendArraySize());
 #endif
@@ -301,6 +303,7 @@ CreateSharedMemoryAndSemaphores(void)
 	SyncScanShmemInit();
 	AsyncShmemInit();
 	StatsShmemInit();
+	WaitEventExtensionShmemInit();
 
 #ifdef EXEC_BACKEND
 
diff --git a/src/backend/utils/activity/generate-wait_event_types.pl b/src/backend/utils/activity/generate-wait_event_types.pl
index f63c991051..56335e8730 100644
--- a/src/backend/utils/activity/generate-wait_event_types.pl
+++ b/src/backend/utils/activity/generate-wait_event_types.pl
@@ -133,10 +133,11 @@ if ($gen_code)
 	foreach my $waitclass (sort { uc($a) cmp uc($b) } keys %hashwe)
 	{
 
-		# Don't generate .c and .h files for LWLock and Lock, these are
-		# handled independently.
+		# Don't generate .c and .h files for Extension, LWLock and
+		# Lock, these are handled independently.
 		next
-		  if ( $waitclass eq 'WaitEventLWLock'
+		  if ( $waitclass eq 'WaitEventExtension'
+			|| $waitclass eq 'WaitEventLWLock'
 			|| $waitclass eq 'WaitEventLock');
 
 		my $last = $waitclass;
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index 59177de7a0..b3596ece80 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,15 +22,18 @@
  */
 #include "postgres.h"
 
+#include "miscadmin.h"
+#include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
 #include "storage/lwlock.h"		/* for GetLWLockIdentifier */
+#include "storage/spin.h"
+#include "utils/memutils.h"
 #include "utils/wait_event.h"
 
 
 static const char *pgstat_get_wait_activity(WaitEventActivity w);
 static const char *pgstat_get_wait_bufferpin(WaitEventBufferPin w);
 static const char *pgstat_get_wait_client(WaitEventClient w);
-static const char *pgstat_get_wait_extension(WaitEventExtension w);
 static const char *pgstat_get_wait_ipc(WaitEventIPC w);
 static const char *pgstat_get_wait_timeout(WaitEventTimeout w);
 static const char *pgstat_get_wait_io(WaitEventIO w);
@@ -42,6 +45,169 @@ uint32	   *my_wait_event_info = &local_my_wait_event_info;
 #define WAIT_EVENT_CLASS_MASK	0xFF000000
 #define WAIT_EVENT_ID_MASK		0x0000FFFF
 
+/* dynamic allocation counter for custom wait events in extensions */
+typedef struct WaitEventExtensionCounterData
+{
+	int			nextId;			/* next ID to assign */
+	slock_t		mutex;			/* protects the counter */
+} WaitEventExtensionCounterData;
+
+/* pointer to the shared memory */
+static WaitEventExtensionCounterData *WaitEventExtensionCounter;
+
+/* first event ID of custom wait events for extensions */
+#define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
+	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
+
+/*
+ * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
+ * stores the names of all dynamically-created event IDs known to the current
+ * process.  Any unused entries in the array will contain NULL.
+ */
+static const char **WaitEventExtensionNames = NULL;
+static int	WaitEventExtensionNamesAllocated = 0;
+
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
+/*
+ *  Return the space for dynamic allocation counter.
+ */
+Size
+WaitEventExtensionShmemSize(void)
+{
+	return sizeof(WaitEventExtensionCounterData);
+}
+
+/*
+ * Allocate shmem space for dynamic allocation counter.
+ */
+void
+WaitEventExtensionShmemInit(void)
+{
+	bool		found;
+
+	WaitEventExtensionCounter = (WaitEventExtensionCounterData *)
+		ShmemInitStruct("WaitEventExtensionCounterData",
+						WaitEventExtensionShmemSize(), &found);
+
+	if (!found)
+	{
+		/* initialize the allocation counter and its spinlock. */
+		WaitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&WaitEventExtensionCounter->mutex);
+	}
+}
+
+/*
+ * Allocate a new event ID and return the wait event.
+ */
+uint32
+WaitEventExtensionNew(void)
+{
+	uint16		eventId;
+
+	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+
+	SpinLockAcquire(&WaitEventExtensionCounter->mutex);
+
+	if (WaitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	{
+		SpinLockRelease(&WaitEventExtensionCounter->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
+	}
+
+	eventId = WaitEventExtensionCounter->nextId++;
+
+	SpinLockRelease(&WaitEventExtensionCounter->mutex);
+
+	return PG_WAIT_EXTENSION | eventId;
+}
+
+/*
+ * Register a dynamic wait event name for extension in the lookup table
+ * of the current process.
+ *
+ * This routine will save a pointer to the wait event name passed as an argument,
+ * so the name should be allocated in a backend-lifetime context
+ * (shared memory, TopMemoryContext, static constant, or similar).
+ *
+ * The "wait_event_name" will be user-visible as a wait event name, so try to
+ * use a name that fits the style for those.
+ */
+void
+WaitEventExtensionRegisterName(uint32 wait_event_info,
+							   const char *wait_event_name)
+{
+	uint32		classId;
+	uint16		eventId;
+
+	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
+	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
+
+	/* Check the wait event class. */
+	if (classId != PG_WAIT_EXTENSION)
+		ereport(ERROR,
+				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("invalid wait event class %u", classId));
+
+	/* This should only be called for user-defined wait event. */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		ereport(ERROR,
+				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("invalid wait event ID %u", eventId));
+
+	/* Convert to array index. */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	/* If necessary, create or enlarge array. */
+	if (eventId >= WaitEventExtensionNamesAllocated)
+	{
+		uint32		newalloc;
+
+		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
+
+		if (WaitEventExtensionNames == NULL)
+			WaitEventExtensionNames = (const char **)
+				MemoryContextAllocZero(TopMemoryContext,
+									   newalloc * sizeof(char *));
+		else
+			WaitEventExtensionNames =
+				repalloc0_array(WaitEventExtensionNames, const char *,
+								WaitEventExtensionNamesAllocated, newalloc);
+		WaitEventExtensionNamesAllocated = newalloc;
+	}
+
+	WaitEventExtensionNames[eventId] = wait_event_name;
+}
+
+/*
+ * Return the name of an wait event ID for extension.
+ */
+static const char *
+GetWaitEventExtensionIdentifier(uint16 eventId)
+{
+	/* Built-in event? */
+	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
+		return "Extension";
+
+	/*
+	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
+	 * However, it is possible that the name has never been registered by
+	 * calling WaitEventExtensionRegisterName() in the current process, in
+	 * which case give up and return "extension".
+	 */
+	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+
+	if (eventId >= WaitEventExtensionNamesAllocated ||
+		WaitEventExtensionNames[eventId] == NULL)
+		return "extension";
+
+	return WaitEventExtensionNames[eventId];
+}
+
+
 /*
  * Configure wait event reporting to report wait events to *wait_event_info.
  * *wait_event_info needs to be valid until pgstat_reset_wait_event_storage()
@@ -151,6 +317,9 @@ pgstat_get_wait_event(uint32 wait_event_info)
 		case PG_WAIT_LOCK:
 			event_name = GetLockNameFromTagType(eventId);
 			break;
+		case PG_WAIT_EXTENSION:
+			event_name = GetWaitEventExtensionIdentifier(eventId);
+			break;
 		case PG_WAIT_BUFFERPIN:
 			{
 				WaitEventBufferPin w = (WaitEventBufferPin) wait_event_info;
@@ -172,13 +341,6 @@ pgstat_get_wait_event(uint32 wait_event_info)
 				event_name = pgstat_get_wait_client(w);
 				break;
 			}
-		case PG_WAIT_EXTENSION:
-			{
-				WaitEventExtension w = (WaitEventExtension) wait_event_info;
-
-				event_name = pgstat_get_wait_extension(w);
-				break;
-			}
 		case PG_WAIT_IPC:
 			{
 				WaitEventIPC w = (WaitEventIPC) wait_event_info;
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 3fabad96d9..2ea4789b00 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -261,7 +261,7 @@ WAIT_EVENT_BUFFER_PIN	BufferPin	"Waiting to acquire an exclusive pin on a buffer
 
 Section: ClassName - WaitEventExtension
 
-WAIT_EVENT_EXTENSION	Extension	"Waiting in an extension."
+WAIT_EVENT_DOCONLY	Extension	"Waiting in an extension."
 
 #
 # Wait events - LWLock
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index 74e109f9a1..c3e7f5fbe6 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,6 +39,28 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
+# Check the wait event used by the dynamic bgworker.  For a session without
+# the state in shared memory known, the default of "extension" is the value
+# waited on.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	'extension');
+is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
+
+# If the shared memory state is loaded (here with worker_spi_init within
+# the same connection as the one querying pg_stat_activity), the wait
+# event is the custom one.
+# The expected result is a special pattern here with a newline coming from the
+# first query where the shared memory state is set.
+$result = $node->poll_query_until(
+	'postgres',
+	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
+	qq[
+worker_spi_main]);
+is($result, 1,
+	'dynamic bgworker has reported "worker_spi_main" as wait event');
+
 note "testing bgworkers loaded with shared_preload_libraries";
 
 # Create the database first so as the workers can connect to it when
@@ -58,9 +80,9 @@ $node->restart;
 # Check that bgworkers have been registered and launched.
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname) FROM pg_stat_activity
-            WHERE backend_type = 'worker_spi' GROUP BY datname;],
-		'mydb|3'),
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
+            WHERE backend_type = 'worker_spi' GROUP BY datname, wait_event;],
+		'mydb|3|worker_spi_main'),
 	'bgworkers all launched'
 ) or die "Timed out while waiting for bgworkers to be launched";
 
@@ -72,10 +94,10 @@ my $worker2_pid = $node->safe_psql('mydb', 'SELECT worker_spi_launch(11);');
 
 ok( $node->poll_query_until(
 		'mydb',
-		qq[SELECT datname, count(datname)  FROM pg_stat_activity
+		qq[SELECT datname, count(datname), wait_event FROM pg_stat_activity
             WHERE backend_type = 'worker_spi dynamic' AND
-            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname;],
-		'mydb|2'),
+            pid IN ($worker1_pid, $worker2_pid) GROUP BY datname, wait_event;],
+		'mydb|2|worker_spi_main'),
 	'dynamic bgworkers all launched'
 ) or die "Timed out while waiting for dynamic bgworkers to be launched";
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index e9d5b07373..f13f7e0f98 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,3 +7,8 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
+
+CREATE FUNCTION worker_spi_init()
+RETURNS VOID STRICT
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 903dcddef9..e030d4b1de 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,22 +44,96 @@
 
 PG_MODULE_MAGIC;
 
+PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
+/* Shared memory state */
+typedef struct worker_spi_state
+{
+	/* the wait event defined during initialization phase */
+	uint32		wait_event;
+} worker_spi_state;
+
+static worker_spi_state *wsstate = NULL;	/* pointer to shared memory */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void worker_spi_shmem_request(void);
+static void worker_spi_shmem_startup(void);
+static void worker_spi_shmem_init(void);
+static Size worker_spi_memsize(void);
+
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database = NULL;
 
-
 typedef struct worktable
 {
 	const char *schema;
 	const char *name;
 } worktable;
 
+static void
+worker_spi_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(worker_spi_memsize());
+}
+
+static void
+worker_spi_shmem_startup(void)
+{
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	worker_spi_shmem_init();
+}
+
+static Size
+worker_spi_memsize(void)
+{
+	return MAXALIGN(sizeof(worker_spi_state));
+}
+
+/*
+ * Initialize the shared memory state of worker_spi.
+ *
+ * This routine allocates a new wait event when called the first time.
+ * On follow-up calls, the name of the wait event associated with the
+ * existing shared memory state is registered.
+ */
+static void
+worker_spi_shmem_init(void)
+{
+	bool		found;
+
+	wsstate = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	wsstate = ShmemInitStruct("worker_spi State",
+							  sizeof(worker_spi_state),
+							  &found);
+
+	/* Define a new wait event */
+	if (!found)
+		wsstate->wait_event = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
+	return;
+}
+
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -149,6 +223,9 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -199,7 +276,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 WAIT_EVENT_EXTENSION);
+						 wsstate->wait_event);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -328,6 +405,11 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = worker_spi_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = worker_spi_shmem_startup;
+
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -351,6 +433,21 @@ _PG_init(void)
 	}
 }
 
+/*
+ * Wrapper to initialize a session with the shared memory state
+ * used by this module.  This is a convenience routine to be able to
+ * see the custom wait event stored in shared memory without loading
+ * through shared_preload_libraries.
+ */
+Datum
+worker_spi_init(PG_FUNCTION_ARGS)
+{
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
+	PG_RETURN_VOID();
+}
+
 /*
  * Dynamically launch an SPI worker.
  */
@@ -363,6 +460,9 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index 588b720f57..f4fc5d814f 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1117,11 +1117,14 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
 
    <note>
     <para>
-     Extensions can add <literal>LWLock</literal> types to the list shown in
-     <xref linkend="wait-event-lwlock-table"/>.  In some cases, the name
+     Extensions can add <literal>Extension</literal> and
+     <literal>LWLock</literal> types
+     to the list shown in <xref linkend="wait-event-extension-table"/> and
+     <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
      assigned by an extension will not be available in all server processes;
-     so an <literal>LWLock</literal> wait event might be reported as
-     just <quote><literal>extension</literal></quote> rather than the
+     so an <literal>Extension</literal> or <literal>LWLock</literal> wait
+     event might be reported as just
+     <quote><literal>extension</literal></quote> rather than the
      extension-assigned name.
     </para>
    </note>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 9620ea9ae3..d6345a775b 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3453,6 +3453,51 @@ if (!ptr)
     </para>
    </sect2>
 
+   <sect2 id="xfunc-addin-wait-events">
+    <title>Shared Memory and Custom Wait Events</title>
+
+    <para>
+     Add-ins can define custom wait events under the wait event type
+     <literal>Extension</literal>. The add-in's shared library must be
+     preloaded by specifying it in <literal>shared_preload_libraries</literal>,
+     and register a <literal>shmem_request_hook</literal> and a
+     <literal>shmem_startup_hook</literal> in its
+     <function>_PG_init</function> function.
+     <literal>shmem_request_hook</literal> can request a shared memory size
+     to be later used at startup by calling:
+<programlisting>
+void RequestAddinShmemSpace(int size)
+</programlisting>
+    </para>
+    <para>
+     <literal>shmem_startup_hook</literal> can allocate in shared memory
+     custom wait events by calling while holding the LWLock
+     <function>AddinShmemInitLock</function> to avoid any race conditions:
+<programlisting>
+uint32 WaitEventExtensionNew(void)
+</programlisting>
+     Next, each process needs to associate the wait event allocated previously
+     to a user-facing custom string, which is something done by calling:
+<programlisting>
+void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+</programlisting>
+     An example can be found in <filename>src/test/modules/worker_spi</filename>
+     in the PostgreSQL source tree.
+    </para>
+    <para>
+     Custom wait events can be viewed in
+     <link linkend="monitoring-pg-stat-activity-view"><structname>pg_stat_activity</structname></link>:
+<screen>
+=# SELECT wait_event_type, wait_event FROM pg_stat_activity
+     WHERE backend_type ~ 'worker_spi';
+ wait_event_type |   wait_event
+-----------------+-----------------
+ Extension       | worker_spi_main
+(1 row)
+</screen>
+    </para>
+   </sect2>
+
    <sect2 id="extend-cpp">
     <title>Using C++ for Extensibility</title>
 
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 11d47294cf..ab97f1794a 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2992,6 +2992,7 @@ WaitEventActivity
 WaitEventBufferPin
 WaitEventClient
 WaitEventExtension
+WaitEventExtensionCounterData
 WaitEventIO
 WaitEventIPC
 WaitEventSet
@@ -3862,6 +3863,7 @@ wchar2mb_with_len_converter
 wchar_t
 win32_deadchild_waitinfo
 wint_t
+worker_spi_state
 worker_state
 worktable
 wrap
-- 
2.40.1

From 0e3ccc6474bfcc51114d9363b7819b68f37fcad3 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <mshr.ikeda@ntt.com>
Date: Tue, 8 Aug 2023 19:24:32 +0900
Subject: [PATCH] Change to manage custom wait events in shared hash

Currently, names of the custom wait event must be
registered per backends. This patch relaxes the
constraints to store the wait events and their names
in the dynamic shared hashtable. So, all backends can
look up the wait event names on pg_stat_activity view
without additional processing by extensions.
---
 doc/src/sgml/monitoring.sgml                  |   7 +-
 doc/src/sgml/xfunc.sgml                       |  10 +-
 src/backend/storage/ipc/ipci.c                |   6 +
 src/backend/storage/lmgr/lwlock.c             |   2 +
 src/backend/utils/activity/wait_event.c       | 275 ++++++++++++------
 src/backend/utils/init/postinit.c             |   3 +
 src/include/storage/lwlock.h                  |   2 +
 src/include/utils/wait_event.h                |  17 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  18 +-
 src/test/modules/worker_spi/worker_spi.c      |   6 +-
 src/tools/pgindent/typedefs.list              |   2 +-
 11 files changed, 216 insertions(+), 132 deletions(-)

diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index f4fc5d814f..19181832d7 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1121,10 +1121,9 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
      <literal>LWLock</literal> types
      to the list shown in <xref linkend="wait-event-extension-table"/> and
      <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
-     assigned by an extension will not be available in all server processes;
-     so an <literal>Extension</literal> or <literal>LWLock</literal> wait
-     event might be reported as just
-     <quote><literal>extension</literal></quote> rather than the
+     of <literal>LWLock</literal> assigned by an extension will not be
+     available in all server processes; so an wait event might be reported
+     as just <quote><literal>extension</literal></quote> rather than the
      extension-assigned name.
     </para>
    </note>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index d6345a775b..7fec034db4 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3470,17 +3470,13 @@ void RequestAddinShmemSpace(int size)
 </programlisting>
     </para>
     <para>
-     <literal>shmem_startup_hook</literal> can allocate in shared memory
+     <literal>shmem_startup_hook</literal> can allocate in dynamic shared memory
      custom wait events by calling while holding the LWLock
      <function>AddinShmemInitLock</function> to avoid any race conditions:
 <programlisting>
-uint32 WaitEventExtensionNew(void)
-</programlisting>
-     Next, each process needs to associate the wait event allocated previously
-     to a user-facing custom string, which is something done by calling:
-<programlisting>
-void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+uint32 WaitEventExtensionNew(const char *wait_event_name)
 </programlisting>
+     The wait event is associated to a user-facing custom string.
      An example can be found in <filename>src/test/modules/worker_spi</filename>
      in the PostgreSQL source tree.
     </para>
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 5551afffc0..b3d2a0b252 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -323,6 +323,12 @@ CreateSharedMemoryAndSemaphores(void)
 	 */
 	if (shmem_startup_hook)
 		shmem_startup_hook();
+
+	/*
+	 * Detach from the dynamic shared hash because postmaster will never
+	 * access these again
+	 */
+	WaitEventExtensionDetachShmem();
 }
 
 /*
diff --git a/src/backend/storage/lmgr/lwlock.c b/src/backend/storage/lmgr/lwlock.c
index 315a78cda9..cf1bdbed85 100644
--- a/src/backend/storage/lmgr/lwlock.c
+++ b/src/backend/storage/lmgr/lwlock.c
@@ -190,6 +190,8 @@ static const char *const BuiltinTrancheNames[] = {
 	"LogicalRepLauncherDSA",
 	/* LWTRANCHE_LAUNCHER_HASH: */
 	"LogicalRepLauncherHash",
+	"LWTRANCHE_WAITEVENTEXTENSION_DSA",
+	"LWTRANCHE_WAITEVENTEXTENSION_HASH",
 };
 
 StaticAssertDecl(lengthof(BuiltinTrancheNames) ==
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index b3596ece80..dbdf681110 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -22,6 +22,7 @@
  */
 #include "postgres.h"
 
+#include "lib/dshash.h"
 #include "miscadmin.h"
 #include "port/pg_bitutils.h"
 #include "storage/lmgr.h"		/* for GetLockNameFromTagType */
@@ -45,141 +46,235 @@ uint32	   *my_wait_event_info = &local_my_wait_event_info;
 #define WAIT_EVENT_CLASS_MASK	0xFF000000
 #define WAIT_EVENT_ID_MASK		0x0000FFFF
 
-/* dynamic allocation counter for custom wait events in extensions */
-typedef struct WaitEventExtensionCounterData
+/* hash table entry for finding the custom wait event name for a key */
+#define MAX_WAIT_EVENT_EXTENSION_NAME		256
+typedef struct WaitEventExtensionNameEntry
 {
-	int			nextId;			/* next ID to assign */
-	slock_t		mutex;			/* protects the counter */
-} WaitEventExtensionCounterData;
+	uint16	event_id;	/* hash key */
+	uint32	status;		/* hash status */
+	char	wait_event_name[MAX_WAIT_EVENT_EXTENSION_NAME]; /* custom wait event name */
+} WaitEventExtensionNameEntry;
+
+/* for references to shared custom wait event entries */
+#define SH_PREFIX		wait_event_extension_name_hash
+#define SH_ELEMENT_TYPE	WaitEventExtensionNameEntry
+#define SH_KEY_TYPE		uint16
+#define	SH_KEY			event_id
+#define SH_HASH_KEY(tb, key)	dshash_memhash(&key, sizeof(uint16), NULL)
+#define SH_EQUAL(tb, a, b)		dshash_memcmp(&a, &b, sizeof(uint16), NULL) == 0
+#define	SH_SCOPE		static inline
+#define SH_DECLARE
+#define SH_DEFINE
+#include "lib/simplehash.h"
+
+/* parameter for the shared hash */
+static const dshash_parameters dsh_params = {
+	sizeof(uint16),
+	sizeof(WaitEventExtensionNameEntry),
+	dshash_memcmp,
+	dshash_memhash,
+	LWTRANCHE_WAITEVENTEXTENSION_HASH
+};
+
+/*
+ * Shared state information for custom wait events in extensions.
+ * This manages a dynamic shared hash and a dynamic allocation counter.
+ */
+typedef struct WaitEventExtensionSharedState
+{
+	/*
+	 * Custom wait event names are registered in shared custom wait event hash.
+	 * Other backends can look up them without additional processing per backend
+	 * like LWLockRegisterTranche().
+	 */
+	void		*raw_dsa_area;
+	dshash_table_handle hash_handle;
+
+	int			nextId;					/* next event ID to assign */
+	slock_t		mutex;					/* protects the counter */
+} WaitEventExtensionSharedState;
 
 /* pointer to the shared memory */
-static WaitEventExtensionCounterData *WaitEventExtensionCounter;
+static WaitEventExtensionSharedState *wee_state;
+
+/* backend-local state to manage dynamic shared hash */
+typedef struct WaitEventExtension_LocalState
+{
+	dsa_area   *dsa;
+	dshash_table *dsh;
+} WaitEventExtension_LocalState;
+
+static WaitEventExtension_LocalState wee_localstate;
 
 /* first event ID of custom wait events for extensions */
 #define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
 	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
 
+static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+
 /*
- * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
- * stores the names of all dynamically-created event IDs known to the current
- * process.  Any unused entries in the array will contain NULL.
+ * The size of the shared memory allocation for custom wait events stored in
+ * the shared hash table. This allocation will be done as part of the main shared
+ * memory, rather than dynamic shared memory, allowing it to be initialized in
+ * postmaster.
  */
-static const char **WaitEventExtensionNames = NULL;
-static int	WaitEventExtensionNamesAllocated = 0;
+static Size
+wait_event_extension_dsa_init_size(void)
+{
+	Size		sz;
 
-static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
+	/* Same size as pgstat_dsa_init_size() for dshash */
+	sz = 256 * 1024;
+	Assert(dsa_minimum_size() <= sz);
+	return MAXALIGN(sz);
+}
 
 /*
- *  Return the space for dynamic allocation counter.
+ *  Return the space for dynamic shared hash and dynamic allocation counter.
  */
 Size
 WaitEventExtensionShmemSize(void)
 {
-	return sizeof(WaitEventExtensionCounterData);
+	Size	sz;
+	sz = MAXALIGN(sizeof(WaitEventExtensionSharedState));
+	sz = add_size(sz, wait_event_extension_dsa_init_size());
+	return sz;
 }
 
 /*
- * Allocate shmem space for dynamic allocation counter.
+ * Allocate shmem space for dynamic shared hash and dynamic allocation counter.
  */
 void
 WaitEventExtensionShmemInit(void)
 {
 	bool		found;
 
-	WaitEventExtensionCounter = (WaitEventExtensionCounterData *)
-		ShmemInitStruct("WaitEventExtensionCounterData",
+	wee_state = (WaitEventExtensionSharedState *)
+		ShmemInitStruct("WaitEventExtensionSharedState",
 						WaitEventExtensionShmemSize(), &found);
 
 	if (!found)
 	{
+		char	*p = (char *) wee_state;
+
 		/* initialize the allocation counter and its spinlock. */
-		WaitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
-		SpinLockInit(&WaitEventExtensionCounter->mutex);
+		wee_state->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+		SpinLockInit(&wee_state->mutex);
+
+		/* initialize dsa and dshash */
+		p += (MAXALIGN(sizeof(WaitEventExtensionSharedState)));
+		wee_state->raw_dsa_area = (void *) p;
+		p += MAXALIGN(wait_event_extension_dsa_init_size());
+		wee_localstate.dsa = dsa_create_in_place(wee_state->raw_dsa_area,
+								  wait_event_extension_dsa_init_size(),
+								  LWTRANCHE_WAITEVENTEXTENSION_DSA,
+								  0);
+		dsa_pin(wee_localstate.dsa);
+
+		/*
+		 * To ensure dshash is created in "plain" shared memory, temporarily
+		 * limit size of dsa to the initial size of the dsa.
+		 */
+		dsa_set_size_limit(wee_localstate.dsa, wait_event_extension_dsa_init_size());
+
+		/*
+		 * With the limit in place, create the dshash table. XXX: It'd be nice
+		 * if there were dshash_create_in_place().
+		 */
+		wee_localstate.dsh = dshash_create(wee_localstate.dsa, &dsh_params, 0);
+		wee_state->hash_handle = dshash_get_hash_table_handle(wee_localstate.dsh);
+
+		dsa_set_size_limit(wee_localstate.dsa, -1);
+
+		/*
+		 * Don't detach from a shared hash table because postmaster can access
+		 * these again via shmem_startup_hooks by extensions.
+		 */
+		Assert(wee_localstate.dsh != NULL);
 	}
 }
 
 /*
- * Allocate a new event ID and return the wait event.
+ * Detach from the dynamic shared hash.
  */
-uint32
-WaitEventExtensionNew(void)
+void
+WaitEventExtensionDetachShmem(void)
 {
-	uint16		eventId;
+	Assert(wee_localstate.dsa);
 
-	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+	dshash_detach(wee_localstate.dsh);
+	wee_localstate.dsh = NULL;
 
-	SpinLockAcquire(&WaitEventExtensionCounter->mutex);
+	dsa_detach(wee_localstate.dsa);
+	wee_localstate.dsa = NULL;
+}
 
-	if (WaitEventExtensionCounter->nextId > PG_UINT16_MAX)
-	{
-		SpinLockRelease(&WaitEventExtensionCounter->mutex);
-		ereport(ERROR,
-				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
-				errmsg("too many wait events for extensions"));
-	}
+/*
+ * Attach the dynamic shared hash for backends. Called From BaseInit().
+ */
+void
+WaitEventExtensionAttachShmem(void)
+{
+	MemoryContext oldcontext;
 
-	eventId = WaitEventExtensionCounter->nextId++;
+	Assert(wee_localstate.dsa == NULL);
 
-	SpinLockRelease(&WaitEventExtensionCounter->mutex);
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
 
-	return PG_WAIT_EXTENSION | eventId;
+	wee_localstate.dsa =
+			dsa_attach_in_place(wee_state->raw_dsa_area,
+								  NULL);
+	dsa_pin_mapping(wee_localstate.dsa);
+
+	wee_localstate.dsh =
+			dshash_attach(wee_localstate.dsa, &dsh_params,
+							wee_state->hash_handle, NULL);
+
+	MemoryContextSwitchTo(oldcontext);
 }
 
 /*
- * Register a dynamic wait event name for extension in the lookup table
- * of the current process.
- *
- * This routine will save a pointer to the wait event name passed as an argument,
- * so the name should be allocated in a backend-lifetime context
- * (shared memory, TopMemoryContext, static constant, or similar).
- *
- * The "wait_event_name" will be user-visible as a wait event name, so try to
- * use a name that fits the style for those.
+ * Allocate a new event ID and return the wait event info.
  */
-void
-WaitEventExtensionRegisterName(uint32 wait_event_info,
-							   const char *wait_event_name)
+uint32
+WaitEventExtensionNew(const char *wait_event_name)
 {
-	uint32		classId;
 	uint16		eventId;
+	bool		found;
+	WaitEventExtensionNameEntry	*entry;
 
-	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
-	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
+	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+	Assert(wee_localstate.dsh != NULL);
 
-	/* Check the wait event class. */
-	if (classId != PG_WAIT_EXTENSION)
+	/* Check the limit of the length of the event name */
+	if (strlen(wait_event_name) >= MAX_WAIT_EVENT_EXTENSION_NAME)
 		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event class %u", classId));
-
-	/* This should only be called for user-defined wait event. */
-	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event ID %u", eventId));
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("wait event name is too long"));
 
-	/* Convert to array index. */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	/* Allocate a new event Id */
+	SpinLockAcquire(&wee_state->mutex);
 
-	/* If necessary, create or enlarge array. */
-	if (eventId >= WaitEventExtensionNamesAllocated)
+	if (wee_state->nextId > PG_UINT16_MAX)
 	{
-		uint32		newalloc;
-
-		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
-
-		if (WaitEventExtensionNames == NULL)
-			WaitEventExtensionNames = (const char **)
-				MemoryContextAllocZero(TopMemoryContext,
-									   newalloc * sizeof(char *));
-		else
-			WaitEventExtensionNames =
-				repalloc0_array(WaitEventExtensionNames, const char *,
-								WaitEventExtensionNamesAllocated, newalloc);
-		WaitEventExtensionNamesAllocated = newalloc;
+		SpinLockRelease(&wee_state->mutex);
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("too many wait events for extensions"));
 	}
 
-	WaitEventExtensionNames[eventId] = wait_event_name;
+	eventId = wee_state->nextId++;
+
+	SpinLockRelease(&wee_state->mutex);
+
+	/* Register a new custom wait event in the shared hash */
+	entry = dshash_find_or_insert(wee_localstate.dsh,
+									&eventId, &found);
+	Assert(!found);
+	strlcpy(entry->wait_event_name, wait_event_name, sizeof(entry->wait_event_name));
+	dshash_release_lock(wee_localstate.dsh, entry);
+
+	return PG_WAIT_EXTENSION | eventId;
 }
 
 /*
@@ -188,23 +283,23 @@ WaitEventExtensionRegisterName(uint32 wait_event_info,
 static const char *
 GetWaitEventExtensionIdentifier(uint16 eventId)
 {
+	WaitEventExtensionNameEntry	*entry;
+
 	/* Built-in event? */
 	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
 		return "Extension";
 
-	/*
-	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
-	 * However, it is possible that the name has never been registered by
-	 * calling WaitEventExtensionRegisterName() in the current process, in
-	 * which case give up and return "extension".
-	 */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
-
-	if (eventId >= WaitEventExtensionNamesAllocated ||
-		WaitEventExtensionNames[eventId] == NULL)
-		return "extension";
+	/* It is a user-defined wait event, so look at the dynamic shared hash */
+	entry = dshash_find(wee_localstate.dsh,
+						&eventId,
+						false);
+	if (entry)
+		dshash_release_lock(wee_localstate.dsh, entry);
+	else
+		ereport(ERROR,
+				errmsg("could not find the name for custom wait event ID %u", eventId));
 
-	return WaitEventExtensionNames[eventId];
+	return entry->wait_event_name;
 }
 
 
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index f31b85c013..3dea5afd8f 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -671,6 +671,9 @@ BaseInit(void)
 	 * drop ephemeral slots, which in turn triggers stats reporting.
 	 */
 	ReplicationSlotInitialize();
+
+    /* Attach the dynamic shared hash for custom wait events */
+	WaitEventExtensionAttachShmem();
 }
 
 
diff --git a/src/include/storage/lwlock.h b/src/include/storage/lwlock.h
index d77410bdea..451c406e47 100644
--- a/src/include/storage/lwlock.h
+++ b/src/include/storage/lwlock.h
@@ -207,6 +207,8 @@ typedef enum BuiltinTrancheIds
 	LWTRANCHE_PGSTATS_DATA,
 	LWTRANCHE_LAUNCHER_DSA,
 	LWTRANCHE_LAUNCHER_HASH,
+	LWTRANCHE_WAITEVENTEXTENSION_DSA,
+	LWTRANCHE_WAITEVENTEXTENSION_HASH,
 	LWTRANCHE_FIRST_USER_DEFINED
 }			BuiltinTrancheIds;
 
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index aad8bc08fa..4c2aa25b78 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -44,12 +44,11 @@ extern PGDLLIMPORT uint32 *my_wait_event_info;
  * Use this category when the server process is waiting for some condition
  * defined by an extension module.
  *
- * Extensions can define their own wait events in this category.  First,
- * they should call WaitEventExtensionNew() to get one or more wait event
- * IDs that are allocated from a shared counter.  These can be used directly
- * with pgstat_report_wait_start() or equivalent.  Next, each individual
- * process should call WaitEventExtensionRegisterName() to associate a wait
- * event string to the number allocated previously.
+ * Extensions can define their own wait events in this category. They should
+ * call WaitEventExtensionNew() with a wait event string. It will get one wait
+ * event ID that is allocated from a shared counter, associate the string to
+ * the number in the shared dynamic hash and return the wait event info. It
+ * can be used directly with pgstat_report_wait_start() or equivalent.
  */
 typedef enum
 {
@@ -59,10 +58,10 @@ typedef enum
 
 extern void WaitEventExtensionShmemInit(void);
 extern Size WaitEventExtensionShmemSize(void);
+extern void WaitEventExtensionDetachShmem(void);
+extern void WaitEventExtensionAttachShmem(void);
 
-extern uint32 WaitEventExtensionNew(void);
-extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
-										   const char *wait_event_name);
+extern uint32 WaitEventExtensionNew(const char *wait_event_name);
 
 /* ----------
  * pgstat_report_wait_start() -
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index c3e7f5fbe6..26b8a49bec 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,25 +39,11 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
-# Check the wait event used by the dynamic bgworker.  For a session without
-# the state in shared memory known, the default of "extension" is the value
-# waited on.
+# Check the wait event used by the dynamic bgworker.
 $result = $node->poll_query_until(
 	'postgres',
 	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	'extension');
-is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
-
-# If the shared memory state is loaded (here with worker_spi_init within
-# the same connection as the one querying pg_stat_activity), the wait
-# event is the custom one.
-# The expected result is a special pattern here with a newline coming from the
-# first query where the shared memory state is set.
-$result = $node->poll_query_until(
-	'postgres',
-	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	qq[
-worker_spi_main]);
+	qq[worker_spi_main]);
 is($result, 1,
 	'dynamic bgworker has reported "worker_spi_main" as wait event');
 
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index c4317351ce..90a762b5e4 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -124,14 +124,10 @@ worker_spi_shmem_init(void)
 
 	/* Define a new wait event */
 	if (!found)
-		wsstate->wait_event = WaitEventExtensionNew();
+		wsstate->wait_event = WaitEventExtensionNew("worker_spi_main");
 
 	LWLockRelease(AddinShmemInitLock);
 
-	/*
-	 * Register the wait event in the lookup table of the current process.
-	 */
-	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
 	return;
 }
 
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 66823bc2a7..0809a28b8f 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2991,7 +2991,7 @@ WaitEventActivity
 WaitEventBufferPin
 WaitEventClient
 WaitEventExtension
-WaitEventExtensionCounterData
+WaitEventExtensionSharedState
 WaitEventIO
 WaitEventIPC
 WaitEventSet
-- 
2.25.1

From ec7d2b0ae47be9bedabb4a4127c914bfb8c361b5 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <mshr.ikeda@ntt.com>
Date: Wed, 9 Aug 2023 19:19:58 +0900
Subject: [PATCH 1/2] Change to manage custom wait events in shared hash

Currently, names of the custom wait event must be registered
per backends. This patch relaxes the constraints by store the
wait events and their names in hash tables in shared memory.

So, all backends can look up the wait event names on
pg_stat_activity view without additional processing by extensions.

In addition, extensions which do not use shared memory for their
use can define custom wait events because WaitEventExtensionNew()
will never allocate new one if the wait event associated to the
name is already allocated. It use hash table to identify uniquness.
---
 doc/src/sgml/monitoring.sgml                  |   7 +-
 doc/src/sgml/xfunc.sgml                       |  10 +-
 src/backend/storage/lmgr/lwlocknames.txt      |   1 +
 src/backend/utils/activity/wait_event.c       | 198 +++++++++++-------
 src/include/utils/wait_event.h                |  17 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  18 +-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 -
 src/test/modules/worker_spi/worker_spi.c      |  22 +-
 8 files changed, 137 insertions(+), 141 deletions(-)

diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index f4fc5d814f..19181832d7 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1121,10 +1121,9 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
      <literal>LWLock</literal> types
      to the list shown in <xref linkend="wait-event-extension-table"/> and
      <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
-     assigned by an extension will not be available in all server processes;
-     so an <literal>Extension</literal> or <literal>LWLock</literal> wait
-     event might be reported as just
-     <quote><literal>extension</literal></quote> rather than the
+     of <literal>LWLock</literal> assigned by an extension will not be
+     available in all server processes; so an wait event might be reported
+     as just <quote><literal>extension</literal></quote> rather than the
      extension-assigned name.
     </para>
    </note>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index d6345a775b..7fec034db4 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3470,17 +3470,13 @@ void RequestAddinShmemSpace(int size)
 </programlisting>
     </para>
     <para>
-     <literal>shmem_startup_hook</literal> can allocate in shared memory
+     <literal>shmem_startup_hook</literal> can allocate in dynamic shared memory
      custom wait events by calling while holding the LWLock
      <function>AddinShmemInitLock</function> to avoid any race conditions:
 <programlisting>
-uint32 WaitEventExtensionNew(void)
-</programlisting>
-     Next, each process needs to associate the wait event allocated previously
-     to a user-facing custom string, which is something done by calling:
-<programlisting>
-void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+uint32 WaitEventExtensionNew(const char *wait_event_name)
 </programlisting>
+     The wait event is associated to a user-facing custom string.
      An example can be found in <filename>src/test/modules/worker_spi</filename>
      in the PostgreSQL source tree.
     </para>
diff --git a/src/backend/storage/lmgr/lwlocknames.txt b/src/backend/storage/lmgr/lwlocknames.txt
index b34b6afecd..7af3e0ba1a 100644
--- a/src/backend/storage/lmgr/lwlocknames.txt
+++ b/src/backend/storage/lmgr/lwlocknames.txt
@@ -53,3 +53,4 @@ XactTruncationLock					44
 # 45 was XactTruncationLock until removal of BackendRandomLock
 WrapLimitsVacuumLock				46
 NotifyQueueTailLock					47
+WaitEventExtensionLock              48
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index b3596ece80..15b2dd06cc 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -45,6 +45,42 @@ uint32	   *my_wait_event_info = &local_my_wait_event_info;
 #define WAIT_EVENT_CLASS_MASK	0xFF000000
 #define WAIT_EVENT_ID_MASK		0x0000FFFF
 
+/*
+ * Hash tables for storing custom wait event ids and their names in
+ * shared memory.
+ *
+ * WaitEventExtensionNameHash is used to find the name from a event id.
+ * It enables all backends look up them without additional processing
+ * per backend like LWLockRegisterTranche().
+ *
+ * WaitEventExtensionIdHash is used to find the event id from a name.
+ * Since it can identify uniquness by the names, extensions that do not
+ * use shared memory also be able to define custom wait events without
+ * defining duplicate wait events.
+ *
+ * The size of hash table is based on assumption. In most of the case,
+ * 16 seem to be sufficient. It seems unlikely that the number of entries
+ * will reach 512.
+ */
+static HTAB *WaitEventExtensionNameHash;	/* find Ids from names */
+static HTAB *WaitEventExtensionIdHash;		/* find names from Ids */
+static const int wee_hash_init_size = 128;
+static const int wee_hash_max_size = 512;
+
+/* hash table entres */
+typedef struct WaitEventExtensionNameEntry
+{
+	uint16	event_id;	/* hash key */
+	char	wait_event_name[NAMEDATALEN]; /* custom wait event name */
+} WaitEventExtensionNameEntry;
+
+typedef struct WaitEventExtensionIdEntry
+{
+	char	wait_event_name[NAMEDATALEN]; /* hash key */
+	uint16	event_id;	/* wait event id */
+} WaitEventExtensionIdEntry;
+
+
 /* dynamic allocation counter for custom wait events in extensions */
 typedef struct WaitEventExtensionCounterData
 {
@@ -59,36 +95,35 @@ static WaitEventExtensionCounterData *WaitEventExtensionCounter;
 #define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
 	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
 
-/*
- * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
- * stores the names of all dynamically-created event IDs known to the current
- * process.  Any unused entries in the array will contain NULL.
- */
-static const char **WaitEventExtensionNames = NULL;
-static int	WaitEventExtensionNamesAllocated = 0;
-
 static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
 
 /*
- *  Return the space for dynamic allocation counter.
+ *  Return the space for dynamic shared hash tables and dynamic allocation counter.
  */
 Size
 WaitEventExtensionShmemSize(void)
 {
-	return sizeof(WaitEventExtensionCounterData);
+	Size	sz;
+	sz = MAXALIGN(sizeof(WaitEventExtensionCounterData));
+	sz = add_size(sz, hash_estimate_size(wee_hash_init_size,
+										sizeof(WaitEventExtensionNameEntry)));
+	sz = add_size(sz, hash_estimate_size(wee_hash_init_size,
+										sizeof(WaitEventExtensionIdEntry)));
+	return sz;
 }
 
 /*
- * Allocate shmem space for dynamic allocation counter.
+ * Allocate shmem space for dynamic shared hash and dynamic allocation counter.
  */
 void
 WaitEventExtensionShmemInit(void)
 {
 	bool		found;
+	HASHCTL		info;
 
 	WaitEventExtensionCounter = (WaitEventExtensionCounterData *)
 		ShmemInitStruct("WaitEventExtensionCounterData",
-						WaitEventExtensionShmemSize(), &found);
+						MAXALIGN(sizeof(WaitEventExtensionCounterData)), &found);
 
 	if (!found)
 	{
@@ -96,21 +131,63 @@ WaitEventExtensionShmemInit(void)
 		WaitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
 		SpinLockInit(&WaitEventExtensionCounter->mutex);
 	}
+
+	/* initialize or atatch the hash tables to store custom wait events */
+	info.keysize = sizeof(uint16);
+	info.entrysize = sizeof(WaitEventExtensionNameEntry);
+	WaitEventExtensionNameHash = ShmemInitHash("Wait Event Extension name hash",
+					wee_hash_init_size,
+					wee_hash_max_size,
+					&info,
+					HASH_ELEM | HASH_BLOBS);
+
+	info.keysize = sizeof(char[NAMEDATALEN]);
+	info.entrysize = sizeof(WaitEventExtensionIdEntry);
+	WaitEventExtensionIdHash = ShmemInitHash("Wait Event Extension id hash",
+					wee_hash_init_size,
+					wee_hash_max_size,
+					&info,
+					HASH_ELEM | HASH_BLOBS);
 }
 
 /*
- * Allocate a new event ID and return the wait event.
+ * Allocate a new event ID and return the wait event info.
+ *
+ * If the wait event name is already defined, this don't
+ * allocate new one, but it return the wait event info
+ * associated to the name.
  */
 uint32
-WaitEventExtensionNew(void)
+WaitEventExtensionNew(const char *wait_event_name)
 {
 	uint16		eventId;
+	bool		found;
+	WaitEventExtensionIdEntry	*id_entry;
+	WaitEventExtensionNameEntry	*name_entry;
 
 	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
 
+	/* Check the limit of the length of the event name */
+	if (strlen(wait_event_name) >= NAMEDATALEN)
+		ereport(ERROR,
+				errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				errmsg("wait event name is too long"));
+
+	/*
+	 * First, check the wait event info associated to the name is
+	 * already defined. Return it if so.
+	 */
+	LWLockAcquire(WaitEventExtensionLock, LW_SHARED);
+	id_entry = (WaitEventExtensionIdEntry *)
+				hash_search(WaitEventExtensionIdHash, &wait_event_name, HASH_FIND, &found);
+	LWLockRelease(WaitEventExtensionLock);
+	if (found)
+		return PG_WAIT_EXTENSION | id_entry->event_id;
+
+	/* Allocate a new event Id */
 	SpinLockAcquire(&WaitEventExtensionCounter->mutex);
 
-	if (WaitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	if (WaitEventExtensionCounter->nextId >= wee_hash_max_size)
 	{
 		SpinLockRelease(&WaitEventExtensionCounter->mutex);
 		ereport(ERROR,
@@ -122,64 +199,22 @@ WaitEventExtensionNew(void)
 
 	SpinLockRelease(&WaitEventExtensionCounter->mutex);
 
-	return PG_WAIT_EXTENSION | eventId;
-}
+	/* Register the new custom wait event in the shared hash table */
+	LWLockAcquire(WaitEventExtensionLock, LW_EXCLUSIVE);
 
-/*
- * Register a dynamic wait event name for extension in the lookup table
- * of the current process.
- *
- * This routine will save a pointer to the wait event name passed as an argument,
- * so the name should be allocated in a backend-lifetime context
- * (shared memory, TopMemoryContext, static constant, or similar).
- *
- * The "wait_event_name" will be user-visible as a wait event name, so try to
- * use a name that fits the style for those.
- */
-void
-WaitEventExtensionRegisterName(uint32 wait_event_info,
-							   const char *wait_event_name)
-{
-	uint32		classId;
-	uint16		eventId;
+	name_entry = (WaitEventExtensionNameEntry *)
+				hash_search(WaitEventExtensionNameHash, &eventId, HASH_ENTER, &found);
+	Assert(!found);
+	strlcpy(name_entry->wait_event_name, wait_event_name, sizeof(name_entry->wait_event_name));
 
-	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
-	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
+	id_entry = (WaitEventExtensionIdEntry *)
+				hash_search(WaitEventExtensionIdHash, &wait_event_name, HASH_ENTER, &found);
+	Assert(!found);
+	id_entry->event_id = eventId;
 
-	/* Check the wait event class. */
-	if (classId != PG_WAIT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event class %u", classId));
+	LWLockRelease(WaitEventExtensionLock);
 
-	/* This should only be called for user-defined wait event. */
-	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event ID %u", eventId));
-
-	/* Convert to array index. */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
-
-	/* If necessary, create or enlarge array. */
-	if (eventId >= WaitEventExtensionNamesAllocated)
-	{
-		uint32		newalloc;
-
-		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
-
-		if (WaitEventExtensionNames == NULL)
-			WaitEventExtensionNames = (const char **)
-				MemoryContextAllocZero(TopMemoryContext,
-									   newalloc * sizeof(char *));
-		else
-			WaitEventExtensionNames =
-				repalloc0_array(WaitEventExtensionNames, const char *,
-								WaitEventExtensionNamesAllocated, newalloc);
-		WaitEventExtensionNamesAllocated = newalloc;
-	}
-
-	WaitEventExtensionNames[eventId] = wait_event_name;
+	return PG_WAIT_EXTENSION | eventId;
 }
 
 /*
@@ -188,23 +223,28 @@ WaitEventExtensionRegisterName(uint32 wait_event_info,
 static const char *
 GetWaitEventExtensionIdentifier(uint16 eventId)
 {
+	bool found;
+	WaitEventExtensionNameEntry	*entry;
+
 	/* Built-in event? */
 	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
 		return "Extension";
 
+	/* It is a user-defined wait event, so lookup hash table. */
+	LWLockAcquire(WaitEventExtensionLock, LW_SHARED);
+	entry = (WaitEventExtensionNameEntry *)
+				hash_search(WaitEventExtensionNameHash, &eventId, HASH_FIND, &found);
+	LWLockRelease(WaitEventExtensionLock);
+
 	/*
-	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
-	 * However, it is possible that the name has never been registered by
-	 * calling WaitEventExtensionRegisterName() in the current process, in
-	 * which case give up and return "extension".
+	 * The entry must be stored because it's registered in
+	 * WaitEventExtensionNew().
 	 */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
-
-	if (eventId >= WaitEventExtensionNamesAllocated ||
-		WaitEventExtensionNames[eventId] == NULL)
-		return "extension";
+	if (!entry)
+		ereport(ERROR,
+				errmsg("could not find the name for custom wait event ID %u", eventId));
 
-	return WaitEventExtensionNames[eventId];
+	return entry->wait_event_name;
 }
 
 
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index aad8bc08fa..51a0b83b79 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -44,12 +44,13 @@ extern PGDLLIMPORT uint32 *my_wait_event_info;
  * Use this category when the server process is waiting for some condition
  * defined by an extension module.
  *
- * Extensions can define their own wait events in this category.  First,
- * they should call WaitEventExtensionNew() to get one or more wait event
- * IDs that are allocated from a shared counter.  These can be used directly
- * with pgstat_report_wait_start() or equivalent.  Next, each individual
- * process should call WaitEventExtensionRegisterName() to associate a wait
- * event string to the number allocated previously.
+ * Extensions can define their own wait events in this category. They should
+ * call WaitEventExtensionNew() with a wait event string. If the wait event
+ * associated the string is already allocated, it returns the wait event info.
+ * If not, it will get one wait event ID that is allocated from a shared counter,
+ * associate the string to the number in the shared dynamic hash and return the
+ * wait event info. The string can be used directly with pgstat_report_wait_start()
+ * or equivalent.
  */
 typedef enum
 {
@@ -60,9 +61,7 @@ typedef enum
 extern void WaitEventExtensionShmemInit(void);
 extern Size WaitEventExtensionShmemSize(void);
 
-extern uint32 WaitEventExtensionNew(void);
-extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
-										   const char *wait_event_name);
+extern uint32 WaitEventExtensionNew(const char *wait_event_name);
 
 /* ----------
  * pgstat_report_wait_start() -
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index c3e7f5fbe6..26b8a49bec 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,25 +39,11 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
-# Check the wait event used by the dynamic bgworker.  For a session without
-# the state in shared memory known, the default of "extension" is the value
-# waited on.
+# Check the wait event used by the dynamic bgworker.
 $result = $node->poll_query_until(
 	'postgres',
 	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	'extension');
-is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
-
-# If the shared memory state is loaded (here with worker_spi_init within
-# the same connection as the one querying pg_stat_activity), the wait
-# event is the custom one.
-# The expected result is a special pattern here with a newline coming from the
-# first query where the shared memory state is set.
-$result = $node->poll_query_until(
-	'postgres',
-	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	qq[
-worker_spi_main]);
+	qq[worker_spi_main]);
 is($result, 1,
 	'dynamic bgworker has reported "worker_spi_main" as wait event');
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index f13f7e0f98..e9d5b07373 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,8 +7,3 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
-
-CREATE FUNCTION worker_spi_init()
-RETURNS VOID STRICT
-AS 'MODULE_PATHNAME'
-LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index c4317351ce..ee1ce13d50 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,7 +44,6 @@
 
 PG_MODULE_MAGIC;
 
-PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
@@ -124,14 +123,10 @@ worker_spi_shmem_init(void)
 
 	/* Define a new wait event */
 	if (!found)
-		wsstate->wait_event = WaitEventExtensionNew();
+		wsstate->wait_event = WaitEventExtensionNew("worker_spi_main");
 
 	LWLockRelease(AddinShmemInitLock);
 
-	/*
-	 * Register the wait event in the lookup table of the current process.
-	 */
-	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
 	return;
 }
 
@@ -434,21 +429,6 @@ _PG_init(void)
 	}
 }
 
-/*
- * Wrapper to initialize a session with the shared memory state
- * used by this module.  This is a convenience routine to be able to
- * see the custom wait event stored in shared memory without loading
- * through shared_preload_libraries.
- */
-Datum
-worker_spi_init(PG_FUNCTION_ARGS)
-{
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
-	PG_RETURN_VOID();
-}
-
 /*
  * Dynamically launch an SPI worker.
  */
-- 
2.25.1

From 13bfa30d804d84502fd5de40c9ab735869a731d8 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <mshr.ikeda@ntt.com>
Date: Wed, 9 Aug 2023 19:40:02 +0900
Subject: [PATCH 2/2] poc: custom wait event for dblink

---
 contrib/dblink/dblink.c | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 41e1f6c91d..1b3dd3736a 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -129,6 +129,7 @@ static void restoreLocalGucs(int nestlevel);
 /* Global */
 static remoteConn *pconn = NULL;
 static HTAB *remoteConnHash = NULL;
+static uint32 *wait_event_info = NULL;
 
 /*
  *	Following is list that holds multiple remote connections.
@@ -203,7 +204,14 @@ dblink_get_conn(char *conname_or_str,
 		dblink_connstr_check(connstr);
 
 		/* OK to make connection */
-		conn = libpqsrv_connect(connstr, WAIT_EVENT_EXTENSION);
+		if (wait_event_info == NULL)
+		{
+			wait_event_info = (uint32 *) MemoryContextAlloc(TopMemoryContext, sizeof(uint32));
+			LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+			*wait_event_info = WaitEventExtensionNew("dblink_get_con");
+			LWLockRelease(AddinShmemInitLock);
+		}
+		conn = libpqsrv_connect(connstr, *wait_event_info);
 
 		if (PQstatus(conn) == CONNECTION_BAD)
 		{
-- 
2.25.1

From 234cc7d852aebf78285ccde383f27ea35a27dad2 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <mshr.ikeda@ntt.com>
Date: Thu, 10 Aug 2023 10:44:41 +0900
Subject: [PATCH 2/2] poc: custom wait event for dblink

---
 contrib/dblink/dblink.c | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 41e1f6c91d..b7158ecd4f 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -129,6 +129,7 @@ static void restoreLocalGucs(int nestlevel);
 /* Global */
 static remoteConn *pconn = NULL;
 static HTAB *remoteConnHash = NULL;
+static uint32 wait_event_info = 0;
 
 /*
  *	Following is list that holds multiple remote connections.
@@ -203,7 +204,9 @@ dblink_get_conn(char *conname_or_str,
 		dblink_connstr_check(connstr);
 
 		/* OK to make connection */
-		conn = libpqsrv_connect(connstr, WAIT_EVENT_EXTENSION);
+		if (wait_event_info == 0)
+			wait_event_info = WaitEventExtensionNew("dblink_get_con");
+		conn = libpqsrv_connect(connstr, wait_event_info);
 
 		if (PQstatus(conn) == CONNECTION_BAD)
 		{
-- 
2.25.1

From 07bbff06a0f58f08d313c02cbb4104ad1db2a0e9 Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <mshr.ikeda@ntt.com>
Date: Thu, 10 Aug 2023 10:43:34 +0900
Subject: [PATCH 1/2] Change to manage custom wait events in shared hash

Currently, names of the custom wait event must be registered
per backends. This patch relaxes the constraints by store the
wait events and their names in hash tables in shared memory.

So, all backends can look up the wait event names on
pg_stat_activity view without additional processing by extensions.

In addition, extensions which do not use shared memory for their
use can define custom wait events because WaitEventExtensionNew()
will never allocate new one if the wait event associated to the
name is already allocated. It use hash table to identify uniqueness.
---
 doc/src/sgml/monitoring.sgml                  |   7 +-
 doc/src/sgml/xfunc.sgml                       |  26 +--
 src/backend/storage/lmgr/lwlocknames.txt      |   1 +
 src/backend/utils/activity/wait_event.c       | 219 +++++++++++-------
 .../utils/activity/wait_event_names.txt       |   1 +
 src/include/utils/wait_event.h                |  17 +-
 .../modules/worker_spi/t/001_worker_spi.pl    |  18 +-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 -
 src/test/modules/worker_spi/worker_spi.c      | 109 +--------
 9 files changed, 165 insertions(+), 238 deletions(-)

diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index f4fc5d814f..19181832d7 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1121,10 +1121,9 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
      <literal>LWLock</literal> types
      to the list shown in <xref linkend="wait-event-extension-table"/> and
      <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
-     assigned by an extension will not be available in all server processes;
-     so an <literal>Extension</literal> or <literal>LWLock</literal> wait
-     event might be reported as just
-     <quote><literal>extension</literal></quote> rather than the
+     of <literal>LWLock</literal> assigned by an extension will not be
+     available in all server processes; so an wait event might be reported
+     as just <quote><literal>extension</literal></quote> rather than the
      extension-assigned name.
     </para>
    </note>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index d6345a775b..281c178b0e 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3454,33 +3454,15 @@ if (!ptr)
    </sect2>
 
    <sect2 id="xfunc-addin-wait-events">
-    <title>Shared Memory and Custom Wait Events</title>
+    <title>Custom Wait Events</title>
 
     <para>
      Add-ins can define custom wait events under the wait event type
-     <literal>Extension</literal>. The add-in's shared library must be
-     preloaded by specifying it in <literal>shared_preload_libraries</literal>,
-     and register a <literal>shmem_request_hook</literal> and a
-     <literal>shmem_startup_hook</literal> in its
-     <function>_PG_init</function> function.
-     <literal>shmem_request_hook</literal> can request a shared memory size
-     to be later used at startup by calling:
+     <literal>Extension</literal> by calling:
 <programlisting>
-void RequestAddinShmemSpace(int size)
-</programlisting>
-    </para>
-    <para>
-     <literal>shmem_startup_hook</literal> can allocate in shared memory
-     custom wait events by calling while holding the LWLock
-     <function>AddinShmemInitLock</function> to avoid any race conditions:
-<programlisting>
-uint32 WaitEventExtensionNew(void)
-</programlisting>
-     Next, each process needs to associate the wait event allocated previously
-     to a user-facing custom string, which is something done by calling:
-<programlisting>
-void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+uint32 WaitEventExtensionNew(const char *wait_event_name)
 </programlisting>
+     The wait event is associated to a user-facing custom string.
      An example can be found in <filename>src/test/modules/worker_spi</filename>
      in the PostgreSQL source tree.
     </para>
diff --git a/src/backend/storage/lmgr/lwlocknames.txt b/src/backend/storage/lmgr/lwlocknames.txt
index b34b6afecd..7af3e0ba1a 100644
--- a/src/backend/storage/lmgr/lwlocknames.txt
+++ b/src/backend/storage/lmgr/lwlocknames.txt
@@ -53,3 +53,4 @@ XactTruncationLock					44
 # 45 was XactTruncationLock until removal of BackendRandomLock
 WrapLimitsVacuumLock				46
 NotifyQueueTailLock					47
+WaitEventExtensionLock              48
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index b3596ece80..d2b21d54a6 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -45,6 +45,42 @@ uint32	   *my_wait_event_info = &local_my_wait_event_info;
 #define WAIT_EVENT_CLASS_MASK	0xFF000000
 #define WAIT_EVENT_ID_MASK		0x0000FFFF
 
+/*
+ * Hash tables for storing custom wait event ids and their names in
+ * shared memory.
+ *
+ * WaitEventExtensionHashById is used to find the name from a event id.
+ * It enables all backends look up them without additional processing
+ * per backend.
+ *
+ * WaitEventExtensionHashByName is used to find the event id from a name.
+ * It's used to ensure that duplicated entries won't be registered
+ * by identify uniqueness by the names
+ *
+ * The size of hash table is based on assumption. In most of the case,
+ * 16 seem to be sufficient. It seems unlikely that the number of entries
+ * will reach 512.
+ */
+static HTAB *WaitEventExtensionHashById;	/* find names from ids */
+static HTAB *WaitEventExtensionHashByName;	/* find ids from names */
+
+#define WaitEventExtensionHashInitSize 128
+#define WaitEventExtensionHashMaxSize 512
+
+/* hash table entres */
+typedef struct WaitEventExtensionEntryById
+{
+	uint16	event_id;						/* hash key */
+	char	wait_event_name[NAMEDATALEN];	/* custom wait event name */
+} WaitEventExtensionEntryById;
+
+typedef struct WaitEventExtensionEntryByName
+{
+	char	wait_event_name[NAMEDATALEN];	/* hash key */
+	uint16	event_id;						/* wait event id */
+} WaitEventExtensionEntryByName;
+
+
 /* dynamic allocation counter for custom wait events in extensions */
 typedef struct WaitEventExtensionCounterData
 {
@@ -59,36 +95,38 @@ static WaitEventExtensionCounterData *WaitEventExtensionCounter;
 #define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
 	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
 
-/*
- * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
- * stores the names of all dynamically-created event IDs known to the current
- * process.  Any unused entries in the array will contain NULL.
- */
-static const char **WaitEventExtensionNames = NULL;
-static int	WaitEventExtensionNamesAllocated = 0;
+/* wait event info for extensions */
+#define WAIT_EVENT_EXTENSION_INFO(eventId)	(PG_WAIT_EXTENSION | eventId)
 
 static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
 
 /*
- *  Return the space for dynamic allocation counter.
+ *  Return the space for dynamic shared hash tables and dynamic allocation counter.
  */
 Size
 WaitEventExtensionShmemSize(void)
 {
-	return sizeof(WaitEventExtensionCounterData);
+	Size	sz;
+	sz = MAXALIGN(sizeof(WaitEventExtensionCounterData));
+	sz = add_size(sz, hash_estimate_size(WaitEventExtensionHashMaxSize,
+										sizeof(WaitEventExtensionEntryById)));
+	sz = add_size(sz, hash_estimate_size(WaitEventExtensionHashMaxSize,
+										sizeof(WaitEventExtensionEntryByName)));
+	return sz;
 }
 
 /*
- * Allocate shmem space for dynamic allocation counter.
+ * Allocate shmem space for dynamic shared hash and dynamic allocation counter.
  */
 void
 WaitEventExtensionShmemInit(void)
 {
 	bool		found;
+	HASHCTL		info;
 
 	WaitEventExtensionCounter = (WaitEventExtensionCounterData *)
 		ShmemInitStruct("WaitEventExtensionCounterData",
-						WaitEventExtensionShmemSize(), &found);
+						MAXALIGN(sizeof(WaitEventExtensionCounterData)), &found);
 
 	if (!found)
 	{
@@ -96,21 +134,77 @@ WaitEventExtensionShmemInit(void)
 		WaitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
 		SpinLockInit(&WaitEventExtensionCounter->mutex);
 	}
+
+	/* initialize or atatch the hash tables to store custom wait events */
+	info.keysize = sizeof(uint16);
+	info.entrysize = sizeof(WaitEventExtensionEntryById);
+	WaitEventExtensionHashById = ShmemInitHash("Wait Event Extension hash by id",
+					WaitEventExtensionHashInitSize,
+					WaitEventExtensionHashMaxSize,
+					&info,
+					HASH_ELEM | HASH_BLOBS);
+
+	info.keysize = sizeof(char[NAMEDATALEN]);
+	info.entrysize = sizeof(WaitEventExtensionEntryByName);
+	WaitEventExtensionHashByName = ShmemInitHash("Wait Event Extension hash by name",
+					WaitEventExtensionHashInitSize,
+					WaitEventExtensionHashMaxSize,
+					&info,
+					HASH_ELEM | HASH_STRINGS);   /* key is Null-terminated C strings */
 }
 
 /*
- * Allocate a new event ID and return the wait event.
+ * Allocate a new event ID and return the wait event info.
+ *
+ * If the wait event name is already defined, this don't
+ * allocate new one, but it return the wait event info
+ * associated to the name.
  */
 uint32
-WaitEventExtensionNew(void)
+WaitEventExtensionNew(const char *wait_event_name)
 {
 	uint16		eventId;
+	bool		found;
+	WaitEventExtensionEntryByName	*entry_by_name;
+	WaitEventExtensionEntryById	*entry_by_id;
+
+	/* Check the limit of the length of the event name */
+	if (strlen(wait_event_name) >= NAMEDATALEN)
+		elog(ERROR,
+			"cannot use custom wait event string longer than %u characters",
+			NAMEDATALEN - 1);
 
-	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+	/*
+	 * Check the wait event info associated to the name is already defined.
+	 * Return it if so.
+	 */
+	LWLockAcquire(WaitEventExtensionLock, LW_SHARED);
+	entry_by_name = (WaitEventExtensionEntryByName *)
+				hash_search(WaitEventExtensionHashByName, wait_event_name,
+							HASH_FIND, &found);
+	LWLockRelease(WaitEventExtensionLock);
+	if (found)
+		return WAIT_EVENT_EXTENSION_INFO(entry_by_name->event_id);
+
+	/*
+	 * Allocate and register a new wait event. But, we need to recheck because
+	 * other processes could already do so while releasing the lock.
+	 */
+	LWLockAcquire(WaitEventExtensionLock, LW_EXCLUSIVE);
+	entry_by_name = (WaitEventExtensionEntryByName *)
+				hash_search(WaitEventExtensionHashByName, wait_event_name,
+							HASH_FIND, &found);
+	/* Recheck */
+	if (found)
+	{
+		LWLockRelease(WaitEventExtensionLock);
+		return WAIT_EVENT_EXTENSION_INFO(entry_by_name->event_id);
+	}
 
+	/* Allocate a new event Id */
 	SpinLockAcquire(&WaitEventExtensionCounter->mutex);
 
-	if (WaitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	if (WaitEventExtensionCounter->nextId >= WaitEventExtensionHashMaxSize)
 	{
 		SpinLockRelease(&WaitEventExtensionCounter->mutex);
 		ereport(ERROR,
@@ -122,64 +216,23 @@ WaitEventExtensionNew(void)
 
 	SpinLockRelease(&WaitEventExtensionCounter->mutex);
 
-	return PG_WAIT_EXTENSION | eventId;
-}
-
-/*
- * Register a dynamic wait event name for extension in the lookup table
- * of the current process.
- *
- * This routine will save a pointer to the wait event name passed as an argument,
- * so the name should be allocated in a backend-lifetime context
- * (shared memory, TopMemoryContext, static constant, or similar).
- *
- * The "wait_event_name" will be user-visible as a wait event name, so try to
- * use a name that fits the style for those.
- */
-void
-WaitEventExtensionRegisterName(uint32 wait_event_info,
-							   const char *wait_event_name)
-{
-	uint32		classId;
-	uint16		eventId;
-
-	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
-	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
-
-	/* Check the wait event class. */
-	if (classId != PG_WAIT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event class %u", classId));
-
-	/* This should only be called for user-defined wait event. */
-	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event ID %u", eventId));
+	/* Register the new wait event */
+	entry_by_id = (WaitEventExtensionEntryById *)
+				hash_search(WaitEventExtensionHashById, &eventId,
+							HASH_ENTER, &found);
+	Assert(!found);
+	strlcpy(entry_by_id->wait_event_name, wait_event_name,
+			sizeof(entry_by_id->wait_event_name));
 
-	/* Convert to array index. */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	entry_by_name = (WaitEventExtensionEntryByName *)
+				hash_search(WaitEventExtensionHashByName, wait_event_name,
+							HASH_ENTER, &found);
+	Assert(!found);
+	entry_by_name->event_id = eventId;
 
-	/* If necessary, create or enlarge array. */
-	if (eventId >= WaitEventExtensionNamesAllocated)
-	{
-		uint32		newalloc;
-
-		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
-
-		if (WaitEventExtensionNames == NULL)
-			WaitEventExtensionNames = (const char **)
-				MemoryContextAllocZero(TopMemoryContext,
-									   newalloc * sizeof(char *));
-		else
-			WaitEventExtensionNames =
-				repalloc0_array(WaitEventExtensionNames, const char *,
-								WaitEventExtensionNamesAllocated, newalloc);
-		WaitEventExtensionNamesAllocated = newalloc;
-	}
+	LWLockRelease(WaitEventExtensionLock);
 
-	WaitEventExtensionNames[eventId] = wait_event_name;
+	return WAIT_EVENT_EXTENSION_INFO(eventId);
 }
 
 /*
@@ -188,23 +241,29 @@ WaitEventExtensionRegisterName(uint32 wait_event_info,
 static const char *
 GetWaitEventExtensionIdentifier(uint16 eventId)
 {
+	bool found;
+	WaitEventExtensionEntryById	*entry;
+
 	/* Built-in event? */
 	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
 		return "Extension";
 
+	/* It is a user-defined wait event, so lookup hash table. */
+	LWLockAcquire(WaitEventExtensionLock, LW_SHARED);
+	entry = (WaitEventExtensionEntryById *)
+				hash_search(WaitEventExtensionHashById, &eventId,
+							HASH_FIND, &found);
+	LWLockRelease(WaitEventExtensionLock);
+
 	/*
-	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
-	 * However, it is possible that the name has never been registered by
-	 * calling WaitEventExtensionRegisterName() in the current process, in
-	 * which case give up and return "extension".
+	 * The entry must be stored because it's registered in
+	 * WaitEventExtensionNew().
 	 */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
-
-	if (eventId >= WaitEventExtensionNamesAllocated ||
-		WaitEventExtensionNames[eventId] == NULL)
-		return "extension";
+	if (!entry)
+		elog(ERROR, "could not find the name for custom wait event ID %u",
+			eventId);
 
-	return WaitEventExtensionNames[eventId];
+	return entry->wait_event_name;
 }
 
 
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 2ea4789b00..a4e5c80bc0 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -317,6 +317,7 @@ WAIT_EVENT_DOCONLY	LogicalRepWorker	"Waiting to read or update the state of logi
 WAIT_EVENT_DOCONLY	XactTruncation	"Waiting to execute <function>pg_xact_status</function> or update the oldest transaction ID available to it."
 WAIT_EVENT_DOCONLY	WrapLimitsVacuum	"Waiting to update limits on transaction id and multixact consumption."
 WAIT_EVENT_DOCONLY	NotifyQueueTail	"Waiting to update limit on <command>NOTIFY</command> message storage."
+WAIT_EVENT_DOCONLY	WaitEventExtension	"Waiting to read or update custom wait events information for extensions."
 
 WAIT_EVENT_DOCONLY	XactBuffer	"Waiting for I/O on a transaction status SLRU buffer."
 WAIT_EVENT_DOCONLY	CommitTsBuffer	"Waiting for I/O on a commit timestamp SLRU buffer."
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index aad8bc08fa..51a0b83b79 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -44,12 +44,13 @@ extern PGDLLIMPORT uint32 *my_wait_event_info;
  * Use this category when the server process is waiting for some condition
  * defined by an extension module.
  *
- * Extensions can define their own wait events in this category.  First,
- * they should call WaitEventExtensionNew() to get one or more wait event
- * IDs that are allocated from a shared counter.  These can be used directly
- * with pgstat_report_wait_start() or equivalent.  Next, each individual
- * process should call WaitEventExtensionRegisterName() to associate a wait
- * event string to the number allocated previously.
+ * Extensions can define their own wait events in this category. They should
+ * call WaitEventExtensionNew() with a wait event string. If the wait event
+ * associated the string is already allocated, it returns the wait event info.
+ * If not, it will get one wait event ID that is allocated from a shared counter,
+ * associate the string to the number in the shared dynamic hash and return the
+ * wait event info. The string can be used directly with pgstat_report_wait_start()
+ * or equivalent.
  */
 typedef enum
 {
@@ -60,9 +61,7 @@ typedef enum
 extern void WaitEventExtensionShmemInit(void);
 extern Size WaitEventExtensionShmemSize(void);
 
-extern uint32 WaitEventExtensionNew(void);
-extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
-										   const char *wait_event_name);
+extern uint32 WaitEventExtensionNew(const char *wait_event_name);
 
 /* ----------
  * pgstat_report_wait_start() -
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index c3e7f5fbe6..26b8a49bec 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,25 +39,11 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
-# Check the wait event used by the dynamic bgworker.  For a session without
-# the state in shared memory known, the default of "extension" is the value
-# waited on.
+# Check the wait event used by the dynamic bgworker.
 $result = $node->poll_query_until(
 	'postgres',
 	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	'extension');
-is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
-
-# If the shared memory state is loaded (here with worker_spi_init within
-# the same connection as the one querying pg_stat_activity), the wait
-# event is the custom one.
-# The expected result is a special pattern here with a newline coming from the
-# first query where the shared memory state is set.
-$result = $node->poll_query_until(
-	'postgres',
-	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	qq[
-worker_spi_main]);
+	qq[worker_spi_main]);
 is($result, 1,
 	'dynamic bgworker has reported "worker_spi_main" as wait event');
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index f13f7e0f98..e9d5b07373 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,8 +7,3 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
-
-CREATE FUNCTION worker_spi_init()
-RETURNS VOID STRICT
-AS 'MODULE_PATHNAME'
-LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index c4317351ce..56e492dd4b 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,33 +44,17 @@
 
 PG_MODULE_MAGIC;
 
-PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
-/* Shared memory state */
-typedef struct worker_spi_state
-{
-	/* the wait event defined during initialization phase */
-	uint32		wait_event;
-} worker_spi_state;
-
-static worker_spi_state *wsstate = NULL;	/* pointer to shared memory */
-
-static shmem_request_hook_type prev_shmem_request_hook = NULL;
-static shmem_request_hook_type prev_shmem_startup_hook = NULL;
-
-static void worker_spi_shmem_request(void);
-static void worker_spi_shmem_startup(void);
-static void worker_spi_shmem_init(void);
-static Size worker_spi_memsize(void);
-
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database = NULL;
 
+/* Local variables */
+static uint32 worker_spi_wait_event = 0;
 
 typedef struct worktable
 {
@@ -78,63 +62,6 @@ typedef struct worktable
 	const char *name;
 } worktable;
 
-static void
-worker_spi_shmem_request(void)
-{
-	if (prev_shmem_request_hook)
-		prev_shmem_request_hook();
-
-	RequestAddinShmemSpace(worker_spi_memsize());
-}
-
-static void
-worker_spi_shmem_startup(void)
-{
-	if (prev_shmem_startup_hook)
-		prev_shmem_startup_hook();
-
-	worker_spi_shmem_init();
-}
-
-static Size
-worker_spi_memsize(void)
-{
-	return MAXALIGN(sizeof(worker_spi_state));
-}
-
-/*
- * Initialize the shared memory state of worker_spi.
- *
- * This routine allocates a new wait event when called the first time.
- * On follow-up calls, the name of the wait event associated with the
- * existing shared memory state is registered.
- */
-static void
-worker_spi_shmem_init(void)
-{
-	bool		found;
-
-	wsstate = NULL;
-
-	/* Create or attach to the shared memory state */
-	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-	wsstate = ShmemInitStruct("worker_spi State",
-							  sizeof(worker_spi_state),
-							  &found);
-
-	/* Define a new wait event */
-	if (!found)
-		wsstate->wait_event = WaitEventExtensionNew();
-
-	LWLockRelease(AddinShmemInitLock);
-
-	/*
-	 * Register the wait event in the lookup table of the current process.
-	 */
-	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
-	return;
-}
-
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -224,9 +151,6 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -268,6 +192,10 @@ worker_spi_main(Datum main_arg)
 	{
 		int			ret;
 
+		/* First time, allocate or get the custom wait event */
+		if (worker_spi_wait_event == 0)
+			worker_spi_wait_event = WaitEventExtensionNew("worker_spi_main");
+
 		/*
 		 * Background workers mustn't call usleep() or any direct equivalent:
 		 * instead, they may wait on their process latch, which sleeps as
@@ -277,7 +205,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 wsstate->wait_event);
+						 worker_spi_wait_event);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -406,11 +334,6 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
-	prev_shmem_request_hook = shmem_request_hook;
-	shmem_request_hook = worker_spi_shmem_request;
-	prev_shmem_startup_hook = shmem_startup_hook;
-	shmem_startup_hook = worker_spi_shmem_startup;
-
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -434,21 +357,6 @@ _PG_init(void)
 	}
 }
 
-/*
- * Wrapper to initialize a session with the shared memory state
- * used by this module.  This is a convenience routine to be able to
- * see the custom wait event stored in shared memory without loading
- * through shared_preload_libraries.
- */
-Datum
-worker_spi_init(PG_FUNCTION_ARGS)
-{
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
-	PG_RETURN_VOID();
-}
-
 /*
  * Dynamically launch an SPI worker.
  */
@@ -461,9 +369,6 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
-- 
2.25.1

From 83e40dace75d0df91e8b7f617bd03d02bd8450d4 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 14 Aug 2023 07:48:15 +0900
Subject: [PATCH v4] Change custom wait events to use shared hash tables

Currently, names of the custom wait event must be registered for each
backends, requiring all these to link to the shared memory area of an
extention, even if not loaded with shared_preload_libraries.

This patch relaxes the constraints by storing the wait events and their
names in hash tables in shared memory.  First, this has the advantage to
simplify the registration of custom wait events to a single routine
call that returns an event ID, either existing or incremented:
uint32 WaitEventExtensionNew(const char *wait_event_name);

Any other backend looking at the wait event information, for instance
via pg_stat_activity, is then able to automatically know about the new
event name.

The code changes done in worker_spi show how things are simplified:
- worker_spi_init() is gone.
- No more shmem hooks.
- The custom wait event ID is cached in the process that needs to set
it, with one single call to WaitEventExtensionNew() to retrieve it.

Per suggestion from Andres Freund.

Author: Masahiro Ikeda
Discussion: https://postgr.es/m/20230801032349.aaiuvhtrcvvcwzcx@awork3.anarazel.de
---
 src/include/utils/wait_event.h                |  18 +-
 src/backend/storage/lmgr/lwlocknames.txt      |   1 +
 src/backend/utils/activity/wait_event.c       | 217 +++++++++++-------
 .../utils/activity/wait_event_names.txt       |   1 +
 .../modules/worker_spi/t/001_worker_spi.pl    |  18 +-
 .../modules/worker_spi/worker_spi--1.0.sql    |   5 -
 src/test/modules/worker_spi/worker_spi.c      | 109 +--------
 doc/src/sgml/monitoring.sgml                  |   5 +-
 doc/src/sgml/xfunc.sgml                       |  26 +--
 src/tools/pgindent/typedefs.list              |   2 +
 10 files changed, 164 insertions(+), 238 deletions(-)

diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index aad8bc08fa..3cffae23e1 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -44,12 +44,14 @@ extern PGDLLIMPORT uint32 *my_wait_event_info;
  * Use this category when the server process is waiting for some condition
  * defined by an extension module.
  *
- * Extensions can define their own wait events in this category.  First,
- * they should call WaitEventExtensionNew() to get one or more wait event
- * IDs that are allocated from a shared counter.  These can be used directly
- * with pgstat_report_wait_start() or equivalent.  Next, each individual
- * process should call WaitEventExtensionRegisterName() to associate a wait
- * event string to the number allocated previously.
+ * Extensions can define their own wait events in this category.  They should
+ * call WaitEventExtensionNew() with a wait event string.  If the wait event
+ * associated the string is already allocated, it returns the wait event info.
+ * If not, it will get one wait event ID that is allocated from a shared
+ + counter, associate the string to the number in the shared dynamic hash and
+ * return the wait event info.
+ *
+ * The ID retrieved can be used with pgstat_report_wait_start() or equivalent.
  */
 typedef enum
 {
@@ -60,9 +62,7 @@ typedef enum
 extern void WaitEventExtensionShmemInit(void);
 extern Size WaitEventExtensionShmemSize(void);
 
-extern uint32 WaitEventExtensionNew(void);
-extern void WaitEventExtensionRegisterName(uint32 wait_event_info,
-										   const char *wait_event_name);
+extern uint32 WaitEventExtensionNew(const char *wait_event_name);
 
 /* ----------
  * pgstat_report_wait_start() -
diff --git a/src/backend/storage/lmgr/lwlocknames.txt b/src/backend/storage/lmgr/lwlocknames.txt
index b34b6afecd..7af3e0ba1a 100644
--- a/src/backend/storage/lmgr/lwlocknames.txt
+++ b/src/backend/storage/lmgr/lwlocknames.txt
@@ -53,3 +53,4 @@ XactTruncationLock					44
 # 45 was XactTruncationLock until removal of BackendRandomLock
 WrapLimitsVacuumLock				46
 NotifyQueueTailLock					47
+WaitEventExtensionLock              48
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index b3596ece80..14750233d4 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -45,6 +45,41 @@ uint32	   *my_wait_event_info = &local_my_wait_event_info;
 #define WAIT_EVENT_CLASS_MASK	0xFF000000
 #define WAIT_EVENT_ID_MASK		0x0000FFFF
 
+/*
+ * Hash tables for storing custom wait event ids and their names in
+ * shared memory.
+ *
+ * WaitEventExtensionHashById is used to find the name from a event id.
+ * Any backend can search it to find custom wait events.
+ *
+ * WaitEventExtensionHashByName is used to find the event ID from a name.
+ * It is used to ensure that no duplicated entries are registered.
+ *
+ * The size of the hash table is based on the assumption that
+ * WAIT_EVENT_EXTENSION_BASH_INIT_SIZE is enough for most cases, and it seems
+ * unlikely that the number of entries will reach
+ * WAIT_EVENT_EXTENSION_BASH_MAX_SIZE.
+ */
+static HTAB *WaitEventExtensionHashById;	/* find names from IDs */
+static HTAB *WaitEventExtensionHashByName;	/* find IDs from names */
+
+#define WAIT_EVENT_EXTENSION_HASH_INIT_SIZE	16
+#define WAIT_EVENT_EXTENSION_HASH_MAX_SIZE	128
+
+/* hash table entries */
+typedef struct WaitEventExtensionEntryById
+{
+	uint16		event_id;		/* hash key */
+	char		wait_event_name[NAMEDATALEN];	/* custom wait event name */
+} WaitEventExtensionEntryById;
+
+typedef struct WaitEventExtensionEntryByName
+{
+	char		wait_event_name[NAMEDATALEN];	/* hash key */
+	uint16		event_id;		/* wait event ID */
+} WaitEventExtensionEntryByName;
+
+
 /* dynamic allocation counter for custom wait events in extensions */
 typedef struct WaitEventExtensionCounterData
 {
@@ -59,36 +94,39 @@ static WaitEventExtensionCounterData *WaitEventExtensionCounter;
 #define NUM_BUILTIN_WAIT_EVENT_EXTENSION	\
 	(WAIT_EVENT_EXTENSION_FIRST_USER_DEFINED - WAIT_EVENT_EXTENSION)
 
-/*
- * This is indexed by event ID minus NUM_BUILTIN_WAIT_EVENT_EXTENSION, and
- * stores the names of all dynamically-created event IDs known to the current
- * process.  Any unused entries in the array will contain NULL.
- */
-static const char **WaitEventExtensionNames = NULL;
-static int	WaitEventExtensionNamesAllocated = 0;
+/* wait event info for extensions */
+#define WAIT_EVENT_EXTENSION_INFO(eventId)	(PG_WAIT_EXTENSION | eventId)
 
 static const char *GetWaitEventExtensionIdentifier(uint16 eventId);
 
 /*
- *  Return the space for dynamic allocation counter.
+ *  Return the space for dynamic shared hash tables and dynamic allocation counter.
  */
 Size
 WaitEventExtensionShmemSize(void)
 {
-	return sizeof(WaitEventExtensionCounterData);
+	Size		sz;
+
+	sz = MAXALIGN(sizeof(WaitEventExtensionCounterData));
+	sz = add_size(sz, hash_estimate_size(WAIT_EVENT_EXTENSION_HASH_MAX_SIZE,
+										 sizeof(WaitEventExtensionEntryById)));
+	sz = add_size(sz, hash_estimate_size(WAIT_EVENT_EXTENSION_HASH_MAX_SIZE,
+										 sizeof(WaitEventExtensionEntryByName)));
+	return sz;
 }
 
 /*
- * Allocate shmem space for dynamic allocation counter.
+ * Allocate shmem space for dynamic shared hash and dynamic allocation counter.
  */
 void
 WaitEventExtensionShmemInit(void)
 {
 	bool		found;
+	HASHCTL		info;
 
 	WaitEventExtensionCounter = (WaitEventExtensionCounterData *)
 		ShmemInitStruct("WaitEventExtensionCounterData",
-						WaitEventExtensionShmemSize(), &found);
+						MAXALIGN(sizeof(WaitEventExtensionCounterData)), &found);
 
 	if (!found)
 	{
@@ -96,21 +134,77 @@ WaitEventExtensionShmemInit(void)
 		WaitEventExtensionCounter->nextId = NUM_BUILTIN_WAIT_EVENT_EXTENSION;
 		SpinLockInit(&WaitEventExtensionCounter->mutex);
 	}
+
+	/* initialize or attach the hash tables to store custom wait events */
+	info.keysize = sizeof(uint16);
+	info.entrysize = sizeof(WaitEventExtensionEntryById);
+	WaitEventExtensionHashById = ShmemInitHash("WaitEventExtension hash by id",
+											   WAIT_EVENT_EXTENSION_HASH_INIT_SIZE,
+											   WAIT_EVENT_EXTENSION_HASH_MAX_SIZE,
+											   &info,
+											   HASH_ELEM | HASH_BLOBS);
+
+	info.keysize = sizeof(char[NAMEDATALEN]);
+	info.entrysize = sizeof(WaitEventExtensionEntryByName);
+	/* key is a NULL-terminated string */
+	WaitEventExtensionHashByName = ShmemInitHash("WaitEventExtension hash by name",
+												 WAIT_EVENT_EXTENSION_HASH_INIT_SIZE,
+												 WAIT_EVENT_EXTENSION_HASH_MAX_SIZE,
+												 &info,
+												 HASH_ELEM | HASH_STRINGS);
 }
 
 /*
- * Allocate a new event ID and return the wait event.
+ * Allocate a new event ID and return the wait event info.
+ *
+ * If the wait event name is already defined, this does not allocate a new
+ * entry; it returns the wait event information associated to the name.
  */
 uint32
-WaitEventExtensionNew(void)
+WaitEventExtensionNew(const char *wait_event_name)
 {
 	uint16		eventId;
+	bool		found;
+	WaitEventExtensionEntryByName *entry_by_name;
+	WaitEventExtensionEntryById *entry_by_id;
 
-	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+	/* Check the limit of the length of the event name */
+	if (strlen(wait_event_name) >= NAMEDATALEN)
+		elog(ERROR,
+			 "cannot use custom wait event string longer than %u characters",
+			 NAMEDATALEN - 1);
 
+	/*
+	 * Check the wait event info associated to the name is already defined.
+	 * Return it if so.
+	 */
+	LWLockAcquire(WaitEventExtensionLock, LW_SHARED);
+	entry_by_name = (WaitEventExtensionEntryByName *)
+		hash_search(WaitEventExtensionHashByName, wait_event_name,
+					HASH_FIND, &found);
+	LWLockRelease(WaitEventExtensionLock);
+	if (found)
+		return WAIT_EVENT_EXTENSION_INFO(entry_by_name->event_id);
+
+	/*
+	 * Allocate and register a new wait event.  Recheck if the event name
+	 * exists, as it could be possible that a concurrent process has inserted
+	 * one with the same name while the lock was previously released.
+	 */
+	LWLockAcquire(WaitEventExtensionLock, LW_EXCLUSIVE);
+	entry_by_name = (WaitEventExtensionEntryByName *)
+		hash_search(WaitEventExtensionHashByName, wait_event_name,
+					HASH_FIND, &found);
+	if (found)
+	{
+		LWLockRelease(WaitEventExtensionLock);
+		return WAIT_EVENT_EXTENSION_INFO(entry_by_name->event_id);
+	}
+
+	/* Allocate a new event Id */
 	SpinLockAcquire(&WaitEventExtensionCounter->mutex);
 
-	if (WaitEventExtensionCounter->nextId > PG_UINT16_MAX)
+	if (WaitEventExtensionCounter->nextId >= WAIT_EVENT_EXTENSION_HASH_MAX_SIZE)
 	{
 		SpinLockRelease(&WaitEventExtensionCounter->mutex);
 		ereport(ERROR,
@@ -122,64 +216,23 @@ WaitEventExtensionNew(void)
 
 	SpinLockRelease(&WaitEventExtensionCounter->mutex);
 
-	return PG_WAIT_EXTENSION | eventId;
-}
+	/* Register the new wait event */
+	entry_by_id = (WaitEventExtensionEntryById *)
+		hash_search(WaitEventExtensionHashById, &eventId,
+					HASH_ENTER, &found);
+	Assert(!found);
+	strlcpy(entry_by_id->wait_event_name, wait_event_name,
+			sizeof(entry_by_id->wait_event_name));
 
-/*
- * Register a dynamic wait event name for extension in the lookup table
- * of the current process.
- *
- * This routine will save a pointer to the wait event name passed as an argument,
- * so the name should be allocated in a backend-lifetime context
- * (shared memory, TopMemoryContext, static constant, or similar).
- *
- * The "wait_event_name" will be user-visible as a wait event name, so try to
- * use a name that fits the style for those.
- */
-void
-WaitEventExtensionRegisterName(uint32 wait_event_info,
-							   const char *wait_event_name)
-{
-	uint32		classId;
-	uint16		eventId;
+	entry_by_name = (WaitEventExtensionEntryByName *)
+		hash_search(WaitEventExtensionHashByName, wait_event_name,
+					HASH_ENTER, &found);
+	Assert(!found);
+	entry_by_name->event_id = eventId;
 
-	classId = wait_event_info & WAIT_EVENT_CLASS_MASK;
-	eventId = wait_event_info & WAIT_EVENT_ID_MASK;
+	LWLockRelease(WaitEventExtensionLock);
 
-	/* Check the wait event class. */
-	if (classId != PG_WAIT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event class %u", classId));
-
-	/* This should only be called for user-defined wait event. */
-	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
-		ereport(ERROR,
-				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				errmsg("invalid wait event ID %u", eventId));
-
-	/* Convert to array index. */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
-
-	/* If necessary, create or enlarge array. */
-	if (eventId >= WaitEventExtensionNamesAllocated)
-	{
-		uint32		newalloc;
-
-		newalloc = pg_nextpower2_32(Max(8, eventId + 1));
-
-		if (WaitEventExtensionNames == NULL)
-			WaitEventExtensionNames = (const char **)
-				MemoryContextAllocZero(TopMemoryContext,
-									   newalloc * sizeof(char *));
-		else
-			WaitEventExtensionNames =
-				repalloc0_array(WaitEventExtensionNames, const char *,
-								WaitEventExtensionNamesAllocated, newalloc);
-		WaitEventExtensionNamesAllocated = newalloc;
-	}
-
-	WaitEventExtensionNames[eventId] = wait_event_name;
+	return WAIT_EVENT_EXTENSION_INFO(eventId);
 }
 
 /*
@@ -188,23 +241,25 @@ WaitEventExtensionRegisterName(uint32 wait_event_info,
 static const char *
 GetWaitEventExtensionIdentifier(uint16 eventId)
 {
+	bool		found;
+	WaitEventExtensionEntryById *entry;
+
 	/* Built-in event? */
 	if (eventId < NUM_BUILTIN_WAIT_EVENT_EXTENSION)
 		return "Extension";
 
-	/*
-	 * It is a user-defined wait event, so look at WaitEventExtensionNames[].
-	 * However, it is possible that the name has never been registered by
-	 * calling WaitEventExtensionRegisterName() in the current process, in
-	 * which case give up and return "extension".
-	 */
-	eventId -= NUM_BUILTIN_WAIT_EVENT_EXTENSION;
+	/* It is a user-defined wait event, so lookup hash table. */
+	LWLockAcquire(WaitEventExtensionLock, LW_SHARED);
+	entry = (WaitEventExtensionEntryById *)
+		hash_search(WaitEventExtensionHashById, &eventId,
+					HASH_FIND, &found);
+	LWLockRelease(WaitEventExtensionLock);
 
-	if (eventId >= WaitEventExtensionNamesAllocated ||
-		WaitEventExtensionNames[eventId] == NULL)
-		return "extension";
+	if (!entry)
+		elog(ERROR, "could not find custom wait event name for ID %u",
+			 eventId);
 
-	return WaitEventExtensionNames[eventId];
+	return entry->wait_event_name;
 }
 
 
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index fcd9d2c63c..f9e01e33b1 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -317,6 +317,7 @@ WAIT_EVENT_DOCONLY	LogicalRepWorker	"Waiting to read or update the state of logi
 WAIT_EVENT_DOCONLY	XactTruncation	"Waiting to execute <function>pg_xact_status</function> or update the oldest transaction ID available to it."
 WAIT_EVENT_DOCONLY	WrapLimitsVacuum	"Waiting to update limits on transaction id and multixact consumption."
 WAIT_EVENT_DOCONLY	NotifyQueueTail	"Waiting to update limit on <command>NOTIFY</command> message storage."
+WAIT_EVENT_DOCONLY	WaitEventExtension	"Waiting to read or update custom wait events information for extensions."
 
 WAIT_EVENT_DOCONLY	XactBuffer	"Waiting for I/O on a transaction status SLRU buffer."
 WAIT_EVENT_DOCONLY	CommitTsBuffer	"Waiting for I/O on a commit timestamp SLRU buffer."
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index c3e7f5fbe6..26b8a49bec 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -39,25 +39,11 @@ $node->poll_query_until('postgres',
 $result = $node->safe_psql('postgres', 'SELECT * FROM schema4.counted;');
 is($result, qq(total|1), 'dynamic bgworker correctly consumed tuple data');
 
-# Check the wait event used by the dynamic bgworker.  For a session without
-# the state in shared memory known, the default of "extension" is the value
-# waited on.
+# Check the wait event used by the dynamic bgworker.
 $result = $node->poll_query_until(
 	'postgres',
 	qq[SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	'extension');
-is($result, 1, 'dynamic bgworker has reported "extension" as wait event');
-
-# If the shared memory state is loaded (here with worker_spi_init within
-# the same connection as the one querying pg_stat_activity), the wait
-# event is the custom one.
-# The expected result is a special pattern here with a newline coming from the
-# first query where the shared memory state is set.
-$result = $node->poll_query_until(
-	'postgres',
-	qq[SELECT worker_spi_init(); SELECT wait_event FROM pg_stat_activity WHERE backend_type ~ 'worker_spi';],
-	qq[
-worker_spi_main]);
+	qq[worker_spi_main]);
 is($result, 1,
 	'dynamic bgworker has reported "worker_spi_main" as wait event');
 
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index f13f7e0f98..e9d5b07373 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -7,8 +7,3 @@ CREATE FUNCTION worker_spi_launch(pg_catalog.int4)
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
-
-CREATE FUNCTION worker_spi_init()
-RETURNS VOID STRICT
-AS 'MODULE_PATHNAME'
-LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index c4317351ce..98f8d4194b 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,33 +44,17 @@
 
 PG_MODULE_MAGIC;
 
-PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
-/* Shared memory state */
-typedef struct worker_spi_state
-{
-	/* the wait event defined during initialization phase */
-	uint32		wait_event;
-} worker_spi_state;
-
-static worker_spi_state *wsstate = NULL;	/* pointer to shared memory */
-
-static shmem_request_hook_type prev_shmem_request_hook = NULL;
-static shmem_request_hook_type prev_shmem_startup_hook = NULL;
-
-static void worker_spi_shmem_request(void);
-static void worker_spi_shmem_startup(void);
-static void worker_spi_shmem_init(void);
-static Size worker_spi_memsize(void);
-
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database = NULL;
 
+/* value cached, fetched from shared memory */
+static uint32 worker_spi_wait_event_main = 0;
 
 typedef struct worktable
 {
@@ -78,63 +62,6 @@ typedef struct worktable
 	const char *name;
 } worktable;
 
-static void
-worker_spi_shmem_request(void)
-{
-	if (prev_shmem_request_hook)
-		prev_shmem_request_hook();
-
-	RequestAddinShmemSpace(worker_spi_memsize());
-}
-
-static void
-worker_spi_shmem_startup(void)
-{
-	if (prev_shmem_startup_hook)
-		prev_shmem_startup_hook();
-
-	worker_spi_shmem_init();
-}
-
-static Size
-worker_spi_memsize(void)
-{
-	return MAXALIGN(sizeof(worker_spi_state));
-}
-
-/*
- * Initialize the shared memory state of worker_spi.
- *
- * This routine allocates a new wait event when called the first time.
- * On follow-up calls, the name of the wait event associated with the
- * existing shared memory state is registered.
- */
-static void
-worker_spi_shmem_init(void)
-{
-	bool		found;
-
-	wsstate = NULL;
-
-	/* Create or attach to the shared memory state */
-	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-	wsstate = ShmemInitStruct("worker_spi State",
-							  sizeof(worker_spi_state),
-							  &found);
-
-	/* Define a new wait event */
-	if (!found)
-		wsstate->wait_event = WaitEventExtensionNew();
-
-	LWLockRelease(AddinShmemInitLock);
-
-	/*
-	 * Register the wait event in the lookup table of the current process.
-	 */
-	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
-	return;
-}
-
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -224,9 +151,6 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -268,6 +192,10 @@ worker_spi_main(Datum main_arg)
 	{
 		int			ret;
 
+		/* First time, allocate or get the custom wait event */
+		if (worker_spi_wait_event_main == 0)
+			worker_spi_wait_event_main = WaitEventExtensionNew("worker_spi_main");
+
 		/*
 		 * Background workers mustn't call usleep() or any direct equivalent:
 		 * instead, they may wait on their process latch, which sleeps as
@@ -277,7 +205,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 wsstate->wait_event);
+						 worker_spi_wait_event_main);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -406,11 +334,6 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
-	prev_shmem_request_hook = shmem_request_hook;
-	shmem_request_hook = worker_spi_shmem_request;
-	prev_shmem_startup_hook = shmem_startup_hook;
-	shmem_startup_hook = worker_spi_shmem_startup;
-
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -434,21 +357,6 @@ _PG_init(void)
 	}
 }
 
-/*
- * Wrapper to initialize a session with the shared memory state
- * used by this module.  This is a convenience routine to be able to
- * see the custom wait event stored in shared memory without loading
- * through shared_preload_libraries.
- */
-Datum
-worker_spi_init(PG_FUNCTION_ARGS)
-{
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
-	PG_RETURN_VOID();
-}
-
 /*
  * Dynamically launch an SPI worker.
  */
@@ -461,9 +369,6 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
-	/* Create (if necessary) and attach to our shared memory area. */
-	worker_spi_shmem_init();
-
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
diff --git a/doc/src/sgml/monitoring.sgml b/doc/src/sgml/monitoring.sgml
index f4fc5d814f..70511a2388 100644
--- a/doc/src/sgml/monitoring.sgml
+++ b/doc/src/sgml/monitoring.sgml
@@ -1121,9 +1121,8 @@ SELECT pid, wait_event_type, wait_event FROM pg_stat_activity WHERE wait_event i
      <literal>LWLock</literal> types
      to the list shown in <xref linkend="wait-event-extension-table"/> and
      <xref linkend="wait-event-lwlock-table"/>. In some cases, the name
-     assigned by an extension will not be available in all server processes;
-     so an <literal>Extension</literal> or <literal>LWLock</literal> wait
-     event might be reported as just
+     of <literal>LWLock</literal> assigned by an extension will not be
+     available in all server processes; It might be reported as just
      <quote><literal>extension</literal></quote> rather than the
      extension-assigned name.
     </para>
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index d6345a775b..281c178b0e 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3454,33 +3454,15 @@ if (!ptr)
    </sect2>
 
    <sect2 id="xfunc-addin-wait-events">
-    <title>Shared Memory and Custom Wait Events</title>
+    <title>Custom Wait Events</title>
 
     <para>
      Add-ins can define custom wait events under the wait event type
-     <literal>Extension</literal>. The add-in's shared library must be
-     preloaded by specifying it in <literal>shared_preload_libraries</literal>,
-     and register a <literal>shmem_request_hook</literal> and a
-     <literal>shmem_startup_hook</literal> in its
-     <function>_PG_init</function> function.
-     <literal>shmem_request_hook</literal> can request a shared memory size
-     to be later used at startup by calling:
+     <literal>Extension</literal> by calling:
 <programlisting>
-void RequestAddinShmemSpace(int size)
-</programlisting>
-    </para>
-    <para>
-     <literal>shmem_startup_hook</literal> can allocate in shared memory
-     custom wait events by calling while holding the LWLock
-     <function>AddinShmemInitLock</function> to avoid any race conditions:
-<programlisting>
-uint32 WaitEventExtensionNew(void)
-</programlisting>
-     Next, each process needs to associate the wait event allocated previously
-     to a user-facing custom string, which is something done by calling:
-<programlisting>
-void WaitEventExtensionRegisterName(uint32 wait_event_info, const char *wait_event_name)
+uint32 WaitEventExtensionNew(const char *wait_event_name)
 </programlisting>
+     The wait event is associated to a user-facing custom string.
      An example can be found in <filename>src/test/modules/worker_spi</filename>
      in the PostgreSQL source tree.
     </para>
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 66823bc2a7..87feb432a0 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2992,6 +2992,8 @@ WaitEventBufferPin
 WaitEventClient
 WaitEventExtension
 WaitEventExtensionCounterData
+WaitEventExtensionEntryById
+WaitEventExtensionEntryByName
 WaitEventIO
 WaitEventIPC
 WaitEventSet
-- 
2.40.1