From 5fc6bc07f2f870f590bd86ac2924aac6a15cd317 Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Fri, 11 Jul 2025 13:51:00 -0500
Subject: [PATCH v1 1/1] Create LWLock tranche in shared memory

---
 src/backend/storage/lmgr/lwlock.c       | 37 +++++++++++++++++-----
 src/backend/utils/activity/wait_event.c | 41 ++++++++++++++++++-------
 src/include/storage/lwlock.h            |  1 +
 src/include/utils/wait_event.h          |  5 +++
 4 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/src/backend/storage/lmgr/lwlock.c b/src/backend/storage/lmgr/lwlock.c
index 46f44bc4511..eeb3c4501c1 100644
--- a/src/backend/storage/lmgr/lwlock.c
+++ b/src/backend/storage/lmgr/lwlock.c
@@ -130,6 +130,12 @@ StaticAssertDecl((LW_VAL_EXCLUSIVE & LW_FLAG_MASK) == 0,
  * or LWLockRegisterTranche.  The names of these that are known in the current
  * process appear in LWLockTrancheNames[].
  *
+ * 4. Extensions can also create a new tranche and define a custom wait event
+ * using LWLockRegisterTrancheWaitEventCustom. This function takes only a
+ * tranche name and returns a trancheId to the caller. Unlike
+ * RequestNamedLWLockTranche or LWLockRegisterTranche, the wait events are
+ * stored in shared memory and are visible to all backends.
+ *
  * All these names are user-visible as wait event names, so choose with care
  * ... and do not forget to update the documentation's list of wait events.
  */
@@ -666,6 +672,12 @@ LWLockRegisterTranche(int tranche_id, const char *tranche_name)
 	LWLockTrancheNames[tranche_id] = tranche_name;
 }
 
+uint16
+LWLockRegisterTrancheWaitEventCustom(const char *tranche_name)
+{
+	return WaitEventLWLockNew(tranche_name) & WAIT_EVENT_ID_MASK;
+}
+
 /*
  * RequestNamedLWLockTranche
  *		Request that extra LWLocks be allocated during postmaster
@@ -754,22 +766,33 @@ LWLockReportWaitEnd(void)
 static const char *
 GetLWTrancheName(uint16 trancheId)
 {
+	const char *wait_event_name = NULL;
+	uint32		wait_event_info;
+	uint16		trancheIdSaved = trancheId;
+
 	/* Built-in tranche or individual LWLock? */
 	if (trancheId < LWTRANCHE_FIRST_USER_DEFINED)
 		return BuiltinTrancheNames[trancheId];
 
 	/*
-	 * It's an extension tranche, so look in LWLockTrancheNames[].  However,
-	 * it's possible that the tranche has never been registered in the current
-	 * process, in which case give up and return "extension".
+	 * It's an extension tranche, so first look in LWLockTrancheNames[] and
+	 * then in custom wait events, in that order, ensuring shared memory is
+	 * only checked if necessary. However, it is possible that the tranche has
+	 * never been registered in the current process; in that case, give up and
+	 * return "extension".
 	 */
 	trancheId -= LWTRANCHE_FIRST_USER_DEFINED;
 
-	if (trancheId >= LWLockTrancheNamesAllocated ||
-		LWLockTrancheNames[trancheId] == NULL)
-		return "extension";
+	if (trancheId < LWLockTrancheNamesAllocated &&
+		LWLockTrancheNames[trancheId] != NULL)
+		return LWLockTrancheNames[trancheId];
+
+	wait_event_info = PG_WAIT_LWLOCK | trancheIdSaved;
+	wait_event_name = GetWaitEventCustomIdentifier(wait_event_info, false);
+	if (wait_event_name)
+		return wait_event_name;
 
