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 in the 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 ID allocated.

Note that this includes an example of how to use this new facility in
worker_spi with tests in TAP for various scenarios, and some
documentation about how to use them.

Any code in the tree that currently uses WAIT_EVENT_EXTENSION could
switch to this new facility to define custom wait events.  This is left
as work for future patches.

Author: Masahiro Ikeda
Reviewed-by: Andres Freund, Michael Paquier, Tristan Partin, Bharath
Rupireddy
Discussion: https://postgr.es/m/b9f5411acda0cf15c8fbb767702ff43e@oss.nttdata.com

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
 

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;

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;

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