From 2e65230485d84b685858c196f5d64b83bfe3bb94 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Fri, 28 Feb 2025 21:19:08 +0200
Subject: [PATCH 2/3] Use ModifyWaitEvent to update exit_on_postmaster_death

This is in preparation for splitting WaitEventSet related functions to
a separate source file. That will hide the details of WaitEventSet
from WaitLatch, so it must use an exposed function instead of
modifying WaitEventSet->exit_on_postmaster_death directly.
---
 src/backend/storage/ipc/latch.c | 42 +++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/backend/storage/ipc/latch.c b/src/backend/storage/ipc/latch.c
index cd09a10f3d1..ab601c748f8 100644
--- a/src/backend/storage/ipc/latch.c
+++ b/src/backend/storage/ipc/latch.c
@@ -154,8 +154,9 @@ struct WaitEventSet
 /* A common WaitEventSet used to implement WaitLatch() */
 static WaitEventSet *LatchWaitSet;
 
-/* The position of the latch in LatchWaitSet. */
+/* The positions of the latch and PM death events in LatchWaitSet */
 #define LatchWaitSetLatchPos 0
+#define LatchWaitSetPostmasterDeathPos 1
 
 #ifndef WIN32
 /* Are we currently in WaitLatch? The signal handler would like to know. */
@@ -353,11 +354,18 @@ InitializeLatchWaitSet(void)
 	LatchWaitSet = CreateWaitEventSet(NULL, 2);
 	latch_pos = AddWaitEventToSet(LatchWaitSet, WL_LATCH_SET, PGINVALID_SOCKET,
 								  MyLatch, NULL);
-	if (IsUnderPostmaster)
-		AddWaitEventToSet(LatchWaitSet, WL_EXIT_ON_PM_DEATH,
-						  PGINVALID_SOCKET, NULL, NULL);
-
 	Assert(latch_pos == LatchWaitSetLatchPos);
+
+	/*
+	 * WaitLatch will modify this to WL_EXIT_ON_PM_DEATH or
+	 * WL_POSTMASTER_DEATH on each call.
+	 */
+	if (IsUnderPostmaster)
+	{
+		latch_pos = AddWaitEventToSet(LatchWaitSet, WL_EXIT_ON_PM_DEATH,
+									  PGINVALID_SOCKET, NULL, NULL);
+		Assert(latch_pos == LatchWaitSetPostmasterDeathPos);
+	}
 }
 
 /*
@@ -505,8 +513,14 @@ WaitLatch(Latch *latch, int wakeEvents, long timeout,
 	if (!(wakeEvents & WL_LATCH_SET))
 		latch = NULL;
 	ModifyWaitEvent(LatchWaitSet, LatchWaitSetLatchPos, WL_LATCH_SET, latch);
-	LatchWaitSet->exit_on_postmaster_death =
-		((wakeEvents & WL_EXIT_ON_PM_DEATH) != 0);
+
+	/*
+	 * Update the event set for whether WL_EXIT_ON_PM_DEATH or
+	 * WL_POSTMASTER_DEATH was requested.  This is also cheap.
+	 */
+	ModifyWaitEvent(LatchWaitSet, LatchWaitSetPostmasterDeathPos,
+					(wakeEvents & (WL_EXIT_ON_PM_DEATH | WL_POSTMASTER_DEATH)),
+					0);
 
 	if (WaitEventSetWait(LatchWaitSet,
 						 (wakeEvents & WL_TIMEOUT) ? timeout : -1,
@@ -1037,15 +1051,19 @@ ModifyWaitEvent(WaitEventSet *set, int pos, uint32 events, Latch *latch)
 		(!(event->events & WL_LATCH_SET) || set->latch == latch))
 		return;
 
-	if (event->events & WL_LATCH_SET &&
-		events != event->events)
+	/* Allow switching between WL_POSTMASTER_DEATH and WL_EXIT_ON_PM_DEATH */
+	if (event->events & WL_POSTMASTER_DEATH)
 	{
-		elog(ERROR, "cannot modify latch event");
+		if (events != WL_POSTMASTER_DEATH && events != WL_EXIT_ON_PM_DEATH)
+			elog(ERROR, "cannot modify postmaster death event");
+		set->exit_on_postmaster_death = ((events & WL_EXIT_ON_PM_DEATH) != 0);
+		return;
 	}
 
-	if (event->events & WL_POSTMASTER_DEATH)
+	if (event->events & WL_LATCH_SET &&
+		events != event->events)
 	{
-		elog(ERROR, "cannot modify postmaster death event");
+		elog(ERROR, "cannot modify latch event");
 	}
 
 	/* FIXME: validate event mask */
-- 
2.39.5

