On 2023-07-19 12:52, Masahiro Ikeda wrote:
Hi,

I updated the patches.
* v6-0001-Support-custom-wait-events-for-extensions.patch

I updated the patch since the cfbot found a bug.
* v7-0001-Support-custom-wait-events-for-extensions.patch

Regards,
--
Masahiro Ikeda
NTT DATA CORPORATION
From e5d63913cbba9f20098252cc1e415eda4d32f5ad Mon Sep 17 00:00:00 2001
From: Masahiro Ikeda <masahiro.ikeda...@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

Reply via email to