-	return LWLockTrancheNames[trancheId];
+	return "extension";
 }
 
 /*
diff --git a/src/backend/utils/activity/wait_event.c b/src/backend/utils/activity/wait_event.c
index d9b8f34a355..5e0bda8714a 100644
--- a/src/backend/utils/activity/wait_event.c
+++ b/src/backend/utils/activity/wait_event.c
@@ -39,9 +39,6 @@ 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_CLASS_MASK	0xFF000000
-#define WAIT_EVENT_ID_MASK		0x0000FFFF
-
 /*
  * Hash tables for storing custom wait event ids and their names in
  * shared memory.
@@ -94,7 +91,6 @@ static WaitEventCustomCounterData *WaitEventCustomCounter;
 #define WAIT_EVENT_CUSTOM_INITIAL_ID	1
 
 static uint32 WaitEventCustomNew(uint32 classId, const char *wait_event_name);
-static const char *GetWaitEventCustomIdentifier(uint32 wait_event_info);
 
 /*
  *  Return the space for dynamic shared hash tables and dynamic allocation counter.
@@ -171,6 +167,12 @@ WaitEventInjectionPointNew(const char *wait_event_name)
 	return WaitEventCustomNew(PG_WAIT_INJECTIONPOINT, wait_event_name);
 }
 
+uint32
+WaitEventLWLockNew(const char *wait_event_name)
+{
+	return WaitEventCustomNew(PG_WAIT_LWLOCK, wait_event_name);
+}
+
 static uint32
 WaitEventCustomNew(uint32 classId, const char *wait_event_name)
 {
@@ -245,7 +247,19 @@ WaitEventCustomNew(uint32 classId, const char *wait_event_name)
 				errmsg("too many custom wait events"));
 	}
 
-	eventId = WaitEventCustomCounter->nextId++;
+	/*
+	 * If classId refers to an LWLock, the corresponding eventId should be a
+	 * trancheId derived from the global LWLockCounter. In this case, we
+	 * increment the WaitEventCustomCounter for tracking and assign the next
+	 * trancheId using LWLockNewTrancheId().
+	 */
+	if (classId == PG_WAIT_LWLOCK)
+	{
+		WaitEventCustomCounter->nextId++;
+		eventId = LWLockNewTrancheId();
+	}
+	else
+		eventId = WaitEventCustomCounter->nextId++;
 
 	SpinLockRelease(&WaitEventCustomCounter->mutex);
 
@@ -272,8 +286,8 @@ WaitEventCustomNew(uint32 classId, const char *wait_event_name)
 /*
  * Return the name of a custom wait event information.
  */
-static const char *
-GetWaitEventCustomIdentifier(uint32 wait_event_info)
+const char *
+GetWaitEventCustomIdentifier(uint32 wait_event_info, bool error_if_not_found)
 {
 	bool		found;
 	WaitEventCustomEntryByInfo *entry;
@@ -290,9 +304,14 @@ GetWaitEventCustomIdentifier(uint32 wait_event_info)
 	LWLockRelease(WaitEventCustomLock);
 
 	if (!entry)
-		elog(ERROR,
-			 "could not find custom name for wait event information %u",
-			 wait_event_info);
+	{
+		if (error_if_not_found)
+			elog(ERROR,
+				 "could not find custom name for wait event information %u",
+				 wait_event_info);
+		else
+			return NULL;
+	}
 
 	return entry->wait_event_name;
 }
@@ -451,7 +470,7 @@ pgstat_get_wait_event(uint32 wait_event_info)
 			break;
 		case PG_WAIT_EXTENSION:
 		case PG_WAIT_INJECTIONPOINT:
-			event_name = GetWaitEventCustomIdentifier(wait_event_info);
+			event_name = GetWaitEventCustomIdentifier(wait_event_info, true);
 			break;
 		case PG_WAIT_BUFFERPIN:
 			{
diff --git a/src/include/storage/lwlock.h b/src/include/storage/lwlock.h
index 08a72569ae5..26067568be6 100644
--- a/src/include/storage/lwlock.h
+++ b/src/include/storage/lwlock.h
@@ -170,6 +170,7 @@ extern LWLockPadded *GetNamedLWLockTranche(const char *tranche_name);
  */
 extern int	LWLockNewTrancheId(void);
 extern void LWLockRegisterTranche(int tranche_id, const char *tranche_name);
+extern uint16 LWLockRegisterTrancheWaitEventCustom(const char *tranche_name);
 extern void LWLockInitialize(LWLock *lock, int tranche_id);
 
 /*
diff --git a/src/include/utils/wait_event.h b/src/include/utils/wait_event.h
index f5815b4994a..e1205cd9375 100644
--- a/src/include/utils/wait_event.h
+++ b/src/include/utils/wait_event.h
@@ -22,6 +22,8 @@ extern void pgstat_reset_wait_event_storage(void);
 
 extern PGDLLIMPORT uint32 *my_wait_event_info;
 
+#define WAIT_EVENT_CLASS_MASK	0xFF000000
+#define WAIT_EVENT_ID_MASK		0x0000FFFF
 
 /*
  * Wait Events - Extension, InjectionPoint
@@ -41,10 +43,13 @@ extern PGDLLIMPORT uint32 *my_wait_event_info;
  */
 extern uint32 WaitEventExtensionNew(const char *wait_event_name);
 extern uint32 WaitEventInjectionPointNew(const char *wait_event_name);
+extern uint32 WaitEventLWLockNew(const char *wait_event_name);
 
 extern void WaitEventCustomShmemInit(void);
 extern Size WaitEventCustomShmemSize(void);
 extern char **GetWaitEventCustomNames(uint32 classId, int *nwaitevents);
+extern const char *GetWaitEventCustomIdentifier(uint32 wait_event_info,
+												bool error_if_not_found);
 
 /* ----------
  * pgstat_report_wait_start() -
-- 
2.39.5 (Apple Git-154)

