v3-0001-Support-custom-wait-events-for-extensions.patch

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